diff options
Diffstat (limited to 'internal/auth')
| -rw-r--r-- | internal/auth/static.go | 124 | ||||
| -rw-r--r-- | internal/auth/static_test.go | 76 |
2 files changed, 200 insertions, 0 deletions
diff --git a/internal/auth/static.go b/internal/auth/static.go new file mode 100644 index 0000000..495f90c --- /dev/null +++ b/internal/auth/static.go @@ -0,0 +1,124 @@ +package auth + +import ( + "crypto/sha256" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "sync/atomic" + + "aigw/internal/domain" +) + +var ErrUnauthorized = errors.New("invalid or missing API key") + +type Authenticator interface { + Authenticate(*http.Request) (domain.Principal, error) +} + +type KeyRecord struct { + Key string `json:"key"` + KeyID string `json:"key_id"` + TenantID string `json:"tenant_id"` + ProjectID string `json:"project_id"` + Scopes []string `json:"scopes"` +} + +type StaticAuthenticator struct { + state atomic.Pointer[keySnapshot] + allowAnonymous bool +} + +type keySnapshot struct { + keys map[[sha256.Size]byte]domain.Principal +} + +type HashedKeyRecord struct { + Hash [sha256.Size]byte + Principal domain.Principal +} + +func NewStatic(raw string, allowAnonymous bool) (*StaticAuthenticator, error) { + result := &StaticAuthenticator{allowAnonymous: allowAnonymous} + if strings.TrimSpace(raw) == "" { + if allowAnonymous { + result.ReplaceHashed(nil) + return result, nil + } + return nil, errors.New("client API key environment variable is empty") + } + + var records []KeyRecord + if err := json.Unmarshal([]byte(raw), &records); err != nil { + return nil, fmt.Errorf("parse client API keys JSON: %w", err) + } + hashed := make([]HashedKeyRecord, 0, len(records)) + seen := make(map[[sha256.Size]byte]struct{}, len(records)) + for i, record := range records { + if record.Key == "" || record.KeyID == "" || record.TenantID == "" || record.ProjectID == "" { + return nil, fmt.Errorf("client API key record %d requires key, key_id, tenant_id, and project_id", i) + } + hash := sha256.Sum256([]byte(record.Key)) + if _, exists := seen[hash]; exists { + return nil, fmt.Errorf("duplicate client API key at record %d", i) + } + seen[hash] = struct{}{} + hashed = append(hashed, HashedKeyRecord{Hash: hash, Principal: domain.Principal{ + KeyID: record.KeyID, TenantID: record.TenantID, ProjectID: record.ProjectID, + Scopes: append([]string(nil), record.Scopes...), + }}) + } + if len(hashed) == 0 && !allowAnonymous { + return nil, errors.New("at least one client API key is required") + } + result.ReplaceHashed(hashed) + return result, nil +} + +func NewDynamic(records []HashedKeyRecord, allowAnonymous bool) *StaticAuthenticator { + result := &StaticAuthenticator{allowAnonymous: allowAnonymous} + result.ReplaceHashed(records) + return result +} + +func (a *StaticAuthenticator) ReplaceHashed(records []HashedKeyRecord) { + keys := make(map[[sha256.Size]byte]domain.Principal, len(records)) + for _, record := range records { + principal := record.Principal + principal.Scopes = append([]string(nil), principal.Scopes...) + keys[record.Hash] = principal + } + a.state.Store(&keySnapshot{keys: keys}) +} + +func (a *StaticAuthenticator) Authenticate(r *http.Request) (domain.Principal, error) { + key := bearerToken(r.Header.Get("Authorization")) + if key == "" { + key = strings.TrimSpace(r.Header.Get("x-api-key")) + } + if key == "" && a.allowAnonymous { + return domain.Principal{KeyID: "anonymous", TenantID: "anonymous", ProjectID: "anonymous"}, nil + } + if key == "" { + return domain.Principal{}, ErrUnauthorized + } + snapshot := a.state.Load() + if snapshot == nil { + return domain.Principal{}, ErrUnauthorized + } + principal, ok := snapshot.keys[sha256.Sum256([]byte(key))] + if !ok { + return domain.Principal{}, ErrUnauthorized + } + return principal, nil +} + +func bearerToken(header string) string { + parts := strings.Fields(header) + if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") { + return "" + } + return parts[1] +} diff --git a/internal/auth/static_test.go b/internal/auth/static_test.go new file mode 100644 index 0000000..cf54ba6 --- /dev/null +++ b/internal/auth/static_test.go @@ -0,0 +1,76 @@ +package auth + +import ( + "crypto/sha256" + "net/http" + "testing" + + "aigw/internal/domain" +) + +func TestStaticAuthenticator(t *testing.T) { + raw := `[{"key":"sk-client","key_id":"key-1","tenant_id":"tenant-1","project_id":"project-1","scopes":["inference"]}]` + authenticator, err := NewStatic(raw, false) + if err != nil { + t.Fatal(err) + } + + request, _ := http.NewRequest(http.MethodGet, "http://gateway.test/v1/models", nil) + request.Header.Set("Authorization", "Bearer sk-client") + principal, err := authenticator.Authenticate(request) + if err != nil { + t.Fatal(err) + } + if principal.TenantID != "tenant-1" || principal.ProjectID != "project-1" || principal.KeyID != "key-1" { + t.Fatalf("unexpected principal: %+v", principal) + } + + request.Header.Set("Authorization", "Bearer wrong") + if _, err := authenticator.Authenticate(request); err != ErrUnauthorized { + t.Fatalf("expected ErrUnauthorized, got %v", err) + } +} + +func TestDynamicAuthenticatorReplacesSnapshot(t *testing.T) { + oldHash := sha256.Sum256([]byte("sk-old")) + newHash := sha256.Sum256([]byte("sk-new")) + authenticator := NewDynamic([]HashedKeyRecord{{ + Hash: oldHash, + Principal: domain.Principal{ + KeyID: "old", TenantID: "tenant-1", ProjectID: "project-1", Scopes: []string{"inference"}, + }, + }}, false) + + request, _ := http.NewRequest(http.MethodGet, "http://gateway.test/v1/models", nil) + request.Header.Set("Authorization", "Bearer sk-old") + if _, err := authenticator.Authenticate(request); err != nil { + t.Fatalf("old key should initially authenticate: %v", err) + } + + authenticator.ReplaceHashed([]HashedKeyRecord{{ + Hash: newHash, + Principal: domain.Principal{ + KeyID: "new", TenantID: "tenant-1", ProjectID: "project-1", Scopes: []string{"inference"}, + }, + }}) + if _, err := authenticator.Authenticate(request); err != ErrUnauthorized { + t.Fatalf("old key remained valid after snapshot replacement: %v", err) + } + request.Header.Set("Authorization", "Bearer sk-new") + principal, err := authenticator.Authenticate(request) + if err != nil || principal.KeyID != "new" { + t.Fatalf("new key was not loaded: principal=%+v err=%v", principal, err) + } +} + +func TestStaticAuthenticatorAcceptsAnthropicHeader(t *testing.T) { + authenticator, err := NewStatic(`[{"key":"sk-client","key_id":"key-1","tenant_id":"tenant-1","project_id":"project-1"}]`, false) + if err != nil { + t.Fatal(err) + } + request, _ := http.NewRequest(http.MethodGet, "http://gateway.test/anthropic/v1/models", nil) + request.Header.Set("x-api-key", "sk-client") + if _, err := authenticator.Authenticate(request); err != nil { + t.Fatal(err) + } +} |
