diff options
Diffstat (limited to 'internal/auth/static.go')
| -rw-r--r-- | internal/auth/static.go | 37 |
1 files changed, 30 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") { |