easyai-ai-gateway/apps/api/internal/auth/oidc_test.go

package auth

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"math/big"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

func TestOIDCVerifierAcceptsRS256AndES256StableClaims(t *testing.T) {
	rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	ecKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	var issuer string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		switch request.URL.Path {
		case "/.well-known/openid-configuration":
			_ = json.NewEncoder(w).Encode(map[string]any{"issuer": issuer, "jwks_uri": issuer + "/jwks"})
		case "/jwks":
			_ = json.NewEncoder(w).Encode(map[string]any{"keys": []any{
				rsaJWK("rsa-key", &rsaKey.PublicKey), ecJWK("ec-key", &ecKey.PublicKey),
			}})
		default:
			http.NotFound(w, request)
		}
	}))
	defer server.Close()
	issuer = server.URL
	verifier, err := NewOIDCVerifier(OIDCConfig{
		Issuer: issuer, Audience: "gateway-api", TenantID: "tenant-1",
		RolePrefix: "gateway.", RequiredScopes: []string{"gateway.access"}, HTTPClient: server.Client(),
	})
	if err != nil {
		t.Fatal(err)
	}
	for _, test := range []struct {
		name   string
		kid    string
		method jwt.SigningMethod
		key    any
	}{{"rs256", "rsa-key", jwt.SigningMethodRS256, rsaKey}, {"es256", "ec-key", jwt.SigningMethodES256, ecKey}} {
		t.Run(test.name, func(t *testing.T) {
			token := signedOIDCToken(t, issuer, test.kid, test.method, test.key, nil)
			user, err := verifier.Verify(context.Background(), token)
			if err != nil {
				t.Fatal(err)
			}
			if user.ID != "platform-subject" || user.TenantID != "tenant-1" || user.Source != "oidc" ||
				len(user.Roles) != 1 || user.Roles[0] != "admin" {
				t.Fatalf("unexpected OIDC user: %#v", user)
			}
		})
	}
}

func TestOIDCVerifierRejectsMissingOrMismatchedSecurityClaims(t *testing.T) {
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	var issuer string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/.well-known/openid-configuration" {
			_ = json.NewEncoder(w).Encode(map[string]any{"issuer": issuer, "jwks_uri": issuer + "/jwks"})
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"keys": []any{ecJWK("ec-key", &key.PublicKey)}})
	}))
	defer server.Close()
	issuer = server.URL
	verifier, _ := NewOIDCVerifier(OIDCConfig{
		Issuer: issuer, Audience: "gateway-api", TenantID: "tenant-1", RolePrefix: "gateway.",
		RequiredScopes: []string{"gateway.access"}, HTTPClient: server.Client(),
	})
	tests := []struct {
		name   string
		mutate func(jwt.MapClaims)
	}{
		{"missing nbf", func(claims jwt.MapClaims) { delete(claims, "nbf") }},
		{"wrong audience", func(claims jwt.MapClaims) { claims["aud"] = "other-api" }},
		{"wrong tenant", func(claims jwt.MapClaims) { claims["tid"] = "tenant-2" }},
		{"missing scope", func(claims jwt.MapClaims) { claims["scope"] = "openid" }},
		{"unmapped role", func(claims jwt.MapClaims) { claims["roles"] = []string{"other.admin"} }},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			raw := signedOIDCToken(t, issuer, "ec-key", jwt.SigningMethodES256, key, test.mutate)
			if _, err := verifier.Verify(context.Background(), raw); err == nil {
				t.Fatal("expected token to be rejected")
			}
		})
	}
}

func TestOIDCVerifierValidatesIDTokenNonceAndPublicClientAudience(t *testing.T) {
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	var issuer string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/.well-known/openid-configuration" {
			_ = json.NewEncoder(w).Encode(map[string]any{"issuer": issuer, "jwks_uri": issuer + "/jwks"})
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"keys": []any{ecJWK("ec-key", &key.PublicKey)}})
	}))
	defer server.Close()
	issuer = server.URL
	verifier, _ := NewOIDCVerifier(OIDCConfig{
		Issuer: issuer, Audience: "gateway-api", TenantID: "tenant-1", RolePrefix: "gateway.",
		RequiredScopes: []string{"gateway.access"}, HTTPClient: server.Client(),
	})
	idToken := signedOIDCToken(t, issuer, "ec-key", jwt.SigningMethodES256, key, func(claims jwt.MapClaims) {
		claims["aud"] = "gateway-public-client"
		claims["nonce"] = "expected-nonce"
	})
	subject, err := verifier.VerifyIDToken(context.Background(), idToken, "gateway-public-client", "expected-nonce")
	if err != nil || subject != "platform-subject" {
		t.Fatalf("VerifyIDToken() subject=%q err=%v", subject, err)
	}
	if _, err := verifier.VerifyIDToken(context.Background(), idToken, "gateway-public-client", "wrong-nonce"); err == nil {
		t.Fatal("ID token with mismatched nonce was accepted")
	}
}

func TestOIDCVerifierFailsClosedWhenIntrospectionMarksSessionInactive(t *testing.T) {
	key, _ := rsa.GenerateKey(rand.Reader, 2048)
	active := true
	var issuer string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		switch request.URL.Path {
		case "/.well-known/openid-configuration":
			_ = json.NewEncoder(w).Encode(map[string]any{
				"issuer": issuer, "jwks_uri": issuer + "/jwks",
				"introspection_endpoint": issuer + "/introspect",
			})
		case "/jwks":
			_ = json.NewEncoder(w).Encode(map[string]any{"keys": []any{rsaJWK("rsa-key", &key.PublicKey)}})
		case "/introspect":
			clientID, clientSecret, ok := request.BasicAuth()
			if !ok || clientID != "gateway-api" || clientSecret != "introspection-secret" {
				w.WriteHeader(http.StatusUnauthorized)
				return
			}
			if err := request.ParseForm(); err != nil || request.Form.Get("token") == "" {
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			_ = json.NewEncoder(w).Encode(map[string]any{"active": active})
		default:
			http.NotFound(w, request)
		}
	}))
	defer server.Close()
	issuer = server.URL
	verifier, err := NewOIDCVerifier(OIDCConfig{
		Issuer: issuer, Audience: "gateway-api", TenantID: "tenant-1", RolePrefix: "gateway.",
		RequiredScopes: []string{"gateway.access"}, HTTPClient: server.Client(),
		IntrospectionEnabled: true, IntrospectionClientID: "gateway-api",
		IntrospectionClientSecret: "introspection-secret",
	})
	if err != nil {
		t.Fatal(err)
	}
	raw := signedOIDCToken(t, issuer, "rsa-key", jwt.SigningMethodRS256, key, nil)
	if _, err := verifier.Verify(context.Background(), raw); err != nil {
		t.Fatalf("active token rejected: %v", err)
	}
	active = false
	if _, err := verifier.Verify(context.Background(), raw); err == nil {
		t.Fatal("inactive token was accepted")
	}
}

func signedOIDCToken(t *testing.T, issuer, kid string, method jwt.SigningMethod, key any, mutate func(jwt.MapClaims)) string {
	t.Helper()
	now := time.Now()
	claims := jwt.MapClaims{
		"iss": issuer, "aud": "gateway-api", "sub": "platform-subject", "tid": "tenant-1",
		"preferred_username": "acceptance", "roles": []string{"gateway.admin"},
		"scope": "openid gateway.access", "iat": now.Unix(), "nbf": now.Add(-time.Second).Unix(),
		"exp": now.Add(time.Hour).Unix(),
	}
	if mutate != nil {
		mutate(claims)
	}
	token := jwt.NewWithClaims(method, claims)
	token.Header["kid"] = kid
	raw, err := token.SignedString(key)
	if err != nil {
		t.Fatal(err)
	}
	return raw
}

func rsaJWK(kid string, key *rsa.PublicKey) map[string]any {
	return map[string]any{
		"kid": kid, "kty": "RSA", "use": "sig", "alg": "RS256",
		"n": base64.RawURLEncoding.EncodeToString(key.N.Bytes()),
		"e": base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key.E)).Bytes()),
	}
}

func ecJWK(kid string, key *ecdsa.PublicKey) map[string]any {
	return map[string]any{
		"kid": kid, "kty": "EC", "use": "sig", "alg": "ES256", "crv": "P-256",
		"x": base64.RawURLEncoding.EncodeToString(key.X.FillBytes(make([]byte, 32))),
		"y": base64.RawURLEncoding.EncodeToString(key.Y.FillBytes(make([]byte, 32))),
	}
}