key-cape/src/internal/server/oidc/token_test.go

package oidc_test

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"keycape/internal/domain"
	profileerrors "keycape/internal/errors"
	"keycape/internal/server/oidc"
	"keycape/internal/server/telemetry"
)

// ---------------------------------------------------------------------------
// Mock UserRepository
// ---------------------------------------------------------------------------

type mockUserRepo struct {
	users map[string]*domain.User
}

func (m *mockUserRepo) LookupUser(_ context.Context, username string) (*domain.User, error) {
	u, ok := m.users[username]
	if !ok {
		return nil, domain.ErrUserNotFound
	}
	return u, nil
}

func (m *mockUserRepo) LookupGroups(_ context.Context, _ string) ([]domain.Group, error) {
	return nil, nil
}

func (m *mockUserRepo) ValidatePassword(_ context.Context, _, _ string) (bool, error) {
	return false, nil
}

func (m *mockUserRepo) ListUsers(_ context.Context) ([]domain.User, error) {
	users := make([]domain.User, 0, len(m.users))
	for _, u := range m.users {
		users = append(users, *u)
	}
	return users, nil
}

// ---------------------------------------------------------------------------
// PKCE helpers
// ---------------------------------------------------------------------------

// makeVerifierAndChallenge returns a code_verifier and its S256 code_challenge.
func makeVerifierAndChallenge() (verifier, challenge string) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	verifier = base64.RawURLEncoding.EncodeToString(b)
	h := sha256.New()
	h.Write([]byte(verifier))
	challenge = base64.RawURLEncoding.EncodeToString(h.Sum(nil))
	return
}

// ---------------------------------------------------------------------------
// Test helpers
// ---------------------------------------------------------------------------

func newTokenHandler(t *testing.T, sessions *oidc.SessionStore, users domain.UserRepository) (*oidc.TokenHandler, *rsa.PrivateKey) {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	emitter := &captureEmitter{}
	h := &oidc.TokenHandler{
		ClientConfig: testClient(),
		Sessions:     sessions,
		Users:        users,
		SigningKey:    key,
		Issuer:       "https://auth.netkingdom.local",
		TokenLifetime: 15 * time.Minute,
		Emitter:      emitter,
	}
	return h, key
}

func tokenRequest(params url.Values) *http.Request {
	req := httptest.NewRequest(http.MethodPost, "/token",
		strings.NewReader(params.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	return req
}

func seededSession(sessions *oidc.SessionStore, verifier string) (code string) {
	challenge := s256Challenge(verifier)
	sess := &oidc.PKCESession{
		ClientID:            "test-client",
		RedirectURI:         "https://app.example.com/callback",
		PKCEChallenge:       challenge,
		PKCEChallengeMethod: "S256",
		State:               "state1",
		Nonce:               "nonce1",
		Username:            "alice",
		Scopes:              []string{"openid", "profile", "email", "groups"},
		ExpiresAt:           time.Now().Add(10 * time.Minute),
	}
	return sessions.Create(sess)
}

func s256Challenge(verifier string) string {
	h := sha256.New()
	h.Write([]byte(verifier))
	return base64.RawURLEncoding.EncodeToString(h.Sum(nil))
}

func decodeTokenResponse(t *testing.T, body string) map[string]interface{} {
	t.Helper()
	var m map[string]interface{}
	if err := json.Unmarshal([]byte(body), &m); err != nil {
		t.Fatalf("could not decode token response: %v (body: %q)", err, body)
	}
	return m
}

func parseJWTPayload(t *testing.T, token string) map[string]interface{} {
	t.Helper()
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		t.Fatalf("expected 3 JWT parts, got %d", len(parts))
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		t.Fatalf("decode JWT payload: %v", err)
	}
	var claims map[string]interface{}
	if err := json.Unmarshal(payload, &claims); err != nil {
		t.Fatalf("unmarshal JWT payload: %v", err)
	}
	return claims
}

func aliceUser() *domain.User {
	return &domain.User{
		ID:       "user-alice",
		Username: "alice",
		Email:    "alice@example.com",
		Groups:   []string{"admin", "users"},
		Enabled:  true,
	}
}

// ---------------------------------------------------------------------------
// T07 Token Endpoint Tests
// ---------------------------------------------------------------------------

func TestTokenHandler_ValidExchange_ReturnsJWT(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{users: map[string]*domain.User{"alice": aliceUser()}}

	h, _ := newTokenHandler(t, sessions, users)

	verifier := "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
	code := seededSession(sessions, verifier)

	params := url.Values{
		"grant_type":    []string{"authorization_code"},
		"code":          []string{code},
		"client_id":     []string{"test-client"},
		"redirect_uri":  []string{"https://app.example.com/callback"},
		"code_verifier": []string{verifier},
	}

	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d (body: %s)", w.Code, w.Body.String())
	}

	resp := decodeTokenResponse(t, w.Body.String())
	if _, ok := resp["access_token"]; !ok {
		t.Error("missing access_token")
	}
	if _, ok := resp["id_token"]; !ok {
		t.Error("missing id_token")
	}
	if resp["token_type"] != "Bearer" {
		t.Errorf("expected token_type Bearer, got %v", resp["token_type"])
	}
	if _, ok := resp["expires_in"]; !ok {
		t.Error("missing expires_in")
	}
}

func TestTokenHandler_WrongGrantType_FeatureNotSupported(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{}

	h, _ := newTokenHandler(t, sessions, users)

	params := url.Values{
		"grant_type": []string{"client_credentials"},
		"client_id":  []string{"test-client"},
	}

	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusBadRequest {
		t.Errorf("expected 400, got %d", w.Code)
	}
	errType := decodeProfileError(t, w.Body.String())
	if errType != profileerrors.ErrFeatureNotSupported {
		t.Errorf("expected feature_not_supported_by_profile, got %q", errType)
	}
}

func TestTokenHandler_PKCEMismatch_InvalidProfileUsage(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{users: map[string]*domain.User{"alice": aliceUser()}}

	h, _ := newTokenHandler(t, sessions, users)

	realVerifier := "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
	code := seededSession(sessions, realVerifier)

	params := url.Values{
		"grant_type":    []string{"authorization_code"},
		"code":          []string{code},
		"client_id":     []string{"test-client"},
		"redirect_uri":  []string{"https://app.example.com/callback"},
		"code_verifier": []string{"wrong-verifier-that-does-not-match"},
	}

	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusBadRequest {
		t.Errorf("expected 400, got %d", w.Code)
	}
	errType := decodeProfileError(t, w.Body.String())
	if errType != profileerrors.ErrInvalidProfileUsage {
		t.Errorf("expected invalid_profile_usage, got %q", errType)
	}
}

func TestTokenHandler_CodeNotFound_InvalidProfileUsage(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{}

	h, _ := newTokenHandler(t, sessions, users)

	params := url.Values{
		"grant_type":    []string{"authorization_code"},
		"code":          []string{"no-such-code"},
		"client_id":     []string{"test-client"},
		"redirect_uri":  []string{"https://app.example.com/callback"},
		"code_verifier": []string{"any-verifier"},
	}

	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusBadRequest {
		t.Errorf("expected 400, got %d", w.Code)
	}
	errType := decodeProfileError(t, w.Body.String())
	if errType != profileerrors.ErrInvalidProfileUsage {
		t.Errorf("expected invalid_profile_usage, got %q", errType)
	}
}

func TestTokenHandler_JWTClaims_CorrectSubAndIssuer(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{users: map[string]*domain.User{"alice": aliceUser()}}

	h, _ := newTokenHandler(t, sessions, users)

	verifier := "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
	code := seededSession(sessions, verifier)

	params := url.Values{
		"grant_type":    []string{"authorization_code"},
		"code":          []string{code},
		"client_id":     []string{"test-client"},
		"redirect_uri":  []string{"https://app.example.com/callback"},
		"code_verifier": []string{verifier},
	}

	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}

	resp := decodeTokenResponse(t, w.Body.String())
	idToken, ok := resp["id_token"].(string)
	if !ok {
		t.Fatal("id_token is not a string")
	}

	claims := parseJWTPayload(t, idToken)

	if claims["sub"] != "user-alice" {
		t.Errorf("sub: expected user-alice, got %v", claims["sub"])
	}
	if claims["iss"] != "https://auth.netkingdom.local" {
		t.Errorf("iss: expected https://auth.netkingdom.local, got %v", claims["iss"])
	}
	if claims["aud"] != "test-client" {
		t.Errorf("aud: expected test-client, got %v", claims["aud"])
	}
	if claims["nonce"] != "nonce1" {
		t.Errorf("nonce: expected nonce1, got %v", claims["nonce"])
	}
}

func TestTokenHandler_ScopeFiltering_ProfileScope(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{users: map[string]*domain.User{"alice": aliceUser()}}

	h, _ := newTokenHandler(t, sessions, users)

	// Seed session with only openid scope (no email, no groups).
	verifier := "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
	challenge := s256Challenge(verifier)
	sess := &oidc.PKCESession{
		ClientID:            "test-client",
		RedirectURI:         "https://app.example.com/callback",
		PKCEChallenge:       challenge,
		PKCEChallengeMethod: "S256",
		Username:            "alice",
		Scopes:              []string{"openid"}, // no profile/email/groups
		ExpiresAt:           time.Now().Add(10 * time.Minute),
	}
	code := sessions.Create(sess)

	params := url.Values{
		"grant_type":    []string{"authorization_code"},
		"code":          []string{code},
		"client_id":     []string{"test-client"},
		"redirect_uri":  []string{"https://app.example.com/callback"},
		"code_verifier": []string{verifier},
	}

	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}

	resp := decodeTokenResponse(t, w.Body.String())
	idToken := resp["id_token"].(string)
	claims := parseJWTPayload(t, idToken)

	// Without profile scope, preferred_username must not be present.
	if _, ok := claims["preferred_username"]; ok {
		t.Error("preferred_username must be absent when profile scope is not granted")
	}
	// Without email scope, email must not be present.
	if _, ok := claims["email"]; ok {
		t.Error("email must be absent when email scope is not granted")
	}
	// Without groups scope, groups must not be present.
	if _, ok := claims["groups"]; ok {
		t.Error("groups must be absent when groups scope is not granted")
	}
}

func TestTokenHandler_ScopeFiltering_AllScopes(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{users: map[string]*domain.User{"alice": aliceUser()}}

	h, _ := newTokenHandler(t, sessions, users)

	verifier := "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
	code := seededSession(sessions, verifier) // has openid, profile, email, groups

	params := url.Values{
		"grant_type":    []string{"authorization_code"},
		"code":          []string{code},
		"client_id":     []string{"test-client"},
		"redirect_uri":  []string{"https://app.example.com/callback"},
		"code_verifier": []string{verifier},
	}

	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}

	resp := decodeTokenResponse(t, w.Body.String())
	idToken := resp["id_token"].(string)
	claims := parseJWTPayload(t, idToken)

	if claims["preferred_username"] != "alice" {
		t.Errorf("preferred_username: expected alice, got %v", claims["preferred_username"])
	}
	if claims["email"] != "alice@example.com" {
		t.Errorf("email: expected alice@example.com, got %v", claims["email"])
	}
	if _, ok := claims["groups"]; !ok {
		t.Error("groups claim must be present when groups scope is granted")
	}
}

func TestTokenHandler_TokenIssuedTelemetry(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{users: map[string]*domain.User{"alice": aliceUser()}}

	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	capture := &captureEmitter{}
	h := &oidc.TokenHandler{
		ClientConfig:  testClient(),
		Sessions:      sessions,
		Users:         users,
		SigningKey:     key,
		Issuer:        "https://auth.netkingdom.local",
		TokenLifetime: 15 * time.Minute,
		Emitter:       capture,
	}

	verifier := "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
	code := seededSession(sessions, verifier)

	params := url.Values{
		"grant_type":    []string{"authorization_code"},
		"code":          []string{code},
		"client_id":     []string{"test-client"},
		"redirect_uri":  []string{"https://app.example.com/callback"},
		"code_verifier": []string{verifier},
	}

	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}

	found := false
	for _, ev := range capture.events {
		if ev.EventType == telemetry.EventTokenIssued {
			found = true
			break
		}
	}
	if !found {
		t.Error("expected token_issued telemetry event")
	}
}

func TestTokenHandler_CodeDeletedAfterUse(t *testing.T) {
	sessions := oidc.NewSessionStore()
	users := &mockUserRepo{users: map[string]*domain.User{"alice": aliceUser()}}

	h, _ := newTokenHandler(t, sessions, users)

	verifier := "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
	code := seededSession(sessions, verifier)

	params := url.Values{
		"grant_type":    []string{"authorization_code"},
		"code":          []string{code},
		"client_id":     []string{"test-client"},
		"redirect_uri":  []string{"https://app.example.com/callback"},
		"code_verifier": []string{verifier},
	}

	// First use — should succeed.
	req := tokenRequest(params)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("first use: expected 200, got %d", w.Code)
	}

	// Second use — code should be gone.
	req2 := tokenRequest(params)
	w2 := httptest.NewRecorder()
	h.ServeHTTP(w2, req2)
	if w2.Code != http.StatusBadRequest {
		t.Errorf("second use: expected 400 (code replay), got %d", w2.Code)
	}
}