mcp-chassis/auth/jwt.go

// Package auth provides the Dex/Authentik-JWT + static-Bearer authentication
// primitives shared by every Mathias-owned MCP server (gitea-mcp, brain-mcp /
// ingestion, future MCPs spawned from template-go-agent).
//
// Replaces ~80 LOC of near-identical jwt.go in each consumer, ~50 LOC of
// Bearer middleware, and ~25 LOC of RFC 9728 protected-resource metadata
// handler. See `gitea.d-ma.be/mathias/infra` docs/superpowers/handoffs/
// 2026-05-22-mcp-chassis-spike.md for the design rationale and the
// abort-criterion check.
//
// A validator trusts a LIST of OIDC issuers (infra ADR-0011): one MCP server can
// accept Authentik JWTs (interactive/web) AND k3s cluster ServiceAccount tokens
// (in-cluster, audience-bound, kubelet-rotated) at once.
package auth

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jwt"
)

// ErrUnavailable indicates JWT validation could not be COMPLETED because the
// JWKS / issuer endpoint was unreachable — a transient condition — as opposed to
// the token being present and invalid. Callers (e.g. BearerMiddleware) map it
// to HTTP 503 temporarily_unavailable rather than a generic 401, so an issuer
// outage is distinguishable from a bad token (gitea-mcp#6).
var ErrUnavailable = errors.New("jwt validation temporarily unavailable")

// IssuerConfig names one trusted OIDC issuer and its (optional) required
// audience. Trusting a list of these is what lets a single MCP server accept
// tokens from Authentik (D3/D4) and the k3s cluster OIDC issuer (D1 SA tokens)
// simultaneously — infra ADR-0011 Decision 5.
type IssuerConfig struct {
	IssuerURL string
	Audience  string // "" = skip audience validation for this issuer
	// HTTPClient, if non-nil, is used to fetch THIS issuer's OIDC discovery
	// document and JWKS. Supply one to reach an issuer that needs a custom CA
	// and/or credential — notably the in-cluster k8s OIDC endpoint, which serves
	// its discovery/JWKS over the cluster CA AND requires an authenticated
	// (ServiceAccount-bearer) request (anonymous access is 401). nil uses
	// http.DefaultClient — correct for public issuers like Authentik.
	HTTPClient *http.Client
}

// issuerEntry is one trusted issuer. jwksURI is resolved LAZILY: OIDC discovery
// may not be answerable when the process starts, and an issuer that is merely
// not-yet-serving must not become a permanent auth outage (mcp-chassis#4,
// gitea-mcp#57 — three recurrences, each ended by a human running
// `kubectl rollout restart`).
//
// Everything downstream of resolution was already correct. Validate reports
// ErrUnavailable when it cannot reach an issuer, and BearerMiddleware already
// answers 503-and-retry for that. Resolving eagerly in the constructor was the
// one step that made the correct path unreachable.
type issuerEntry struct {
	issuer     string
	audience   string
	httpClient *http.Client

	mu          sync.Mutex
	jwksURI     string    // "" until discovery has succeeded
	lastAttempt time.Time // zero until the first attempt
	lastErr     error
}

// resolve returns this issuer's jwks_uri, running OIDC discovery if it has not
// succeeded yet. Failures are retried, no sooner than retryInterval apart, so a
// down issuer costs one discovery request per interval rather than one per
// inbound request.
func (e *issuerEntry) resolve(ctx context.Context, cache *jwk.Cache, retryInterval, refreshInterval time.Duration) (string, error) {
	e.mu.Lock()
	defer e.mu.Unlock()

	if e.jwksURI != "" {
		return e.jwksURI, nil
	}
	if !e.lastAttempt.IsZero() && time.Since(e.lastAttempt) < retryInterval {
		// Back off, but report the real reason rather than a generic one — the
		// operator needs the issuer's actual error, not "throttled".
		return "", e.lastErr
	}
	e.lastAttempt = time.Now()

	client := e.httpClient
	if client == nil {
		client = http.DefaultClient
	}
	jwksURI, err := discoverJWKSURI(ctx, client, e.issuer)
	if err != nil {
		e.lastErr = fmt.Errorf("%w: issuer %s: %v", ErrUnavailable, e.issuer, err)
		return "", e.lastErr
	}

	regOpts := []jwk.RegisterOption{jwk.WithMinRefreshInterval(refreshInterval)}
	if e.httpClient != nil {
		regOpts = append(regOpts, jwk.WithHTTPClient(e.httpClient))
	}
	if err := cache.Register(jwksURI, regOpts...); err != nil {
		e.lastErr = fmt.Errorf("%w: register jwks cache (%s): %v", ErrUnavailable, e.issuer, err)
		return "", e.lastErr
	}
	if _, err := cache.Refresh(ctx, jwksURI); err != nil {
		// Registered but unfetchable. Unregister so the next attempt re-runs
		// discovery cleanly rather than inheriting a half-built entry.
		_ = cache.Unregister(jwksURI)
		e.lastErr = fmt.Errorf("%w: initial jwks fetch (%s): %v", ErrUnavailable, e.issuer, err)
		return "", e.lastErr
	}

	e.jwksURI, e.lastErr = jwksURI, nil
	return jwksURI, nil
}

// Option configures a validator. Variadic so every existing caller compiles
// unchanged — a shared library cannot afford a breaking signature for this.
type Option func(*options)

type options struct {
	retryInterval   time.Duration
	refreshInterval time.Duration
}

// defaultRetryInterval bounds how often a down issuer is re-probed. A readiness
// probe typically runs every 10s, so recovery lands within one interval of the
// issuer coming back — without turning an IdP outage into a discovery flood.
const defaultRetryInterval = 30 * time.Second

// defaultRefreshInterval is the JWKS cache's minimum refresh interval. 1h
// preserves the value that was hardcoded here before options existed; making it
// configurable is mcp-chassis#1.
const defaultRefreshInterval = time.Hour

// WithRetryInterval sets how long to wait before re-attempting OIDC discovery
// for an issuer that is not yet resolved. Zero disables the wait, which is for
// tests; production wants the default.
func WithRetryInterval(d time.Duration) Option {
	return func(o *options) { o.retryInterval = d }
}

// JWTValidator validates Bearer JWTs against one or more trusted OIDC issuers.
//
// A nil *JWTValidator behaves as "JWT auth disabled" — Validate returns an
// error without panicking. Callers can construct one validator at startup keyed
// on whether any issuer is configured, and pass nil through the rest of the
// codebase without further branching.
type JWTValidator struct {
	entries []*issuerEntry
	cache   *jwk.Cache
	opts    options
}

// NewJWTValidator builds a single-issuer validator — the common case, and
// backward compatible with every existing caller. Empty issuerURL returns
// (nil, nil) so callers can use one constructor regardless of whether an issuer
// is configured.
func NewJWTValidator(ctx context.Context, issuerURL, audience string, opts ...Option) (*JWTValidator, error) {
	if issuerURL == "" {
		return nil, nil
	}
	return NewMultiJWTValidator(ctx, []IssuerConfig{{IssuerURL: issuerURL, Audience: audience}}, opts...)
}

// NewMultiJWTValidator builds a validator that trusts every issuer in the list.
// For each, it fetches the OIDC discovery document, registers jwks_uri in a
// shared cache, and warms it. Entries with an empty IssuerURL are skipped; an
// empty/nil resulting set returns (nil, nil) — "JWT auth disabled".
func NewMultiJWTValidator(ctx context.Context, issuers []IssuerConfig, opts ...Option) (*JWTValidator, error) {
	o := options{retryInterval: defaultRetryInterval, refreshInterval: defaultRefreshInterval}
	for _, fn := range opts {
		fn(&o)
	}

	var entries []*issuerEntry
	for _, ic := range issuers {
		if ic.IssuerURL == "" {
			continue
		}
		entries = append(entries, &issuerEntry{
			issuer:     ic.IssuerURL,
			audience:   ic.Audience,
			httpClient: ic.HTTPClient,
		})
	}
	if len(entries) == 0 {
		// No issuer configured: JWT auth is genuinely disabled, and nil is the
		// right answer. This is the ONE case where nil is not a lie, which is
		// why the boot-race fix must not collapse it — see
		// TestNewMultiJWTValidator_NoIssuer_StillReturnsNilNil.
		return nil, nil
	}

	v := &JWTValidator{entries: entries, cache: jwk.NewCache(ctx), opts: o}

	// Try once, eagerly, so a healthy start still warms the cache before the
	// first request — but do NOT fail on it. A validator that exists and is
	// temporarily unresolved is strictly better than no validator: the first
	// answers 503-and-retry, the second answers 401 forever.
	for _, e := range v.entries {
		_, _ = e.resolve(ctx, v.cache, 0, o.refreshInterval)
	}
	return v, nil
}

// Ready reports whether every trusted issuer has been resolved, attempting
// resolution for any that have not. It is safe and intended to be called from a
// readiness probe: the probe's own cadence then drives retry, so recovery needs
// no background goroutine and no process restart.
//
// A nil validator is ready — "JWT auth disabled" is a configuration, not a
// fault.
func (v *JWTValidator) Ready(ctx context.Context) (bool, error) {
	if v == nil {
		return true, nil
	}
	for _, e := range v.entries {
		if _, err := e.resolve(ctx, v.cache, v.opts.retryInterval, v.opts.refreshInterval); err != nil {
			return false, err
		}
	}
	return true, nil
}

// discoverJWKSURI fetches the OIDC discovery document from issuerURL and returns
// its jwks_uri.
func discoverJWKSURI(ctx context.Context, client *http.Client, issuerURL string) (string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet,
		issuerURL+"/.well-known/openid-configuration", nil)
	if err != nil {
		return "", fmt.Errorf("build oidc discovery request: %w", err)
	}
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("fetch oidc discovery: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("oidc discovery: status %d", resp.StatusCode)
	}

	var doc struct {
		JWKSURI string `json:"jwks_uri"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&doc); err != nil {
		return "", fmt.Errorf("decode oidc discovery: %w", err)
	}
	if doc.JWKSURI == "" {
		return "", fmt.Errorf("oidc discovery: empty jwks_uri")
	}
	return doc.JWKSURI, nil
}

// Validate parses rawToken against each trusted issuer and returns the subject
// claim on the first success. If no issuer accepts it: when EVERY issuer's JWKS
// was unreachable the error wraps ErrUnavailable (caller answers 503); otherwise
// at least one issuer reached a signature/claims verdict and rejected it, so the
// error is a definitive rejection (401). A nil receiver returns errDisabled.
func (v *JWTValidator) Validate(ctx context.Context, rawToken string) (string, error) {
	if v == nil {
		return "", errDisabled
	}

	var lastErr error
	sawDecision := false // at least one issuer reached a real verify verdict
	for _, e := range v.entries {
		jwksURI, err := e.resolve(ctx, v.cache, v.opts.retryInterval, v.opts.refreshInterval)
		if err != nil {
			// Not resolved yet — transient by construction (resolve wraps
			// ErrUnavailable). Try the next issuer.
			lastErr = err
			continue
		}
		keySet, err := v.cache.Get(ctx, jwksURI)
		if err != nil {
			// This issuer's JWKS is unreachable — transient. Try the next; only
			// if ALL issuers are unreachable do we surface ErrUnavailable.
			lastErr = fmt.Errorf("%w: get jwks (%s): %v", ErrUnavailable, e.issuer, err)
			continue
		}

		opts := []jwt.ParseOption{
			jwt.WithKeySet(keySet),
			jwt.WithValidate(true),
			jwt.WithIssuer(e.issuer),
		}
		if e.audience != "" {
			opts = append(opts, jwt.WithAudience(e.audience))
		}

		tok, err := jwt.ParseString(rawToken, opts...)
		if err == nil {
			return tok.Subject(), nil
		}
		sawDecision = true
		lastErr = fmt.Errorf("validate jwt: %w", err)
	}

	if lastErr == nil {
		lastErr = errDisabled
	}
	if !sawDecision {
		// Every issuer's JWKS was unreachable — lastErr already wraps ErrUnavailable.
		return "", lastErr
	}
	return "", lastErr
}

// errDisabled is the sentinel returned by Validate on a nil receiver.
// Not exported because callers care about "not authorized" not "why";
// this lets consumers treat any err the same way without depending on
// the specific value.
var errDisabled = fmt.Errorf("jwt auth disabled")