Close the gateway: API keys, tenancy and model allowlists

Until now anything that could reach the client port could use every model in
the catalog. A key now resolves to an identity, and that identity decides what
it may do.

Authentication wraps the whole mux rather than checking inside each handler,
so adding a route cannot add an unauthenticated one by omission. Health stays
open, because a load balancer has no key and a health check that needs one
reports the gateway down.

Nothing on the request path reads the database. A key is resolved once and
answered from memory until the entry expires; rejections are cached too, so a
client looping with a bad key cannot turn into a database read per request.
The price is that a revocation takes up to cache_ttl to take effect, which is
the trade stated in the configuration next to the knob. A store failure is
not an authentication answer: it is never cached as one and returns 503, not
401 — sending someone hunting for a key problem that does not exist wastes the
one thing an incident is short of. A revoked key is refused with the same
message as an unknown one, because saying otherwise confirms the key was real.

With auth on the band comes from the key and the band header is ignored. The
header was trusted only while there was nothing to consult, and it let any
client promote itself into the top priority band. Tenant fairness within a
band is now per key rather than per session, so one caller cannot take a band
by opening many sessions.

Allowlists are decided per model, not per string: a model's canonical name and
every alias name the same thing, so either form works in an allowlist and no
alias can reach a model that was not allowed. The first version of that test
passed without the resolution being there at all, which is why it now checks
all four combinations.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5 committed 2026-09-16 08:44:33 +02:00
1 parent 29617f8cb1
commit 45c8cd23a3
11 files changed
+1491 -14

No files matched your search

+25 -3
View File
@@ -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
+267
View File
@@ -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"))
}
+382
View File
@@ -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)
}
}
+97
View File
@@ -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,
}})
}
+288
View File
@@ -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 <name>",
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 <key-id>",
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[:]))
}
+1
View File
@@ -40,6 +40,7 @@ func Root() *cobra.Command {
newServerCmd(),
newAgentCmd(),
newCtlCmd(),
newTenancyCmd(),
newSimEngineCmd(),
)
return root
+23
View File
@@ -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 {
+66 -9
View File
@@ -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
+25
View File
@@ -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 {
+7 -2
View File
@@ -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
+310
View File
@@ -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)
}
}