luxferre/dynagate
1
0
Fork 0
Code Actions Activity

added multi-key support for optional auth logic

Browse Source
This commit is contained in:
Luxferre
2026-07-11 10:02:48 +03:00
parent 2439257729
commit 1c25a98208
4 changed files with 144 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+5 -4
View File
@@ -76,9 +76,10 @@ ollama-llama3,-,http://localhost:11434/v1,
```
### 3. Setup gateway key (optional)
To restrict access to your gateway, write a single plaintext key to a file (e.g. `gatekey.txt`):
To restrict access to your gateway, write one or more plaintext keys (one per line) to a file (e.g. `gatekey.txt`):
```text
my-secure-gateway-token
my-secure-gateway-token-1
my-secure-gateway-token-2
```
### 4. Run the server
@@ -94,7 +95,7 @@ my-secure-gateway-token
| ---- | ---- | ------- | ----------- |
| `-port` | int | `8080` | Port the gateway listens on. |
| `-csv` | string | `models.csv` | Path to the model configuration CSV. |
| `-key` | string | `""` | Path to a file containing the gateway's access token. If blank, authentication is disabled. |
| `-key` | string | `""` | Path to a file containing the gateway's access tokens (one per line). If blank, authentication is disabled. |
| `-csv-updater` | string | `""` | Command/script to run periodically to update `models.csv`. |
| `-csv-update-interval`| int | `10` | Frequency in minutes to invoke `-csv-updater`. |
@@ -152,7 +153,7 @@ curl -i -X POST http://localhost:8080/v1/chat/completions \
### 401 Unauthorized errors
- **Symptom**: Requests return `{"error": {"code": "invalid_api_key", ...}}`.
- **Solution**: Check that you are passing the correct Bearer token configured in the `-key` token file. Trim any whitespace, and verify it starts with `Bearer `.
- **Solution**: Check that you are passing one of the correct Bearer tokens configured in the `-key` token file. Trim any whitespace, and verify it starts with `Bearer `.
## FAQ
gateway_test.go
+108 -14
View File
@@ -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")
}
}
handler.go
+14 -9
View File
@@ -14,8 +14,8 @@ import (
var httpClient = &http.Client{}
func checkAuth(expectedToken string, r *http.Request) bool {
if expectedToken == "" {
func checkAuth(expectedTokens []string, r *http.Request) bool {
if expectedTokens == nil {
return true
}
authHeader := r.Header.Get("Authorization")
@@ -23,7 +23,12 @@ func checkAuth(expectedToken string, r *http.Request) bool {
return false
}
token := strings.TrimPrefix(authHeader, "Bearer ")
return token == expectedToken
for _, expected := range expectedTokens {
if token == expected {
return true
}
}
return false
}
func sendUnauthorized(w http.ResponseWriter) {
@@ -39,7 +44,7 @@ func sendUnauthorized(w http.ResponseWriter) {
})
}
func handleModels(cm *ConfigManager, expectedToken string) http.HandlerFunc {
func handleModels(cm *ConfigManager, expectedTokens []string) http.HandlerFunc {
type ModelData struct {
ID string `json:"id"`
Object string `json:"object"`
@@ -53,7 +58,7 @@ func handleModels(cm *ConfigManager, expectedToken string) http.HandlerFunc {
}
return func(w http.ResponseWriter, r *http.Request) {
if !checkAuth(expectedToken, r) {
if !checkAuth(expectedTokens, r) {
sendUnauthorized(w)
return
}
@@ -91,9 +96,9 @@ func handleModels(cm *ConfigManager, expectedToken string) http.HandlerFunc {
}
}
func handleChatCompletions(cm *ConfigManager, expectedToken string) http.HandlerFunc {
func handleChatCompletions(cm *ConfigManager, expectedTokens []string) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if !checkAuth(expectedToken, r) {
if !checkAuth(expectedTokens, r) {
sendUnauthorized(w)
return
}
@@ -324,9 +329,9 @@ func handleChatCompletions(cm *ConfigManager, expectedToken string) http.Handler
}
}
func handleImageGenerations(cm *ConfigManager, expectedToken string) http.HandlerFunc {
func handleImageGenerations(cm *ConfigManager, expectedTokens []string) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if !checkAuth(expectedToken, r) {
if !checkAuth(expectedTokens, r) {
sendUnauthorized(w)
return
}
main.go
+17 -6
View File
@@ -27,14 +27,25 @@ func main() {
csvUpdateInterval := flag.Int("csv-update-interval", 10, "Interval in minutes with which to run the CSV updater command")
flag.Parse()
var authToken string
var authTokens []string
if *keyPath != "" {
content, err := os.ReadFile(*keyPath)
if err != nil {
log.Fatalf("Failed to read auth token file: %v", err)
}
authToken = strings.TrimSpace(string(content))
log.Printf("Starting Dynagate LLM Gateway. Config: %s, Port: %d, Auth: Required", *csvPath, *port)
var tokens []string
for _, line := range strings.Split(string(content), "\n") {
token := strings.TrimSpace(line)
if token != "" {
tokens = append(tokens, token)
}
}
if len(tokens) > 0 {
authTokens = tokens
log.Printf("Starting Dynagate LLM Gateway. Config: %s, Port: %d, Auth: Required (%d tokens loaded)", *csvPath, *port, len(authTokens))
} else {
log.Printf("Starting Dynagate LLM Gateway. Config: %s, Port: %d, Auth: None (empty key file)", *csvPath, *port)
}
} else {
log.Printf("Starting Dynagate LLM Gateway. Config: %s, Port: %d, Auth: None", *csvPath, *port)
}
@@ -62,9 +73,9 @@ func main() {
// Register routes
mux := http.NewServeMux()
mux.HandleFunc("/v1/models", handleModels(cm, authToken))
mux.HandleFunc("/v1/chat/completions", handleChatCompletions(cm, authToken))
mux.HandleFunc("/v1/images/generations", handleImageGenerations(cm, authToken))
mux.HandleFunc("/v1/models", handleModels(cm, authTokens))
mux.HandleFunc("/v1/chat/completions", handleChatCompletions(cm, authTokens))
mux.HandleFunc("/v1/images/generations", handleImageGenerations(cm, authTokens))
server := &http.Server{
Addr: fmt.Sprintf(":%d", *port),