mcp-chassis/auth/jwt_test.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"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
}

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, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(set)
	})
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, _ *http.Request) {
		_ = 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)
}