1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
|
package auth
import (
"crypto/sha256"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"sync/atomic"
"time"
"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"`
AllowedModels []string `json:"allowed_models,omitempty"`
MonthlySpendMicros int64 `json:"monthly_spend_micros,omitempty"`
ExpiresAt *time.Time `json:"expires_at,omitempty"`
}
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{}{}
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...), AllowedModels: allowedModels,
MonthlySpendMicros: record.MonthlySpendMicros, ExpiresAt: record.ExpiresAt,
}})
}
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...)
principal.AllowedModels = cloneSet(principal.AllowedModels)
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 || (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") {
return ""
}
return parts[1]
}
|
