knowfoolery/backend/shared/infra/auth/zitadel/client_test.go

package zitadel

// Tests for Zitadel client user info calls, JWT validation, and token operations.

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

	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/require"
)

type jwksKey struct {
	Kty string `json:"kty"`
	Kid string `json:"kid"`
	Use string `json:"use"`
	Alg string `json:"alg"`
	N   string `json:"n"`
	E   string `json:"e"`
}

type jwks struct {
	Keys []jwksKey `json:"keys"`
}

// generateJWKS is a test helper.
func generateJWKS(t *testing.T) (*rsa.PrivateKey, jwks, string) {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)

	kid := "kid-1"
	n := base64.RawURLEncoding.EncodeToString(key.N.Bytes())
	e := base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key.E)).Bytes())

	return key, jwks{
		Keys: []jwksKey{
			{
				Kty: "RSA",
				Kid: kid,
				Use: "sig",
				Alg: "RS256",
				N:   n,
				E:   e,
			},
		},
	}, kid
}

// signToken is a test helper.
func signToken(t *testing.T, key *rsa.PrivateKey, kid string, claims jwt.MapClaims) string {
	t.Helper()
	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	token.Header["kid"] = kid
	signed, err := token.SignedString(key)
	require.NoError(t, err)
	return signed
}

// newOIDCServer is a test helper.
func newOIDCServer(t *testing.T, jwksDoc jwks) *httptest.Server {
	t.Helper()
	var baseURL string

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/.well-known/openid-configuration":
			_ = json.NewEncoder(w).Encode(map[string]string{
				"issuer":              baseURL,
				"jwks_uri":            baseURL + "/jwks",
				"token_endpoint":      baseURL + "/token",
				"revocation_endpoint": baseURL + "/revoke",
			})
		case "/jwks":
			_ = json.NewEncoder(w).Encode(jwksDoc)
		case "/token":
			if err := r.ParseForm(); err != nil {
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			if r.FormValue("grant_type") != "refresh_token" {
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			_ = json.NewEncoder(w).Encode(TokenResponse{
				AccessToken:  "new-access",
				RefreshToken: "new-refresh",
				TokenType:    "Bearer",
				ExpiresIn:    3600,
			})
		case "/revoke":
			w.WriteHeader(http.StatusOK)
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	})

	server := httptest.NewServer(handler)
	baseURL = server.URL
	return server
}

// TestGetUserInfo_Success verifies user info retrieval on a 200 response.
func TestGetUserInfo_Success(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "Bearer token", r.Header.Get("Authorization"))
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(UserInfo{
			ID:       "user-1",
			Email:    "a@b.com",
			Name:     "Alice",
			Verified: true,
			Roles:    []string{"player"},
		})
	}))
	defer server.Close()

	client := NewClient(Config{BaseURL: server.URL, Timeout: 2 * time.Second})
	info, err := client.GetUserInfo(context.Background(), "token")
	require.NoError(t, err)
	require.Equal(t, "user-1", info.ID)
}

// TestGetUserInfo_NonOK verifies non-200 responses return an error.
func TestGetUserInfo_NonOK(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	defer server.Close()

	client := NewClient(Config{BaseURL: server.URL, Timeout: 2 * time.Second})
	_, err := client.GetUserInfo(context.Background(), "token")
	require.Error(t, err)
}

// TestValidateToken_Success verifies JWT validation using JWKS and discovery.
func TestValidateToken_Success(t *testing.T) {
	key, jwksDoc, kid := generateJWKS(t)
	server := newOIDCServer(t, jwksDoc)
	defer server.Close()

	claims := jwt.MapClaims{
		"sub":   "user-1",
		"email": "a@b.com",
		"name":  "Alice",
		"aud":   "client-1",
		"iss":   server.URL,
		"iat":   time.Now().Unix(),
		"exp":   time.Now().Add(10 * time.Minute).Unix(),
		"urn:zitadel:iam:org:project:roles": map[string]interface{}{
			"admin": map[string]interface{}{},
		},
		"amr": []string{"otp"},
	}
	token := signToken(t, key, kid, claims)

	client := NewClient(Config{BaseURL: server.URL, Timeout: 2 * time.Second})
	parsed, err := client.ValidateToken(context.Background(), token, ValidationOptions{
		Issuer:   server.URL,
		Audience: "client-1",
	})
	require.NoError(t, err)
	require.Equal(t, "user-1", parsed.Subject)
	require.True(t, parsed.MFAVerified)
	require.Contains(t, parsed.Roles, "admin")
}

// TestValidateToken_InvalidAudience verifies audience validation.
func TestValidateToken_InvalidAudience(t *testing.T) {
	key, jwksDoc, kid := generateJWKS(t)
	server := newOIDCServer(t, jwksDoc)
	defer server.Close()

	token := signToken(t, key, kid, jwt.MapClaims{
		"sub": "user-1",
		"aud": "other",
		"iss": server.URL,
		"exp": time.Now().Add(5 * time.Minute).Unix(),
	})

	client := NewClient(Config{BaseURL: server.URL, Timeout: 2 * time.Second})
	_, err := client.ValidateToken(context.Background(), token, ValidationOptions{
		Issuer:   server.URL,
		Audience: "client-1",
	})
	require.Error(t, err)
}

// TestValidateToken_MissingRequiredClaim verifies required claim enforcement.
func TestValidateToken_MissingRequiredClaim(t *testing.T) {
	key, jwksDoc, kid := generateJWKS(t)
	server := newOIDCServer(t, jwksDoc)
	defer server.Close()

	token := signToken(t, key, kid, jwt.MapClaims{
		"sub": "user-1",
		"aud": "client-1",
		"iss": server.URL,
		"exp": time.Now().Add(5 * time.Minute).Unix(),
	})

	client := NewClient(Config{BaseURL: server.URL, Timeout: 2 * time.Second})
	_, err := client.ValidateToken(context.Background(), token, ValidationOptions{
		Issuer:         server.URL,
		Audience:       "client-1",
		RequiredClaims: []string{"email"},
	})
	require.Error(t, err)
}

// TestRefreshToken_Success verifies token refresh against discovery endpoint.
func TestRefreshToken_Success(t *testing.T) {
	key, jwksDoc, _ := generateJWKS(t)
	_ = key
	server := newOIDCServer(t, jwksDoc)
	defer server.Close()

	client := NewClient(Config{
		BaseURL:      server.URL,
		ClientID:     "client",
		ClientSecret: "secret",
		Timeout:      2 * time.Second,
	})

	resp, err := client.RefreshToken(context.Background(), "refresh-token")
	require.NoError(t, err)
	require.Equal(t, "new-access", resp.AccessToken)
}

// TestRevokeToken_Success verifies token revocation against discovery endpoint.
func TestRevokeToken_Success(t *testing.T) {
	key, jwksDoc, _ := generateJWKS(t)
	_ = key
	server := newOIDCServer(t, jwksDoc)
	defer server.Close()

	client := NewClient(Config{
		BaseURL:      server.URL,
		ClientID:     "client",
		ClientSecret: "secret",
		Timeout:      2 * time.Second,
	})

	err := client.RevokeToken(context.Background(), "token")
	require.NoError(t, err)
}

// TestRefreshToken_UsesForm ensures refresh uses form-encoded requests.
func TestRefreshToken_UsesForm(t *testing.T) {
	key, jwksDoc, _ := generateJWKS(t)
	_ = key

	var captured string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/.well-known/openid-configuration" {
			_ = json.NewEncoder(w).Encode(map[string]string{
				"issuer":         "test",
				"jwks_uri":       serverURL(r),
				"token_endpoint": serverURL(r) + "/token",
			})
			return
		}
		if r.URL.Path == "/jwks" {
			_ = json.NewEncoder(w).Encode(jwksDoc)
			return
		}
		if r.URL.Path == "/token" {
			body, _ := io.ReadAll(r.Body)
			captured = string(body)
			_ = json.NewEncoder(w).Encode(TokenResponse{
				AccessToken: "access",
			})
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	defer server.Close()

	client := NewClient(Config{
		BaseURL:      server.URL,
		ClientID:     "client",
		ClientSecret: "secret",
		Timeout:      2 * time.Second,
	})
	_, err := client.RefreshToken(context.Background(), "refresh-token")
	require.NoError(t, err)
	require.Contains(t, captured, url.QueryEscape("refresh-token"))
	require.Contains(t, captured, "client_id=client")
}

// serverURL is a test helper.
func serverURL(r *http.Request) string {
	return "http://" + r.Host
}