mcp-chassis/auth/jwt_test.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jwt"
	"github.com/stretchr/testify/require"
)

// testIssuer is an in-process OIDC issuer: it serves an openid-configuration
// discovery doc + a JWKS, and mints signed JWTs — enough to exercise the real
// signature/issuer/audience validation path (which the chassis previously had
// no happy-path test for). Used to prove multi-issuer validation (ADR-0011).
type testIssuer struct {
	url           string
	priv          jwk.Key
	requireBearer string // if set, discovery+JWKS return 401 without this bearer

	// mu guards the boot-race fields below, which a test mutates while the
	// httptest server is serving.
	mu sync.Mutex
	// discoveryFailures is how many initial discovery requests answer 503 —
	// the IdP-not-yet-serving case that disabled OAuth for six days
	// (mcp-chassis#4, gitea-mcp#57).
	discoveryFailures int
	discoveryHits     int
}

func (ti *testIssuer) authed(r *http.Request) bool {
	return ti.requireBearer == "" || r.Header.Get("Authorization") == "Bearer "+ti.requireBearer
}

// discoveryDown reports whether this discovery request should answer 503, and
// counts the attempt. Attempts are counted even while down, so a test can assert
// that resolution was RETRIED rather than cached as a permanent failure.
func (ti *testIssuer) discoveryDown() bool {
	ti.mu.Lock()
	defer ti.mu.Unlock()
	ti.discoveryHits++
	return ti.discoveryHits <= ti.discoveryFailures
}

func (ti *testIssuer) discoveryAttempts() int {
	ti.mu.Lock()
	defer ti.mu.Unlock()
	return ti.discoveryHits
}

// bearerRT adds a static Authorization header to every request — a minimal stand-in
// for the in-cluster case where the JWKS fetch carries the pod's SA token.
type bearerRT struct{ token string }

func (b bearerRT) RoundTrip(r *http.Request) (*http.Response, error) {
	r.Header.Set("Authorization", "Bearer "+b.token)
	return http.DefaultTransport.RoundTrip(r)
}

func newTestIssuer(t *testing.T) *testIssuer {
	t.Helper()
	raw, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	priv, err := jwk.FromRaw(raw)
	require.NoError(t, err)
	require.NoError(t, priv.Set(jwk.KeyIDKey, "test-kid"))
	require.NoError(t, priv.Set(jwk.AlgorithmKey, jwa.RS256))
	pub, err := priv.PublicKey()
	require.NoError(t, err)
	set := jwk.NewSet()
	require.NoError(t, set.AddKey(pub))

	ti := &testIssuer{priv: priv}
	mux := http.NewServeMux()
	mux.HandleFunc("/jwks", func(w http.ResponseWriter, r *http.Request) {
		if !ti.authed(r) {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		_ = json.NewEncoder(w).Encode(set)
	})
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
		if ti.discoveryDown() {
			w.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		if !ti.authed(r) {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]string{
			"issuer":   ti.url,
			"jwks_uri": ti.url + "/jwks",
		})
	})
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	ti.url = srv.URL
	return ti
}

func (ti *testIssuer) mint(t *testing.T, aud, sub string) string {
	t.Helper()
	tok, err := jwt.NewBuilder().
		Issuer(ti.url).
		Subject(sub).
		Audience([]string{aud}).
		IssuedAt(time.Now()).
		Expiration(time.Now().Add(time.Hour)).
		Build()
	require.NoError(t, err)
	signed, err := jwt.Sign(tok, jwt.WithKey(jwa.RS256, ti.priv))
	require.NoError(t, err)
	return string(signed)
}

// The core ADR-0011 claim: one validator, several trusted issuers (e.g.
// Authentik + the k3s cluster OIDC), a token from any trusted issuer validates.
func TestNewMultiJWTValidator_AcceptsEitherTrustedIssuer(t *testing.T) {
	ctx := context.Background()
	a := newTestIssuer(t)
	b := newTestIssuer(t)

	v, err := NewMultiJWTValidator(ctx, []IssuerConfig{
		{IssuerURL: a.url, Audience: "brain-mcp"},
		{IssuerURL: b.url, Audience: "brain-mcp"},
	})
	require.NoError(t, err)
	require.NotNil(t, v)

	subA, err := v.Validate(ctx, a.mint(t, "brain-mcp", "alice"))
	require.NoError(t, err)
	require.Equal(t, "alice", subA)

	subB, err := v.Validate(ctx, b.mint(t, "brain-mcp", "bob"))
	require.NoError(t, err)
	require.Equal(t, "bob", subB)
}

// A token from an issuer NOT in the trusted list is rejected — definitively
// (401), not as a transient JWKS outage (503).
func TestNewMultiJWTValidator_RejectsUntrustedIssuer(t *testing.T) {
	ctx := context.Background()
	trusted := newTestIssuer(t)
	untrusted := newTestIssuer(t)

	v, err := NewMultiJWTValidator(ctx, []IssuerConfig{{IssuerURL: trusted.url, Audience: "brain-mcp"}})
	require.NoError(t, err)

	_, err = v.Validate(ctx, untrusted.mint(t, "brain-mcp", "eve"))
	require.Error(t, err)
	require.NotErrorIs(t, err, ErrUnavailable)
}

// Audience is enforced per issuer — the k8s-SA-token replay guard from the ADR
// spike. A token minted for a different audience is rejected.
func TestNewMultiJWTValidator_EnforcesPerIssuerAudience(t *testing.T) {
	ctx := context.Background()
	a := newTestIssuer(t)

	v, err := NewMultiJWTValidator(ctx, []IssuerConfig{{IssuerURL: a.url, Audience: "brain-mcp"}})
	require.NoError(t, err)

	_, err = v.Validate(ctx, a.mint(t, "some-other-service", "alice"))
	require.Error(t, err)
}

// Backward compatibility: the existing single-issuer constructor still works and
// validates a real signed token end-to-end.
func TestNewJWTValidator_SingleIssuer_BackwardCompatible(t *testing.T) {
	ctx := context.Background()
	a := newTestIssuer(t)

	v, err := NewJWTValidator(ctx, a.url, "brain-mcp")
	require.NoError(t, err)
	require.NotNil(t, v)

	sub, err := v.Validate(ctx, a.mint(t, "brain-mcp", "carol"))
	require.NoError(t, err)
	require.Equal(t, "carol", sub)
}

func TestNewMultiJWTValidator_EmptyList_ReturnsNilNil(t *testing.T) {
	v, err := NewMultiJWTValidator(context.Background(), nil)
	require.NoError(t, err)
	require.Nil(t, v)
}

// The in-cluster k8s issuer requires an AUTHENTICATED discovery/JWKS fetch
// (anonymous is 401, confirmed against k3s). IssuerConfig.HTTPClient must carry
// that credential (the server pod's SA token) + trust the cluster CA.
func TestNewMultiJWTValidator_UsesPerIssuerHTTPClientForAuthedFetch(t *testing.T) {
	ctx := context.Background()
	iss := newTestIssuer(t)
	iss.requireBearer = "fetch-secret" // discovery + JWKS now demand this bearer

	// nil client (http.DefaultClient, no credential) → discovery 401 → the issuer
	// never resolves.
	//
	// This used to assert that the CONSTRUCTOR errored. It no longer does: since
	// mcp-chassis#4, construction cannot fail on an unreachable issuer, because
	// "could not resolve" and "JWT auth not wanted" must stop collapsing into the
	// same nil. The property that actually matters is unchanged and asserted more
	// directly here — without the credential, a correctly-minted token from this
	// very issuer still gets nowhere, and the validator says it is not ready.
	unauth, err := NewMultiJWTValidator(ctx,
		[]IssuerConfig{{IssuerURL: iss.url, Audience: "sa-poc"}},
		WithRetryInterval(0),
	)
	require.NoError(t, err)
	require.NotNil(t, unauth)

	ready, why := unauth.Ready(ctx)
	require.False(t, ready, "an issuer whose discovery 401s is not ready")
	require.Error(t, why)

	_, err = unauth.Validate(ctx, iss.mint(t, "sa-poc", "poc-pod"))
	require.Error(t, err, "unauthenticated fetch must authenticate nobody, even with a genuine token")
	require.ErrorIs(t, err, ErrUnavailable)

	// A client that presents the fetch credential → build succeeds, token validates.
	client := &http.Client{Transport: bearerRT{token: "fetch-secret"}}
	v, err := NewMultiJWTValidator(ctx, []IssuerConfig{
		{IssuerURL: iss.url, Audience: "sa-poc", HTTPClient: client},
	})
	require.NoError(t, err)
	sub, err := v.Validate(ctx, iss.mint(t, "sa-poc", "poc-pod"))
	require.NoError(t, err)
	require.Equal(t, "poc-pod", sub)
}