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"
)

// 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"))
}