ai-gateway-go/internal/identity/oidc_test.go

package identity

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

func TestVerifyIDTokenValidatesOIDCSecurityClaims(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	kid := "test-key"
	jwks := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{"keys": []map[string]string{{
			"kid": kid,
			"kty": "RSA",
			"n":   base64.RawURLEncoding.EncodeToString(key.N.Bytes()),
			"e":   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key.E)).Bytes()),
		}}})
	}))
	defer jwks.Close()

	now := time.Unix(1_800_000_000, 0)
	service := &Service{
		allowPrivateIdentityProvider: true,
		oidcClient:                   newOIDCHTTPClient(true),
		now:                          func() time.Time { return now },
	}
	discovery := oidcDiscovery{Issuer: "https://issuer.example", JWKSURI: jwks.URL}
	baseClaims := map[string]any{
		"iss":   discovery.Issuer,
		"sub":   "subject-1",
		"aud":   "client-1",
		"exp":   now.Add(time.Minute).Unix(),
		"iat":   now.Unix(),
		"nonce": "nonce-1",
	}

	if _, err := service.verifyIDToken(context.Background(), discovery, "client-1", "nonce-1", signTestIDToken(t, key, kid, baseClaims)); err != nil {
		t.Fatalf("valid ID token rejected: %v", err)
	}

	tests := []struct {
		name   string
		change func(map[string]any)
	}{
		{name: "wrong issuer", change: func(c map[string]any) { c["iss"] = "https://attacker.example" }},
		{name: "wrong nonce", change: func(c map[string]any) { c["nonce"] = "other" }},
		{name: "expired", change: func(c map[string]any) { c["exp"] = now.Add(-time.Minute).Unix() }},
		{name: "issued in future", change: func(c map[string]any) { c["iat"] = now.Add(time.Minute).Unix() }},
		{name: "missing subject", change: func(c map[string]any) { delete(c, "sub") }},
		{name: "wrong audience", change: func(c map[string]any) { c["aud"] = "other-client" }},
		{name: "multiple audiences without azp", change: func(c map[string]any) { c["aud"] = []string{"client-1", "other-client"} }},
		{name: "wrong authorized party", change: func(c map[string]any) { c["aud"] = []string{"client-1", "other-client"}; c["azp"] = "other-client" }},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			claims := cloneClaims(baseClaims)
			test.change(claims)
			if _, err := service.verifyIDToken(context.Background(), discovery, "client-1", "nonce-1", signTestIDToken(t, key, kid, claims)); err == nil {
				t.Fatal("expected ID token to be rejected")
			}
		})
	}

	multiple := cloneClaims(baseClaims)
	multiple["aud"] = []string{"client-1", "other-client"}
	multiple["azp"] = "client-1"
	if _, err := service.verifyIDToken(context.Background(), discovery, "client-1", "nonce-1", signTestIDToken(t, key, kid, multiple)); err != nil {
		t.Fatalf("valid multi-audience ID token rejected: %v", err)
	}
}

func TestOIDCURLAndScopeValidation(t *testing.T) {
	got, err := validateOIDCURL(context.Background(), "http://127.0.0.1:9090/issuer/", true)
	if err != nil || got != "http://127.0.0.1:9090/issuer" {
		t.Fatalf("private development issuer rejected: %q %v", got, err)
	}
	for _, raw := range []string{
		"http://127.0.0.1:9090/issuer?tenant=1",
		"http://127.0.0.1:9090/issuer#fragment",
		"http://8.8.8.8/issuer",
	} {
		allowPrivate := raw != "http://8.8.8.8/issuer"
		if _, err := validateOIDCURL(context.Background(), raw, allowPrivate); err == nil {
			t.Fatalf("unsafe issuer accepted: %s", raw)
		}
	}
	if _, err := validateAbsoluteURL("https://user:secret@example.com/callback"); err == nil {
		t.Fatal("URL containing credentials was accepted")
	}

	scopes, err := normalizeOIDCScopes([]string{"openid", " profile ", "openid", "email"})
	if err != nil || len(scopes) != 3 || scopes[0] != "openid" || scopes[1] != "profile" || scopes[2] != "email" {
		t.Fatalf("unexpected normalized scopes: %#v %v", scopes, err)
	}
	for _, scopes := range [][]string{{"openid email"}, {"openid", ""}} {
		if _, err := normalizeOIDCScopes(scopes); err == nil {
			t.Fatalf("invalid scopes accepted: %#v", scopes)
		}
	}
}

func signTestIDToken(t *testing.T, key *rsa.PrivateKey, kid string, claims map[string]any) string {
	t.Helper()
	header, err := json.Marshal(map[string]string{"alg": "RS256", "kid": kid, "typ": "JWT"})
	if err != nil {
		t.Fatal(err)
	}
	payload, err := json.Marshal(claims)
	if err != nil {
		t.Fatal(err)
	}
	signingInput := base64.RawURLEncoding.EncodeToString(header) + "." + base64.RawURLEncoding.EncodeToString(payload)
	digest := sha256.Sum256([]byte(signingInput))
	signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
	if err != nil {
		t.Fatal(err)
	}
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(signature)
}

func cloneClaims(source map[string]any) map[string]any {
	result := make(map[string]any, len(source))
	for key, value := range source {
		result[key] = value
	}
	return result
}