added multi-key support for optional auth logic

gateway_test.go

@@ -211,7 +211,7 @@ func TestGatewayCompletionsFailover(t *testing.T) {
     uniqueModels: []string{"model-1", "model-2"},
   }
 
-  handler := handleChatCompletions(cm, "")
+  handler := handleChatCompletions(cm, nil)
 
   // 1. Test non-streaming failover request starting with model-1
   reqBodyObj := map[string]any{
@@ -284,7 +284,7 @@ func TestGatewayCompletionsStreamingFailover(t *testing.T) {
     uniqueModels: []string{"model-1", "model-2"},
   }
 
-  handler := handleChatCompletions(cm, "")
+  handler := handleChatCompletions(cm, nil)
 
   reqBodyObj := map[string]any{
     "model":  "model-1",
@@ -328,7 +328,7 @@ func TestModelsEndpoint(t *testing.T) {
     uniqueModels: []string{"model-1", "model-2"},
   }
 
-  handler := handleModels(cm, "")
+  handler := handleModels(cm, nil)
 
   req := httptest.NewRequest("GET", "/v1/models", nil)
   w := httptest.NewRecorder()
@@ -367,9 +367,10 @@ func TestGatewayAuthentication(t *testing.T) {
     uniqueModels: []string{"model-1"},
   }
 
-  // Set up handlers expecting "secret-token"
-  modelsHandler := handleModels(cm, "secret-token")
-  completionsHandler := handleChatCompletions(cm, "secret-token")
+  // Set up handlers expecting multiple secret tokens
+  expectedTokens := []string{"secret-token-1", "secret-token-2"}
+  modelsHandler := handleModels(cm, expectedTokens)
+  completionsHandler := handleChatCompletions(cm, expectedTokens)
 
   // Case 1: No Authorization header
   {
@@ -392,14 +393,25 @@ func TestGatewayAuthentication(t *testing.T) {
     }
   }
 
-  // Case 3: Correct Authorization header
+  // Case 3: Correct Authorization header (first token)
   {
     req := httptest.NewRequest("GET", "/v1/models", nil)
-    req.Header.Set("Authorization", "Bearer secret-token")
+    req.Header.Set("Authorization", "Bearer secret-token-1")
     w := httptest.NewRecorder()
     modelsHandler.ServeHTTP(w, req)
     if w.Result().StatusCode != http.StatusOK {
-      t.Errorf("Expected 200 OK for correct token, got %d", w.Result().StatusCode)
+      t.Errorf("Expected 200 OK for correct token 1, got %d", w.Result().StatusCode)
+    }
+  }
+
+  // Case 3b: Correct Authorization header (second token)
+  {
+    req := httptest.NewRequest("GET", "/v1/models", nil)
+    req.Header.Set("Authorization", "Bearer secret-token-2")
+    w := httptest.NewRecorder()
+    modelsHandler.ServeHTTP(w, req)
+    if w.Result().StatusCode != http.StatusOK {
+      t.Errorf("Expected 200 OK for correct token 2, got %d", w.Result().StatusCode)
     }
   }
 
@@ -410,7 +422,7 @@ func TestGatewayAuthentication(t *testing.T) {
     }
     reqBytes, _ := json.Marshal(reqBodyObj)
     req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
-    req.Header.Set("Authorization", "Bearer secret-token")
+    req.Header.Set("Authorization", "Bearer secret-token-2")
     req.Header.Set("Content-Type", "application/json")
     w := httptest.NewRecorder()
     completionsHandler.ServeHTTP(w, req)
@@ -492,7 +504,7 @@ func TestGatewayImageGenerations(t *testing.T) {
     uniqueModels: []string{"dall-e-3"},
   }
 
-  handler := handleImageGenerations(cm, "")
+  handler := handleImageGenerations(cm, nil)
 
   reqBodyObj := map[string]any{
     "prompt": "a beautiful kitten",
@@ -551,7 +563,7 @@ func TestGatewayNoAuthHeaderOnEmptyKey(t *testing.T) {
     uniqueModels: []string{"model-no-key", "model-no-key-2"},
   }
 
-  handler := handleChatCompletions(cm, "")
+  handler := handleChatCompletions(cm, nil)
 
   // Case 1: Empty Key
   {
@@ -615,7 +627,7 @@ func TestGatewayBlankKey(t *testing.T) {
     uniqueModels: []string{"model-blank-key"},
   }
 
-  handler := handleChatCompletions(cm, "")
+  handler := handleChatCompletions(cm, nil)
 
   reqBodyObj := map[string]any{
     "model": "model-blank-key",
@@ -664,7 +676,7 @@ func TestGatewayExtraHeaders(t *testing.T) {
     uniqueModels: []string{"model-extra-headers"},
   }
 
-  handler := handleChatCompletions(cm, "")
+  handler := handleChatCompletions(cm, nil)
 
   reqBodyObj := map[string]any{
     "model": "model-extra-headers",
@@ -685,5 +697,87 @@ func TestGatewayExtraHeaders(t *testing.T) {
   }
 }
 
+func TestMultipleTokensKeyFile(t *testing.T) {
+  // Create a temporary key file with multiple tokens, some empty lines, and comments/spaces
+  tmpDir, err := os.MkdirTemp("", "dynagate-keyfile-test")
+  if err != nil {
+    t.Fatalf("failed to create temp dir: %v", err)
+  }
+  defer os.RemoveAll(tmpDir)
+
+  keyFilePath := filepath.Join(tmpDir, "keys.txt")
+  keyFileContent := "\n  token-a  \n\ntoken-b\n   \ntoken-c\n"
+  if err := os.WriteFile(keyFilePath, []byte(keyFileContent), 0644); err != nil {
+    t.Fatalf("failed to write key file: %v", err)
+  }
+
+  // Parse using the same logic as main.go
+  content, err := os.ReadFile(keyFilePath)
+  if err != nil {
+    t.Fatalf("failed to read key file: %v", err)
+  }
+  var authTokens []string
+  for _, line := range strings.Split(string(content), "\n") {
+    token := strings.TrimSpace(line)
+    if token != "" {
+      authTokens = append(authTokens, token)
+    }
+  }
+
+  expected := []string{"token-a", "token-b", "token-c"}
+  if !reflect.DeepEqual(authTokens, expected) {
+    t.Errorf("Expected parsed tokens to be %v, got %v", expected, authTokens)
+  }
+
+  // Verify that the helper checkAuth works with these parsed tokens
+  req := httptest.NewRequest("GET", "/v1/models", nil)
+  req.Header.Set("Authorization", "Bearer token-b")
+  if !checkAuth(authTokens, req) {
+    t.Errorf("Expected token-b to authenticate successfully")
+  }
+
+  reqWrong := httptest.NewRequest("GET", "/v1/models", nil)
+  reqWrong.Header.Set("Authorization", "Bearer token-wrong")
+  if checkAuth(authTokens, reqWrong) {
+    t.Errorf("Expected token-wrong to fail authentication")
+  }
+
+  // Test empty key file content behavior
+  emptyKeyFilePath := filepath.Join(tmpDir, "empty_keys.txt")
+  emptyKeyFileContent := "\n   \n\n  \n"
+  if err := os.WriteFile(emptyKeyFilePath, []byte(emptyKeyFileContent), 0644); err != nil {
+    t.Fatalf("failed to write empty key file: %v", err)
+  }
+
+  // Parse using empty key file logic as in main.go
+  emptyContent, err := os.ReadFile(emptyKeyFilePath)
+  if err != nil {
+    t.Fatalf("failed to read empty key file: %v", err)
+  }
+  var emptyAuthTokens []string
+  var tokens []string
+  for _, line := range strings.Split(string(emptyContent), "\n") {
+    token := strings.TrimSpace(line)
+    if token != "" {
+      tokens = append(tokens, token)
+    }
+  }
+  if len(tokens) > 0 {
+    emptyAuthTokens = tokens
+  } else {
+    emptyAuthTokens = nil
+  }
+
+  if emptyAuthTokens != nil {
+    t.Errorf("Expected emptyAuthTokens to be nil, got %v", emptyAuthTokens)
+  }
+
+  // Verify that checkAuth returns true when expectedTokens is nil
+  reqNoAuth := httptest.NewRequest("GET", "/v1/models", nil)
+  if !checkAuth(emptyAuthTokens, reqNoAuth) {
+    t.Errorf("Expected checkAuth to return true for nil expectedTokens")
+  }
+}
+