llmgw/internal/auth/auth.go

// 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"
	"golang.org/x/sync/singleflight"
)

// Principal is who a request is, once authenticated.
type Principal struct {
	KeyID     string
	KeyPrefix string
	// UserID is the person the key belongs to, empty for a key an operator
	// issued in a project. A personal key has a person and no project.
	UserID    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")
	// ErrThrottled is a client presenting more unknown keys than it is
	// allowed to have looked up.
	ErrThrottled = errors.New("too many unrecognised api keys from this address")
)

// 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

	// lookups rations the database reads of keys not in the cache, per
	// client address.
	lookups *buckets

	// resolving shares one database read between the requests that miss the
	// cache for the same key at once: when a busy key's entry expired, every
	// request in flight missed together, each took from the client's ration,
	// and a valid key was told to slow down.
	resolving singleflight.Group
}

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{},
		lookups: newBuckets(20, 50),
	}
}

// Authenticate resolves a presented credential.
func (a *Authenticator) Authenticate(ctx context.Context, presented string) (Principal, error) {
	return a.AuthenticateFrom(ctx, presented, "")
}

// AuthenticateFrom resolves a credential presented by client, an address:
// a key not in the cache costs a database read, so those are rationed per
// client, and one that has made too many is told to slow down rather than
// read for. Empty client is not rationed.
func (a *Authenticator) AuthenticateFrom(ctx context.Context, presented, client 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
	}

	// One read for every request missing this key at once, rationed once.
	v, err, _ := a.resolving.Do(id, func() (any, error) {
		if e, ok := a.lookup(id); ok { // filled while this waited its turn
			return e, nil
		}
		if client != "" && !a.lookups.allow(client, a.now()) {
			return nil, ErrThrottled
		}
		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 nil, err
		}
		ttl := a.cfg.TTL
		if err != nil {
			ttl = a.cfg.NegativeTTL
		}
		e := entry{principal: p, err: err, expires: a.now().Add(ttl)}
		a.store_(id, e)
		return e, nil
	})
	if err != nil {
		return Principal{}, err
	}
	e := v.(entry)
	if e.err == nil {
		a.markUsed(e.principal.KeyID)
	}
	return e.principal, e.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 errors.Is(err, store.ErrOwnerInactive) {
		// Said as a revoked key is: whether the key exists is not the
		// caller's to learn.
		return Principal{}, ErrRevoked
	}
	if err != nil {
		return Principal{}, err
	}
	return Principal{
		KeyID: k.ID, KeyPrefix: k.Prefix, UserID: eff.UserID, 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
}

// maxCached bounds the cache. Every well-formed key anyone presents is an
// entry, a rejected one included, so without a bound a stream of made-up keys
// grows it until the process runs out of memory.
const maxCached = 100_000

func (a *Authenticator) store_(id string, e entry) {
	a.mu.Lock()
	defer a.mu.Unlock()
	if len(a.cache) >= maxCached {
		now := a.now()
		for k, v := range a.cache {
			if now.After(v.expires) {
				delete(a.cache, k)
			}
		}
		// Still full of live entries: half of them go, which costs only a
		// lookup each for the keys in use.
		if len(a.cache) >= maxCached {
			n := 0
			for k := range a.cache {
				if n%2 == 0 {
					delete(a.cache, k)
				}
				n++
			}
		}
	}
	a.cache[id] = e
}

// 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)
	}
	if k := r.Header.Get("api-key"); k != "" {
		return strings.TrimSpace(k)
	}
	// Anthropic's clients send their key here.
	return strings.TrimSpace(r.Header.Get("x-api-key"))
}