diff --git a/deploy/examples/server.yaml b/deploy/examples/server.yaml index e2668dba..563a0856 100644 --- a/deploy/examples/server.yaml +++ b/deploy/examples/server.yaml @@ -36,6 +36,28 @@ agent_listen: ":7443" # cert: /etc/llmgw/server.pem # key: /etc/llmgw/server-key.pem +# API keys on the client-facing port (plan §7.9). Off by default, because a +# gateway on a trusted network with no keys issued still has to serve — but +# nothing else protects that port, so turning it on is what makes it safe to +# expose. +# +# llmgw keys init --store /var/lib/llmgw/llmgw.db +# llmgw keys create --store /var/lib/llmgw/llmgw.db --project default "laptop" +# +# The key is printed once. Only its hash is stored, so there is nothing to +# recover it from. +# +# auth: +# enabled: true +# store: /var/lib/llmgw/llmgw.db +# # How long a resolved key is trusted before being read again. This is the +# # price of keeping the database off the request path: it bounds how long a +# # revoked key keeps working. +# cache_ttl: 30s +# # How long a rejected key is remembered as rejected, so a client looping +# # with a bad key cannot become a database read per request. +# negative_ttl: 5s + router: scrape_interval: 2s scrape_timeout: 1s @@ -75,9 +97,9 @@ router: # traffic, because an engine queue is FIFO and knows nothing about tenants. admission: enabled: true - # Until authentication lands, the band comes from a header and is trusted. - # Do not expose this port to callers who should not be able to promote - # themselves. + # With auth on, the key decides the band and this header is ignored. With + # auth off the header is trusted, so any client can promote itself into the + # top band — do not expose the port in that configuration. band_header: x-llmgw-band default_band: gold # A band waits once pool saturation passes its ceiling, so the cluster diff --git a/internal/auth/auth.go b/internal/auth/auth.go new file mode 100644 index 00000000..161da154 --- /dev/null +++ b/internal/auth/auth.go @@ -0,0 +1,267 @@ +// Package auth authenticates client requests against the key store (plan §7.9). +// +// The rule that shapes everything here: authentication runs on every request, +// and the request path does not touch the database. A key is looked up once, +// resolved against its tenancy chain once, and then answered from memory until +// the entry expires or is invalidated. +package auth + +import ( + "context" + "errors" + "net/http" + "strings" + "sync" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// Principal is who a request is, once authenticated. +type Principal struct { + KeyID string + KeyPrefix string + ProjectID string + TeamID string + OrgID string + + // Band and RoutingProfile come from the key, and are authoritative: with + // authentication on, a client can no longer promote itself with a header. + Band string + RoutingProfile string + + Limits store.Effective +} + +// Allows reports whether this principal may use a model. +func (p Principal) Allows(model string) bool { return p.Limits.Allows(model) } + +type contextKey struct{} + +// WithPrincipal attaches an authenticated identity to a context. +func WithPrincipal(ctx context.Context, p Principal) context.Context { + return context.WithValue(ctx, contextKey{}, p) +} + +// FromContext returns the authenticated identity, if there is one. +// +// Absent means authentication is off, which is a deployment choice rather than +// a failure: a gateway on a trusted network with no keys configured still +// serves. Callers must therefore treat the missing case as "unrestricted", and +// the middleware is what guarantees it cannot mean "unchecked". +func FromContext(ctx context.Context) (Principal, bool) { + p, ok := ctx.Value(contextKey{}).(Principal) + return p, ok +} + +// Config tunes the authenticator. +type Config struct { + // TTL is how long a resolved key is trusted before being read again. It + // bounds how stale a revocation can be, so it trades a database read per + // key per TTL against how long a revoked key keeps working. + TTL time.Duration + // NegativeTTL is how long a rejected key is remembered as rejected, so a + // client looping with a bad key cannot turn into a database read per + // request. + NegativeTTL time.Duration + // TouchInterval is how often a key's last-used time is written, at most. + TouchInterval time.Duration +} + +// DefaultConfig returns sensible values. +func DefaultConfig() Config { + return Config{TTL: 30 * time.Second, NegativeTTL: 5 * time.Second, TouchInterval: time.Minute} +} + +func (c *Config) normalize() { + d := DefaultConfig() + if c.TTL <= 0 { + c.TTL = d.TTL + } + if c.NegativeTTL <= 0 { + c.NegativeTTL = d.NegativeTTL + } + if c.TouchInterval <= 0 { + c.TouchInterval = d.TouchInterval + } +} + +// Errors an authenticator returns. They are distinguished because the status +// codes differ: a missing credential is 401, a forbidden model is 403. +var ( + ErrMissing = errors.New("no api key") + ErrInvalid = errors.New("invalid api key") + ErrExpired = errors.New("api key expired") + ErrRevoked = errors.New("api key revoked") +) + +// Authenticator resolves keys and caches the result. +type Authenticator struct { + store store.Store + cfg Config + now func() time.Time + + mu sync.RWMutex + cache map[string]entry + + // touch batches last-used timestamps; nothing reads the field in real time + // and writing it per request would put a database write back on the path. + touchMu sync.Mutex + touch map[string]time.Time +} + +type entry struct { + principal Principal + err error + expires time.Time +} + +// New builds an authenticator over a store. +func New(s store.Store, cfg Config) *Authenticator { + cfg.normalize() + return &Authenticator{ + store: s, cfg: cfg, now: time.Now, + cache: map[string]entry{}, + touch: map[string]time.Time{}, + } +} + +// Authenticate resolves a presented credential. +func (a *Authenticator) Authenticate(ctx context.Context, presented string) (Principal, error) { + if presented == "" { + return Principal{}, ErrMissing + } + // Shape first: a bearer token from another system is rejected without a + // lookup, and never reaches the store. + if err := store.ParseKey(presented); err != nil { + return Principal{}, ErrInvalid + } + + hash := store.HashKey(presented) + id := string(hash) + + if e, ok := a.lookup(id); ok { + if e.err == nil { + a.markUsed(e.principal.KeyID) + } + return e.principal, e.err + } + + p, err := a.resolve(ctx, hash) + if err != nil && !isRejection(err) { + // A database problem is not an authentication answer, so it must not be + // cached as one. Failing the request is right — the alternative is + // serving traffic whose authorization is unknown. + return Principal{}, err + } + + ttl := a.cfg.TTL + if err != nil { + ttl = a.cfg.NegativeTTL + } + a.store_(id, entry{principal: p, err: err, expires: a.now().Add(ttl)}) + if err == nil { + a.markUsed(p.KeyID) + } + return p, err +} + +func (a *Authenticator) resolve(ctx context.Context, hash []byte) (Principal, error) { + k, err := a.store.KeyByHash(ctx, hash) + if errors.Is(err, store.ErrNotFound) { + return Principal{}, ErrInvalid + } + if err != nil { + return Principal{}, err + } + switch { + case !k.RevokedAt.IsZero(): + return Principal{}, ErrRevoked + case !k.ExpiresAt.IsZero() && !a.now().Before(k.ExpiresAt): + return Principal{}, ErrExpired + } + + eff, err := a.store.Resolve(ctx, k.ID) + if err != nil { + return Principal{}, err + } + return Principal{ + KeyID: k.ID, KeyPrefix: k.Prefix, ProjectID: eff.ProjectID, + TeamID: eff.TeamID, OrgID: eff.OrgID, + Band: eff.Band, RoutingProfile: eff.RoutingProfile, + Limits: eff, + }, nil +} + +func (a *Authenticator) lookup(id string) (entry, bool) { + a.mu.RLock() + e, ok := a.cache[id] + a.mu.RUnlock() + if !ok || a.now().After(e.expires) { + return entry{}, false + } + return e, true +} + +func (a *Authenticator) store_(id string, e entry) { + a.mu.Lock() + a.cache[id] = e + a.mu.Unlock() +} + +// Invalidate drops every cached decision. +// +// Coarse on purpose: a limit change anywhere in the tree can change any key +// below it, and working out which keys those are costs more than re-resolving +// the few that are actually in use. +func (a *Authenticator) Invalidate() { + a.mu.Lock() + a.cache = map[string]entry{} + a.mu.Unlock() +} + +func (a *Authenticator) markUsed(keyID string) { + if keyID == "" { + return + } + now := a.now() + a.touchMu.Lock() + if last, ok := a.touch[keyID]; !ok || now.Sub(last) >= a.cfg.TouchInterval { + a.touch[keyID] = now + } + a.touchMu.Unlock() +} + +// FlushTouches writes the pending last-used timestamps. The batch writer calls +// it on its own schedule. +func (a *Authenticator) FlushTouches(ctx context.Context) error { + a.touchMu.Lock() + pending := a.touch + a.touch = map[string]time.Time{} + a.touchMu.Unlock() + + if len(pending) == 0 { + return nil + } + return a.store.TouchKeys(ctx, pending) +} + +func isRejection(err error) bool { + return errors.Is(err, ErrInvalid) || errors.Is(err, ErrExpired) || + errors.Is(err, ErrRevoked) || errors.Is(err, ErrMissing) +} + +// BearerToken extracts the credential from a request. +// +// Both spellings are accepted because both are in the wild: the OpenAI SDKs +// send Authorization, and some tools only manage to set an api-key header. +func BearerToken(r *http.Request) string { + if h := r.Header.Get("Authorization"); h != "" { + if token, ok := strings.CutPrefix(h, "Bearer "); ok { + return strings.TrimSpace(token) + } + // A bare value is common enough in hand-written clients to accept. + return strings.TrimSpace(h) + } + return strings.TrimSpace(r.Header.Get("api-key")) +} diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go new file mode 100644 index 00000000..38fdd64a --- /dev/null +++ b/internal/auth/auth_test.go @@ -0,0 +1,382 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "path/filepath" + "sync" + "sync/atomic" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// fixture builds a store with one org/team/project and one key. +func fixture(t *testing.T) (*store.SQLite, string, string) { + t.Helper() + ctx := context.Background() + s, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "t.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = s.Close() }) + + must(t, s.PutOrg(ctx, store.Org{ID: "org1", Name: "Acme"})) + must(t, s.PutTeam(ctx, store.Team{ID: "team1", OrgID: "org1", Name: "R"})) + must(t, s.PutProject(ctx, store.Project{ID: "proj1", TeamID: "team1", Name: "P"})) + + secret, k, err := store.NewKey("key1", "laptop", "proj1") + if err != nil { + t.Fatal(err) + } + k.CreatedAt = time.Now() + k.Band = "gold" + must(t, s.CreateKey(ctx, k)) + return s, secret, k.ID +} + +func must(t *testing.T, err error) { + t.Helper() + if err != nil { + t.Fatal(err) + } +} + +func TestAValidKeyResolvesItsWholeChain(t *testing.T) { + s, secret, keyID := fixture(t) + a := New(s, DefaultConfig()) + + p, err := a.Authenticate(context.Background(), secret) + if err != nil { + t.Fatal(err) + } + if p.KeyID != keyID || p.ProjectID != "proj1" || p.TeamID != "team1" || p.OrgID != "org1" { + t.Errorf("principal %+v", p) + } + if p.Band != "gold" { + t.Errorf("band %q", p.Band) + } +} + +func TestRejections(t *testing.T) { + s, secret, keyID := fixture(t) + ctx := context.Background() + a := New(s, DefaultConfig()) + + if _, err := a.Authenticate(ctx, ""); !errors.Is(err, ErrMissing) { + t.Errorf("no credential: %v", err) + } + // A bearer token from some other system must be rejected on shape alone. + if _, err := a.Authenticate(ctx, "ghp_0123456789abcdefghij"); !errors.Is(err, ErrInvalid) { + t.Errorf("foreign token: %v", err) + } + if _, err := a.Authenticate(ctx, secret[:len(secret)-1]+"X"); !errors.Is(err, ErrInvalid) { + t.Errorf("wrong key: %v", err) + } + + must(t, s.RevokeKey(ctx, keyID, time.Now())) + a.Invalidate() + if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrRevoked) { + t.Errorf("revoked key: %v", err) + } +} + +// countingStore counts how often the database is actually read. +type countingStore struct { + store.Store + lookups atomic.Int64 + fail atomic.Bool +} + +func (c *countingStore) KeyByHash(ctx context.Context, hash []byte) (store.Key, error) { + c.lookups.Add(1) + if c.fail.Load() { + return store.Key{}, errors.New("database is down") + } + return c.Store.KeyByHash(ctx, hash) +} + +func TestTheRequestPathDoesNotReadTheDatabase(t *testing.T) { + // The rule the whole package is shaped around. Without the cache this is a + // database read per request, on every request. + s, secret, _ := fixture(t) + c := &countingStore{Store: s} + a := New(c, DefaultConfig()) + + for range 100 { + if _, err := a.Authenticate(context.Background(), secret); err != nil { + t.Fatal(err) + } + } + if n := c.lookups.Load(); n != 1 { + t.Errorf("100 requests caused %d database reads", n) + } +} + +func TestARejectedKeyIsNotADatabaseReadPerRequest(t *testing.T) { + // A client looping with a bad key would otherwise be a free way to load the + // database from outside. + s, _, _ := fixture(t) + c := &countingStore{Store: s} + a := New(c, DefaultConfig()) + + bad := store.KeyPrefix + "0000000000000000000000000000000000000000000" + for range 50 { + if _, err := a.Authenticate(context.Background(), bad); !errors.Is(err, ErrInvalid) { + t.Fatal(err) + } + } + if n := c.lookups.Load(); n != 1 { + t.Errorf("50 bad requests caused %d database reads", n) + } +} + +func TestCachedDecisionsExpire(t *testing.T) { + // The cache bounds how long a revocation can take to take effect, so it has + // to actually expire. + s, secret, keyID := fixture(t) + ctx := context.Background() + + now := time.Now() + a := New(s, Config{TTL: time.Second, NegativeTTL: time.Second}) + a.now = func() time.Time { return now } + + if _, err := a.Authenticate(ctx, secret); err != nil { + t.Fatal(err) + } + must(t, s.RevokeKey(ctx, keyID, now)) + + // Still inside the TTL: the stale answer is the point of the cache. + if _, err := a.Authenticate(ctx, secret); err != nil { + t.Errorf("the cache expired early: %v", err) + } + + now = now.Add(2 * time.Second) + if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrRevoked) { + t.Errorf("a revoked key still works after the TTL: %v", err) + } +} + +func TestAStoreFailureIsNotAnAuthenticationAnswer(t *testing.T) { + // Two things must not happen: the failure must not be cached as a + // rejection, and it must not read to the caller as a bad key. + s, secret, _ := fixture(t) + c := &countingStore{Store: s} + a := New(c, DefaultConfig()) + + c.fail.Store(true) + _, err := a.Authenticate(context.Background(), secret) + if err == nil { + t.Fatal("a broken database authenticated a request") + } + if isRejection(err) { + t.Errorf("a database failure was reported as a credential problem: %v", err) + } + + c.fail.Store(false) + if _, err := a.Authenticate(context.Background(), secret); err != nil { + t.Errorf("the failure was cached: %v", err) + } +} + +func TestExpiryIsNotCached(t *testing.T) { + ctx := context.Background() + s, _, _ := fixture(t) + + secret, k, _ := store.NewKey("key2", "temp", "proj1") + k.CreatedAt = time.Now() + k.ExpiresAt = time.Now().Add(time.Hour) + must(t, s.CreateKey(ctx, k)) + + now := time.Now() + a := New(s, Config{TTL: time.Minute}) + a.now = func() time.Time { return now } + + if _, err := a.Authenticate(ctx, secret); err != nil { + t.Fatal(err) + } + // Past the key's expiry but inside the cache TTL. The cached decision is + // allowed to be stale by at most the TTL, which is the documented trade. + now = now.Add(2 * time.Hour) + if _, err := a.Authenticate(ctx, secret); !errors.Is(err, ErrExpired) { + t.Errorf("an expired key authenticated: %v", err) + } +} + +func TestConcurrentAuthentication(t *testing.T) { + s, secret, _ := fixture(t) + a := New(s, DefaultConfig()) + + var wg sync.WaitGroup + errs := make(chan error, 64) + for range 64 { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := a.Authenticate(context.Background(), secret); err != nil { + errs <- err + } + }() + } + wg.Wait() + close(errs) + for err := range errs { + t.Fatal(err) + } +} + +func TestMiddlewareGuardsEveryRouteButTheExemptOnes(t *testing.T) { + s, secret, _ := fixture(t) + a := New(s, DefaultConfig()) + + var reached atomic.Int64 + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + reached.Add(1) + if _, ok := FromContext(r.Context()); !ok && r.URL.Path != "/health" { + t.Error("an authenticated request carried no principal") + } + w.WriteHeader(http.StatusOK) + }) + srv := httptest.NewServer(Middleware(a, nil, "/health")(inner)) + defer srv.Close() + + get := func(path, token string) int { + req, _ := http.NewRequest(http.MethodGet, srv.URL+path, nil) //nolint:noctx // test + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + return resp.StatusCode + } + + if code := get("/health", ""); code != http.StatusOK { + t.Errorf("health needs no key, got %d", code) + } + if code := get("/v1/models", ""); code != http.StatusUnauthorized { + t.Errorf("an unauthenticated request got %d", code) + } + if code := get("/v1/models", "sk-llmgw-bogus"); code != http.StatusUnauthorized { + t.Errorf("a bad key got %d", code) + } + if code := get("/v1/models", secret); code != http.StatusOK { + t.Errorf("a good key got %d", code) + } + if n := reached.Load(); n != 2 { + t.Errorf("the handler was reached %d times; only health and the valid key should pass", n) + } +} + +func TestARevokedKeyLooksExactlyLikeAnUnknownOne(t *testing.T) { + // Telling a caller that a key exists but was revoked confirms the key was + // real, which is information worth withholding. + s, secret, keyID := fixture(t) + must(t, s.RevokeKey(context.Background(), keyID, time.Now())) + + a := New(s, DefaultConfig()) + srv := httptest.NewServer(Middleware(a, nil)(http.HandlerFunc( + func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }))) + defer srv.Close() + + body := func(token string) string { + req, _ := http.NewRequest(http.MethodGet, srv.URL, nil) //nolint:noctx // test + req.Header.Set("Authorization", "Bearer "+token) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + var e struct { + Error struct{ Message string } `json:"error"` + } + _ = json.NewDecoder(resp.Body).Decode(&e) + return e.Error.Message + } + + unknown := store.KeyPrefix + "0000000000000000000000000000000000000000000" + if body(secret) != body(unknown) { + t.Errorf("revoked says %q, unknown says %q", body(secret), body(unknown)) + } +} + +func TestAStoreFailureDoesNotReadAsABadKey(t *testing.T) { + s, secret, _ := fixture(t) + c := &countingStore{Store: s} + c.fail.Store(true) + a := New(c, DefaultConfig()) + + srv := httptest.NewServer(Middleware(a, nil)(http.HandlerFunc( + func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }))) + defer srv.Close() + + req, _ := http.NewRequest(http.MethodGet, srv.URL, nil) //nolint:noctx // test + req.Header.Set("Authorization", "Bearer "+secret) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + // 503 rather than 401: sending someone hunting for a key problem that does + // not exist wastes the one thing an incident is short of. + if resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("status %d", resp.StatusCode) + } +} + +func TestBearerTokenAcceptsWhatClientsActuallySend(t *testing.T) { + cases := map[string]string{ + "Bearer sk-llmgw-abc": "sk-llmgw-abc", + "sk-llmgw-abc": "sk-llmgw-abc", + "Bearer sk-llmgw-abc": "sk-llmgw-abc", + "": "", + } + for header, want := range cases { + r := httptest.NewRequest(http.MethodGet, "/", nil) + if header != "" { + r.Header.Set("Authorization", header) + } + if got := BearerToken(r); got != want { + t.Errorf("Authorization %q -> %q, want %q", header, got, want) + } + } + + // Some tools only manage to set this one. + r := httptest.NewRequest(http.MethodGet, "/", nil) + r.Header.Set("api-key", "sk-llmgw-xyz") + if got := BearerToken(r); got != "sk-llmgw-xyz" { + t.Errorf("api-key header -> %q", got) + } +} + +func TestLastUsedIsBatchedNotWrittenPerRequest(t *testing.T) { + s, secret, keyID := fixture(t) + ctx := context.Background() + a := New(s, Config{TTL: time.Minute, TouchInterval: time.Hour}) + + for range 20 { + if _, err := a.Authenticate(ctx, secret); err != nil { + t.Fatal(err) + } + } + // Nothing written yet: the field is not read in real time. + k, _ := s.KeyByHash(ctx, store.HashKey(secret)) + if !k.LastUsedAt.IsZero() { + t.Error("last-used was written on the request path") + } + + must(t, a.FlushTouches(ctx)) + k, _ = s.KeyByHash(ctx, store.HashKey(secret)) + if k.LastUsedAt.IsZero() { + t.Error("last-used was never written") + } + if k.ID != keyID { + t.Errorf("wrong key touched: %s", k.ID) + } +} diff --git a/internal/auth/middleware.go b/internal/auth/middleware.go new file mode 100644 index 00000000..fe050c1e --- /dev/null +++ b/internal/auth/middleware.go @@ -0,0 +1,97 @@ +package auth + +import ( + "encoding/json" + "errors" + "log/slog" + "net/http" + "strings" +) + +// Middleware authenticates every request before it reaches the router. +// +// It is a wrapper rather than a check inside the handler so there is no path +// into the API that forgets to call it. Adding a route cannot accidentally add +// an unauthenticated one. +func Middleware(a *Authenticator, log *slog.Logger, exempt ...string) func(http.Handler) http.Handler { + if log == nil { + log = slog.Default() + } + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if isExempt(r.URL.Path, exempt) { + next.ServeHTTP(w, r) + return + } + + p, err := a.Authenticate(r.Context(), BearerToken(r)) + if err != nil { + writeAuthError(w, err) + // The key prefix, never the key: a rejected credential in a log + // is still a credential. + log.Debug("rejected a request", "path", r.URL.Path, + "reason", err, "remote", r.RemoteAddr) + return + } + next.ServeHTTP(w, r.WithContext(WithPrincipal(r.Context(), p))) + }) + } +} + +// isExempt reports whether a path is served without a key. +// +// Health is the only one by default. Metrics live on the admin listener, which +// is a separate port precisely so it can be firewalled rather than +// authenticated. +func isExempt(path string, exempt []string) bool { + for _, e := range exempt { + if path == e || (strings.HasSuffix(e, "/") && strings.HasPrefix(path, e)) { + return true + } + } + return false +} + +func writeAuthError(w http.ResponseWriter, err error) { + status, kind, msg := http.StatusUnauthorized, "invalid_request_error", "invalid api key" + switch { + case errors.Is(err, ErrMissing): + msg = "no api key was provided" + case errors.Is(err, ErrExpired): + msg = "this api key has expired" + case errors.Is(err, ErrRevoked): + // Deliberately the same message as an unknown key. Telling a caller + // that a key exists but was revoked confirms the key was real. + msg = "invalid api key" + case errors.Is(err, ErrInvalid): + default: + // A store failure is ours, not the caller's, and must not read as a + // rejected credential — that would send someone hunting for a key + // problem that does not exist. + status, kind, msg = http.StatusServiceUnavailable, "server_error", + "authentication is temporarily unavailable" + } + + if status == http.StatusUnauthorized { + // What the OpenAI SDKs look for to stop retrying. + w.Header().Set("WWW-Authenticate", `Bearer realm="llmgw"`) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{ + "message": msg, "type": kind, "code": nil, + }}) +} + +// WriteForbiddenModel reports a model this key may not use. +// +// Separate from authentication because it happens later: the model is only +// known once the body has been parsed. +func WriteForbiddenModel(w http.ResponseWriter, model string) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{ + "message": "this api key is not permitted to use model " + model, + "type": "invalid_request_error", "code": nil, + }}) +} diff --git a/internal/cli/keys.go b/internal/cli/keys.go new file mode 100644 index 00000000..03e2e7a8 --- /dev/null +++ b/internal/cli/keys.go @@ -0,0 +1,288 @@ +package cli + +import ( + "context" + "crypto/rand" + "encoding/base32" + "fmt" + "strings" + "text/tabwriter" + "time" + + "github.com/spf13/cobra" + + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +func newTenancyCmd() *cobra.Command { + var dbPath string + + cmd := &cobra.Command{ + Use: "keys", + Short: "Manage tenants and API keys", + Long: "Works directly on the store file rather than through the admin API, so keys\n" + + "can be issued before the server is running and without exposing an endpoint\n" + + "that mints credentials.", + } + cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store") + + open := func(ctx context.Context) (*store.SQLite, error) { return store.OpenSQLite(ctx, dbPath) } + + cmd.AddCommand( + newKeyInitCmd(open), + newKeyCreateCmd(open), + newKeyListCmd(open), + newKeyRevokeCmd(open), + newLimitsCmd(open), + ) + return cmd +} + +type opener func(context.Context) (*store.SQLite, error) + +func newKeyInitCmd(open opener) *cobra.Command { + var org, team, project string + cmd := &cobra.Command{ + Use: "init", + Short: "Create the tenancy tree a key hangs off", + Long: "Idempotent: running it again renames rather than duplicating, so it can be\n" + + "kept in whatever provisions the machine.", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + s, err := open(cmd.Context()) + if err != nil { + return err + } + defer s.Close() + + ctx := cmd.Context() + if err := s.PutOrg(ctx, store.Org{ID: slug(org), Name: org}); err != nil { + return err + } + if err := s.PutTeam(ctx, store.Team{ID: slug(team), OrgID: slug(org), Name: team}); err != nil { + return err + } + if err := s.PutProject(ctx, store.Project{ + ID: slug(project), TeamID: slug(team), Name: project, + }); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "%s / %s / %s\n", slug(org), slug(team), slug(project)) + return nil + }, + } + cmd.Flags().StringVar(&org, "org", "default", "organization name") + cmd.Flags().StringVar(&team, "team", "default", "team name") + cmd.Flags().StringVar(&project, "project", "default", "project name") + return cmd +} + +func newKeyCreateCmd(open opener) *cobra.Command { + var project, band, profile string + var expires time.Duration + + cmd := &cobra.Command{ + Use: "create ", + Short: "Issue an API key", + Long: "The key is printed once and never again: only its hash is stored, so there\n" + + "is nothing to recover it from. That is the point — a leaked database hands\n" + + "over no working credential.", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + s, err := open(cmd.Context()) + if err != nil { + return err + } + defer s.Close() + + secret, k, err := store.NewKey(randomID("key"), args[0], project) + if err != nil { + return err + } + k.CreatedAt = time.Now() + k.Band, k.RoutingProfile = band, profile + if expires > 0 { + k.ExpiresAt = k.CreatedAt.Add(expires) + } + if err := s.CreateKey(cmd.Context(), k); err != nil { + return fmt.Errorf("creating the key (does project %q exist? run `llmgw keys init`): %w", + project, err) + } + + // The secret goes to stdout and everything else to stderr, so + // `keys create ... > key.txt` captures the key and nothing else. + // Cobra's Print writes to stderr, which would have made the one + // output that matters the one that cannot be piped. + fmt.Fprintln(cmd.OutOrStdout(), secret) + + cmd.PrintErrf("\nid %s\nproject %s\n", k.ID, k.ProjectID) + if !k.ExpiresAt.IsZero() { + cmd.PrintErrf("expires %s\n", k.ExpiresAt.Format(time.RFC3339)) + } + cmd.PrintErrln("\nThis is the only time the key is shown. Store it now.") + return nil + }, + } + cmd.Flags().StringVar(&project, "project", "default", "project the key belongs to") + cmd.Flags().StringVar(&band, "band", "", "admission band; empty inherits") + cmd.Flags().StringVar(&profile, "profile", "", "routing profile; empty inherits") + cmd.Flags().DurationVar(&expires, "expires-in", 0, "lifetime; 0 never expires") + return cmd +} + +func newKeyListCmd(open opener) *cobra.Command { + var project string + cmd := &cobra.Command{ + Use: "list", + Short: "List a project's keys", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + s, err := open(cmd.Context()) + if err != nil { + return err + } + defer s.Close() + + keys, err := s.Keys(cmd.Context(), project) + if err != nil { + return err + } + if len(keys) == 0 { + fmt.Fprintf(cmd.OutOrStdout(), "no keys in project %q\n", project) + return nil + } + + w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0) + fmt.Fprintln(w, "ID\tNAME\tKEY\tBAND\tSTATE\tLAST USED") + now := time.Now() + for _, k := range keys { + fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n", + k.ID, k.Name, k.Prefix, dash(k.Band), keyState(k, now), ago(k.LastUsedAt, now)) + } + return w.Flush() + }, + } + cmd.Flags().StringVar(&project, "project", "default", "project to list") + return cmd +} + +func newKeyRevokeCmd(open opener) *cobra.Command { + return &cobra.Command{ + Use: "revoke ", + Short: "Stop a key working, without deleting it", + Long: "A soft delete. Usage rows reference the key, and an audit trail that loses\n" + + "its subject is not an audit trail.", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + s, err := open(cmd.Context()) + if err != nil { + return err + } + defer s.Close() + + if err := s.RevokeKey(cmd.Context(), args[0], time.Now()); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "revoked %s\n", args[0]) + cmd.PrintErrln("A running server keeps accepting it until its cache expires (auth.cache_ttl).") + return nil + }, + } +} + +func newLimitsCmd(open opener) *cobra.Command { + var scope, id, period string + var rpm, tpm int + var budget int64 + var models []string + + cmd := &cobra.Command{ + Use: "limits", + Short: "Set limits at one level of the tenancy tree", + Long: "Limits compose downward by taking the most restrictive value, so a key\n" + + "cannot raise what its organization set. Model allowlists intersect: a key\n" + + "cannot add a model its team withheld.", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + s, err := open(cmd.Context()) + if err != nil { + return err + } + defer s.Close() + + l := store.Limits{Scope: store.Scope(scope), ScopeID: id} + if cmd.Flags().Changed("rpm") { + l.RPM = &rpm + } + if cmd.Flags().Changed("tpm") { + l.TPM = &tpm + } + if cmd.Flags().Changed("budget") { + l.Budget = &budget + l.BudgetPeriod = store.Period(period) + } + if cmd.Flags().Changed("models") { + // An explicitly empty list is an allowlist that permits + // nothing, which is different from setting no list at all. + if models == nil { + models = []string{} + } + l.AllowedModels = models + } + if err := s.PutLimits(cmd.Context(), l); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "limits set on %s %s\n", scope, id) + return nil + }, + } + cmd.Flags().StringVar(&scope, "scope", "project", "org | team | project | key") + cmd.Flags().StringVar(&id, "id", "default", "id at that scope") + cmd.Flags().IntVar(&rpm, "rpm", 0, "requests per minute") + cmd.Flags().IntVar(&tpm, "tpm", 0, "tokens per minute") + cmd.Flags().Int64Var(&budget, "budget", 0, "budget per period, in micro-units") + cmd.Flags().StringVar(&period, "period", "month", "day | month") + cmd.Flags().StringSliceVar(&models, "models", nil, "allowed models; empty allows none") + return cmd +} + +func keyState(k store.Key, now time.Time) string { + switch { + case !k.RevokedAt.IsZero(): + return "revoked" + case !k.ExpiresAt.IsZero() && !now.Before(k.ExpiresAt): + return "expired" + default: + return "active" + } +} + +func ago(t, now time.Time) string { + if t.IsZero() { + return "never" + } + return now.Sub(t).Round(time.Second).String() + " ago" +} + +// slug turns a human name into an id. +func slug(name string) string { + var b strings.Builder + for _, r := range strings.ToLower(strings.TrimSpace(name)) { + switch { + case r >= 'a' && r <= 'z', r >= '0' && r <= '9': + b.WriteRune(r) + case r == ' ' || r == '-' || r == '_': + b.WriteByte('-') + } + } + if b.Len() == 0 { + return "default" + } + return b.String() +} + +func randomID(prefix string) string { + var buf [10]byte + _, _ = rand.Read(buf[:]) + return prefix + "-" + strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:])) +} diff --git a/internal/cli/root.go b/internal/cli/root.go index 3ebb4414..fc0e2788 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -40,6 +40,7 @@ func Root() *cobra.Command { newServerCmd(), newAgentCmd(), newCtlCmd(), + newTenancyCmd(), newSimEngineCmd(), ) return root diff --git a/internal/config/config.go b/internal/config/config.go index d92a3a03..a0663c2c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -32,6 +32,9 @@ type Server struct { // so the fleet-facing surface can be firewalled independently. AgentListen string `koanf:"agent_listen"` + // Auth turns on API keys on the client-facing port. + Auth Auth `koanf:"auth"` + // AgentTLS authenticates the fleet-facing port. It deliberately does not // cover the client-facing one: that surface is fronted by a publicly // trusted certificate, while this one is mutual and private. @@ -47,6 +50,23 @@ type Server struct { Endpoints []Endpoint `koanf:"endpoints"` } +// Auth configures client authentication (plan §7.9). +// +// Off by default, because a gateway on a trusted network with no keys issued +// still has to serve. Turning it on is what makes the client-facing port safe +// to expose. +type Auth struct { + Enabled bool `koanf:"enabled"` + // Store is the SQLite file holding keys, tenancy and usage. + Store string `koanf:"store"` + // CacheTTL bounds how long a revoked key keeps working, traded against a + // database read per key per TTL. + CacheTTL time.Duration `koanf:"cache_ttl"` + // NegativeTTL is how long a rejected key is remembered as rejected, so a + // client looping with a bad key cannot become a database read per request. + NegativeTTL time.Duration `koanf:"negative_ttl"` +} + // Control configures the control plane. type Control struct { // Interval is how often desired and actual state are compared. @@ -256,6 +276,9 @@ func (s Server) Validate() error { if err := s.AgentTLS.Validate(); err != nil { return err } + if s.Auth.Enabled && s.Auth.Store == "" { + return fmt.Errorf("auth.store is required when auth is enabled") + } seen := make(map[string]struct{}, len(s.Endpoints)) for i, e := range s.Endpoints { diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 9be63c11..ae01bae6 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -12,6 +12,7 @@ import ( "strings" "time" + "gitea.mynodes.xyz/ilker/llmgw/internal/auth" "gitea.mynodes.xyz/ilker/llmgw/internal/router/admission" "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" ) @@ -22,13 +23,18 @@ func (r *Router) admit(ctx context.Context, req *parsedRequest) (*admission.Leas if r.adm == nil { return nil, nil } + // The key is the tenant when there is one: fairness within a band has to be + // per caller, and a session id is something the caller picks, so one client + // could take a whole band by opening many sessions. Without authentication + // the session is still the best available grouping. + tenant := req.Tenant + if tenant == "" { + tenant = req.SessionID + } return r.adm.Acquire(ctx, admission.Request{ - Model: req.Model, - Band: req.Band, - // Until authentication arrives the session is the best tenant key we - // have: it groups one caller's traffic without pretending to identify - // them (plan §7.9). - Tenant: req.SessionID, + Model: req.Model, + Band: req.Band, + Tenant: tenant, }) } @@ -74,7 +80,13 @@ func (r *Router) Handler() http.Handler { mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }) - return mux + if r.auth == nil { + return mux + } + // Wrapping the whole mux rather than checking inside each handler: adding a + // route must not be able to add an unauthenticated one by omission. Health + // stays open because a load balancer has no key. + return auth.Middleware(r.auth, r.log, "/health")(mux) } // parsedRequest is what the router needs from a request. Everything else in the @@ -89,8 +101,12 @@ type parsedRequest struct { UpstreamBody []byte Header http.Header - SessionID string - Band string + SessionID string + Band string + // KeyID, ProjectID and Tenant are empty when authentication is off. + KeyID string + ProjectID string + Tenant string ProfileOverride string PromptPrefix string PromptTokens int @@ -127,6 +143,12 @@ func (r *Router) handleInference(path string) http.HandlerFunc { return } + if p, ok := auth.FromContext(req.Context()); ok && !r.permits(p, parsed.Model) { + auth.WriteForbiddenModel(w, parsed.Model) + r.m.observeRequest(parsed.Model, "", http.StatusForbidden, time.Since(start), nil, 0) + return + } + // Step 5 of the lifecycle: "now?" before "where?" (plan §7.1). lease, err := r.admit(req.Context(), parsed) if err != nil { @@ -145,6 +167,27 @@ func (r *Router) handleInference(path string) http.HandlerFunc { } } +// permits reports whether a key's allowlist covers the requested model. +// +// The decision is per model, not per string. A model's canonical name and every +// alias of it name the same thing, so the allowlist is matched against all of +// them: writing either form works, and no alias can reach a model that was not +// allowed. Matching only what the client sent would make every alias a way +// around the allowlist; matching only the canonical name would silently ignore +// an allowlist written in the names clients actually use. +func (r *Router) permits(p auth.Principal, requested string) bool { + canonical, m := r.resolve(requested) + if p.Allows(canonical) { + return true + } + for _, alias := range m.Aliases { + if p.Allows(alias) { + return true + } + } + return false +} + // finish records metrics and, when nothing was written, produces the error body. func (r *Router) finish(w http.ResponseWriter, req *parsedRequest, res dispatchResult, start time.Time) { status := res.outcome.Status @@ -203,6 +246,20 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) { MaxTokens: max(b.MaxTokens, b.MaxCompletionTokens), PromptPrefix: promptPrefix(&b), } + // With authentication on, the key decides the band and the routing profile, + // and the headers stop being believed. Trusting the header was only ever + // acceptable while there was no identity to consult — it let any client + // promote itself into the top priority band. + if principal, ok := auth.FromContext(req.Context()); ok { + p.Band = principal.Band + p.ProfileOverride = principal.RoutingProfile + p.KeyID = principal.KeyID + p.ProjectID = principal.ProjectID + // Tenant fairness within a band is per key rather than per session, so + // one key cannot take a whole band by opening many sessions. + p.Tenant = principal.KeyID + } + // Bytes over four: a crude stand-in for a tokenizer, good enough for // scoring. It becomes exact only with plan §7.5 tier 2 (see plan §14). p.PromptTokens = len(p.PromptPrefix) / 4 diff --git a/internal/router/router.go b/internal/router/router.go index 22479678..f8945e80 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -15,6 +15,7 @@ import ( "strings" "time" + "gitea.mynodes.xyz/ilker/llmgw/internal/auth" "gitea.mynodes.xyz/ilker/llmgw/internal/config" "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" "gitea.mynodes.xyz/ilker/llmgw/internal/router/admission" @@ -23,6 +24,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy" "gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler" "gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" ) // Router serves the OpenAI API in front of a set of engine replicas. @@ -41,8 +43,19 @@ type Router struct { // models maps every name a client may use — canonical names and aliases — // onto its catalog entry. models map[string]catalog.Model + + // auth and store are nil when authentication is disabled. + auth *auth.Authenticator + store *store.SQLite } +// Auth exposes the authenticator so ctl and the batch writer can reach it. +// Nil when authentication is off. +func (r *Router) Auth() *auth.Authenticator { return r.auth } + +// Store exposes the key store, or nil when authentication is off. +func (r *Router) Store() *store.SQLite { return r.store } + // New builds a router from configuration. Start must be called to begin scraping. func New(cfg config.Server, log *slog.Logger) (*Router, error) { if log == nil { @@ -101,6 +114,18 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { Timeout: cfg.Router.ScrapeTimeout, }, log) + if cfg.Auth.Enabled { + st, err := store.OpenSQLite(context.Background(), cfg.Auth.Store) + if err != nil { + return nil, fmt.Errorf("auth store: %w", err) + } + r.store = st + r.auth = auth.New(st, auth.Config{ + TTL: cfg.Auth.CacheTTL, + NegativeTTL: cfg.Auth.NegativeTTL, + }) + } + if cfg.Router.Admission.Enabled { adm, err := admission.New(admissionConfig(cfg.Router), r.saturation) if err != nil { diff --git a/llmgw-plan.md b/llmgw-plan.md index 2f6dd463..40ddda89 100644 --- a/llmgw-plan.md +++ b/llmgw-plan.md @@ -1405,8 +1405,13 @@ The launch path is now built too: because everything that genuinely needs persisting — keys, usage, capacity profiles — is a phase 5 feature. Writing it earlier would have meant designing a schema against guesses (§6.5). -- [ ] Tenancy: org/team/project/key hierarchy, budgets, RPM/TPM, usage records (§7.9) -- [ ] Budget fallback, model allowlists, key→band mapping +- [x] Tenancy: org/team/project/key hierarchy, key storage and authentication, + limit resolution (most restrictive along the chain), `llmgw keys` +- [ ] Budgets, RPM/TPM enforcement and the usage batch writer (§7.9.3, §7.9.4) +- [x] Model allowlists and key→band mapping. With auth on, the band comes from + the key and the header is ignored — it was trusted only while there was + no identity to consult, which let any client promote itself. +- [ ] Budget fallback - [ ] Routing profiles from configuration plus the `x-llmgw-profile` header - [ ] Approximate radix tree (§7.5 tier 1), with its thresholds and decision metrics - [ ] Fallback models and pool spillover diff --git a/test/e2e/auth_test.go b/test/e2e/auth_test.go new file mode 100644 index 00000000..5c99cd14 --- /dev/null +++ b/test/e2e/auth_test.go @@ -0,0 +1,310 @@ +package e2e + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// provision builds a store with a tenancy tree and returns its path plus a +// function that mints keys in it. +func provision(t *testing.T) (string, func(t *testing.T, name string, tune ...func(*store.Key)) string) { + t.Helper() + path := filepath.Join(t.TempDir(), "auth.db") + ctx := context.Background() + + s, err := store.OpenSQLite(ctx, path) + if err != nil { + t.Fatal(err) + } + for _, err := range []error{ + s.PutOrg(ctx, store.Org{ID: "org", Name: "Org"}), + s.PutTeam(ctx, store.Team{ID: "team", OrgID: "org", Name: "Team"}), + s.PutProject(ctx, store.Project{ID: "proj", TeamID: "team", Name: "Proj"}), + } { + if err != nil { + t.Fatal(err) + } + } + _ = s.Close() + + mint := func(t *testing.T, name string, tune ...func(*store.Key)) string { + t.Helper() + s, err := store.OpenSQLite(context.Background(), path) + if err != nil { + t.Fatal(err) + } + defer s.Close() + + secret, k, err := store.NewKey("key-"+name, name, "proj") + if err != nil { + t.Fatal(err) + } + k.CreatedAt = time.Now() + for _, f := range tune { + f(&k) + } + if err := s.CreateKey(context.Background(), k); err != nil { + t.Fatal(err) + } + return secret + } + return path, mint +} + +// withStore sets limits on an existing store. +func withStore(t *testing.T, path string, fn func(*store.SQLite)) { + t.Helper() + s, err := store.OpenSQLite(context.Background(), path) + if err != nil { + t.Fatal(err) + } + defer s.Close() + fn(s) +} + +func (h *harness) call(t *testing.T, key, model string) (int, string) { + t.Helper() + body, _ := json.Marshal(map[string]any{ + "model": model, + "messages": []map[string]string{{"role": "user", "content": "hi"}}, + "max_tokens": 4, + }) + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, + h.router.URL+"/v1/chat/completions", bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + if key != "" { + req.Header.Set("Authorization", "Bearer "+key) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + var e struct { + Error struct{ Message string } `json:"error"` + } + _ = json.NewDecoder(resp.Body).Decode(&e) + return resp.StatusCode, e.Error.Message +} + +func TestTheGatewayIsClosedWithoutAKey(t *testing.T) { + // The whole point of phase 5's first half: before this, anything that could + // reach the port could use every model in the catalog. + path, mint := provision(t) + key := mint(t, "good") + + h := start(t, 2, func(c *config.Server) { + c.Auth = config.Auth{Enabled: true, Store: path, CacheTTL: 50 * time.Millisecond} + }) + + if code, _ := h.call(t, "", testModel); code != http.StatusUnauthorized { + t.Errorf("no key got %d", code) + } + if code, _ := h.call(t, "sk-llmgw-"+"0123456789012345678901234567890123456789012", testModel); code != http.StatusUnauthorized { + t.Errorf("an unknown key got %d", code) + } + if code, msg := h.call(t, key, testModel); code != http.StatusOK { + t.Errorf("a valid key got %d: %s", code, msg) + } +} + +func TestRevocationTakesEffectWithinTheCacheTTL(t *testing.T) { + // The cache is what keeps the database off the request path, so the price + // is that a revocation is not instant. It has to be bounded and it has to + // actually happen. + path, mint := provision(t) + key := mint(t, "doomed") + + h := start(t, 2, func(c *config.Server) { + c.Auth = config.Auth{Enabled: true, Store: path, CacheTTL: 100 * time.Millisecond} + }) + if code, _ := h.call(t, key, testModel); code != http.StatusOK { + t.Fatal("the key did not work to begin with") + } + + withStore(t, path, func(s *store.SQLite) { + if err := s.RevokeKey(context.Background(), "key-doomed", time.Now()); err != nil { + t.Fatal(err) + } + }) + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if code, _ := h.call(t, key, testModel); code == http.StatusUnauthorized { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatal("a revoked key still works well past the cache TTL") +} + +func TestAnAllowlistCoversEveryNameOfAModel(t *testing.T) { + // The decision is per model, not per string. Matching only what the client + // sent would make every alias a way around the allowlist; matching only the + // canonical name would silently ignore an allowlist written in the names + // clients actually use. Both directions have to work. + path, mint := provision(t) + byCanonical := mint(t, "canonical") + byAlias := mint(t, "alias") + other := mint(t, "other") + + withStore(t, path, func(s *store.SQLite) { + for _, l := range []store.Limits{ + {Scope: store.ScopeKey, ScopeID: "key-canonical", AllowedModels: []string{testModel}}, + {Scope: store.ScopeKey, ScopeID: "key-alias", AllowedModels: []string{"alias-for-test-model"}}, + {Scope: store.ScopeKey, ScopeID: "key-other", AllowedModels: []string{"something-else"}}, + } { + if err := s.PutLimits(context.Background(), l); err != nil { + t.Fatal(err) + } + } + }) + + h := start(t, 2, func(c *config.Server) { + c.Auth = config.Auth{Enabled: true, Store: path, CacheTTL: 50 * time.Millisecond} + c.Models = []catalog.Model{{Name: testModel, Aliases: []string{"alias-for-test-model"}}} + }) + + cases := []struct { + name, key, model string + want int + }{ + {"canonical allowlist, canonical request", byCanonical, testModel, http.StatusOK}, + {"canonical allowlist, alias request", byCanonical, "alias-for-test-model", http.StatusOK}, + {"alias allowlist, alias request", byAlias, "alias-for-test-model", http.StatusOK}, + {"alias allowlist, canonical request", byAlias, testModel, http.StatusOK}, + {"unrelated allowlist, canonical request", other, testModel, http.StatusForbidden}, + {"unrelated allowlist, alias request", other, "alias-for-test-model", http.StatusForbidden}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if code, msg := h.call(t, c.key, c.model); code != c.want { + t.Errorf("got %d (%s), want %d", code, msg, c.want) + } + }) + } +} + +func TestTheKeyDecidesTheBandNotTheClient(t *testing.T) { + // Until authentication there was nothing to consult, so the band came from + // a header and any client could promote itself into the top priority band. + path, mint := provision(t) + bronze := mint(t, "bronze", func(k *store.Key) { k.Band = "bronze" }) + + h := start(t, 2, func(c *config.Server) { + c.Auth = config.Auth{Enabled: true, Store: path, CacheTTL: 50 * time.Millisecond} + c.Router.Admission.Enabled = true + c.Router.Admission.BandHeader = "x-llmgw-band" + c.Router.Admission.DefaultBand = "bronze" + c.Router.Admission.Bands = []config.Band{ + {Name: "platinum", Priority: 400, SaturationCeiling: 2}, + {Name: "bronze", Priority: 100, SaturationCeiling: 0.3}, + } + }) + + body, _ := json.Marshal(map[string]any{ + "model": testModel, "max_tokens": 4, + "messages": []map[string]string{{"role": "user", "content": "hi"}}, + }) + req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, + h.router.URL+"/v1/chat/completions", bytes.NewReader(body)) + req.Header.Set("Authorization", "Bearer "+bronze) + // The client claims the top band. The key says otherwise. + req.Header.Set("x-llmgw-band", "platinum") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status %d", resp.StatusCode) + } + + // The band a request was admitted in shows up in the admission metrics. + if got := admitted(t, h, "platinum"); got != 0 { + t.Errorf("a bronze key was admitted into platinum %d times", got) + } + if got := admitted(t, h, "bronze"); got == 0 { + t.Error("the request was not admitted in the key's own band") + } +} + +// admitted counts requests the router admitted in a band, read from its own +// metrics rather than inferred. +func admitted(t *testing.T, h *harness, band string) int { + t.Helper() + mux, err := h.rt.AdminHandler() + if err != nil { + t.Fatal(err) + } + srv := httptest.NewServer(mux) + defer srv.Close() + + req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, srv.URL+"/metrics", nil) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + + want := fmt.Sprintf("llmgw_admission_wait_seconds_count{band=\"%s\"", band) + for _, line := range strings.Split(string(body), "\n") { + if strings.HasPrefix(line, want) { + fields := strings.Fields(line) + n, err := strconv.Atoi(fields[len(fields)-1]) + if err != nil { + t.Fatalf("parsing %q: %v", line, err) + } + return n + } + } + return 0 +} + +func TestHealthStaysOpen(t *testing.T) { + // A load balancer has no key, and a health check that needs one is a health + // check that reports the gateway down. + path, _ := provision(t) + h := start(t, 1, func(c *config.Server) { + c.Auth = config.Auth{Enabled: true, Store: path} + }) + + req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, h.router.URL+"/health", nil) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Errorf("health needs a key: %d", resp.StatusCode) + } +} + +func TestWithoutAuthTheGatewayStillServes(t *testing.T) { + // Turning authentication on is a deployment choice. A gateway on a trusted + // network with no keys issued has to keep working. + h := start(t, 2) + if code, msg := h.call(t, "", testModel); code != http.StatusOK { + t.Errorf("an unauthenticated gateway refused a request: %d %s", code, msg) + } +}