diff options
Diffstat (limited to '')
| -rw-r--r-- | internal/auth/static.go | 37 | ||||
| -rw-r--r-- | internal/auth/static_test.go | 31 |
2 files changed, 61 insertions, 7 deletions
diff --git a/internal/auth/static.go b/internal/auth/static.go index 495f90c..2b44158 100644 --- a/internal/auth/static.go +++ b/internal/auth/static.go @@ -8,6 +8,7 @@ import ( "net/http" "strings" "sync/atomic" + "time" "aigw/internal/domain" ) @@ -19,11 +20,14 @@ type Authenticator interface { } 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"` + Key string `json:"key"` + KeyID string `json:"key_id"` + TenantID string `json:"tenant_id"` + ProjectID string `json:"project_id"` + Scopes []string `json:"scopes"` + AllowedModels []string `json:"allowed_models,omitempty"` + MonthlySpendMicros int64 `json:"monthly_spend_micros,omitempty"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` } type StaticAuthenticator struct { @@ -65,9 +69,16 @@ func NewStatic(raw string, allowAnonymous bool) (*StaticAuthenticator, error) { return nil, fmt.Errorf("duplicate client API key at record %d", i) } seen[hash] = struct{}{} + allowedModels := make(map[string]struct{}, len(record.AllowedModels)) + for _, model := range record.AllowedModels { + if model = strings.TrimSpace(model); model != "" { + allowedModels[model] = 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...), + Scopes: append([]string(nil), record.Scopes...), AllowedModels: allowedModels, + MonthlySpendMicros: record.MonthlySpendMicros, ExpiresAt: record.ExpiresAt, }}) } if len(hashed) == 0 && !allowAnonymous { @@ -88,6 +99,7 @@ func (a *StaticAuthenticator) ReplaceHashed(records []HashedKeyRecord) { for _, record := range records { principal := record.Principal principal.Scopes = append([]string(nil), principal.Scopes...) + principal.AllowedModels = cloneSet(principal.AllowedModels) keys[record.Hash] = principal } a.state.Store(&keySnapshot{keys: keys}) @@ -109,12 +121,23 @@ func (a *StaticAuthenticator) Authenticate(r *http.Request) (domain.Principal, e return domain.Principal{}, ErrUnauthorized } principal, ok := snapshot.keys[sha256.Sum256([]byte(key))] - if !ok { + if !ok || (principal.ExpiresAt != nil && !principal.ExpiresAt.After(time.Now())) { return domain.Principal{}, ErrUnauthorized } return principal, nil } +func cloneSet(source map[string]struct{}) map[string]struct{} { + if len(source) == 0 { + return nil + } + result := make(map[string]struct{}, len(source)) + for key := range source { + result[key] = struct{}{} + } + return result +} + func bearerToken(header string) string { parts := strings.Fields(header) if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") { diff --git a/internal/auth/static_test.go b/internal/auth/static_test.go index cf54ba6..a037ce1 100644 --- a/internal/auth/static_test.go +++ b/internal/auth/static_test.go @@ -4,6 +4,7 @@ import ( "crypto/sha256" "net/http" "testing" + "time" "aigw/internal/domain" ) @@ -74,3 +75,33 @@ func TestStaticAuthenticatorAcceptsAnthropicHeader(t *testing.T) { t.Fatal(err) } } + +func TestStaticAuthenticatorLoadsRestrictionsAndRejectsExpiredKey(t *testing.T) { + future := time.Now().Add(time.Hour).UTC().Format(time.RFC3339Nano) + authenticator, err := NewStatic(`[{"key":"sk-limited","key_id":"key-1","tenant_id":"tenant-1","project_id":"project-1","allowed_models":["model/allowed"],"monthly_spend_micros":1250000,"expires_at":"`+future+`"}]`, false) + if err != nil { + t.Fatal(err) + } + request, _ := http.NewRequest(http.MethodGet, "http://gateway.test/v1/models", nil) + request.Header.Set("Authorization", "Bearer sk-limited") + principal, err := authenticator.Authenticate(request) + if err != nil { + t.Fatal(err) + } + if principal.MonthlySpendMicros != 1_250_000 { + t.Fatalf("monthly spend limit = %d", principal.MonthlySpendMicros) + } + if _, ok := principal.AllowedModels["model/allowed"]; !ok { + t.Fatalf("allowed model was not loaded: %+v", principal.AllowedModels) + } + + past := time.Now().Add(-time.Minute) + hash := sha256.Sum256([]byte("sk-expired")) + authenticator.ReplaceHashed([]HashedKeyRecord{{Hash: hash, Principal: domain.Principal{ + KeyID: "expired", TenantID: "tenant-1", ProjectID: "project-1", ExpiresAt: &past, + }}}) + request.Header.Set("Authorization", "Bearer sk-expired") + if _, err := authenticator.Authenticate(request); err != ErrUnauthorized { + t.Fatalf("expired key error = %v, want ErrUnauthorized", err) + } +} |