pocket-id-pocket-id-1/backend/internal/service/jwt_service_test.go

package service

import (
	"context"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"sync"
	"testing"
	"time"

	"github.com/lestrrat-go/jwx/v3/jwa"
	"github.com/lestrrat-go/jwx/v3/jwk"
	"github.com/lestrrat-go/jwx/v3/jwt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"

	"github.com/pocket-id/pocket-id/backend/internal/common"
	"github.com/pocket-id/pocket-id/backend/internal/model"
	"github.com/pocket-id/pocket-id/backend/internal/utils"
	jwkutils "github.com/pocket-id/pocket-id/backend/internal/utils/jwk"
	testutils "github.com/pocket-id/pocket-id/backend/internal/utils/testing"
)

const testEncryptionKey = "0123456789abcdef0123456789abcdef"

func newTestEnvConfig() *common.EnvConfigSchema {
	return &common.EnvConfigSchema{
		AppURL:        "https://test.example.com",
		EncryptionKey: []byte(testEncryptionKey),
	}
}

func initJwtService(t *testing.T, db *gorm.DB, appConfig *AppConfigService, envConfig *common.EnvConfigSchema) *JwtService {
	t.Helper()

	service := &JwtService{}
	err := service.init(db, appConfig, envConfig)
	require.NoError(t, err, "Failed to initialize JWT service")

	return service
}

func setupJwtService(t *testing.T, appConfig *AppConfigService) (*JwtService, *gorm.DB, *common.EnvConfigSchema) {
	t.Helper()

	db := testutils.NewDatabaseForTest(t)
	envConfig := newTestEnvConfig()

	return initJwtService(t, db, appConfig, envConfig), db, envConfig
}

func newTestDbAndEnv(t *testing.T) (*gorm.DB, *common.EnvConfigSchema) {
	t.Helper()

	return testutils.NewDatabaseForTest(t), newTestEnvConfig()
}

func saveKeyToDatabase(t *testing.T, db *gorm.DB, envConfig *common.EnvConfigSchema, appConfig *AppConfigService, key jwk.Key) string {
	t.Helper()

	keyProvider, err := jwkutils.GetKeyProvider(db, envConfig, appConfig.GetDbConfig().InstanceID.Value)
	require.NoError(t, err, "Failed to init key provider")

	err = keyProvider.SaveKey(key)
	require.NoError(t, err, "Failed to save key")

	kid, ok := key.KeyID()
	require.True(t, ok, "Key ID must be set")
	require.NotEmpty(t, kid, "Key ID must not be empty")

	return kid
}

func TestJwtService_Init(t *testing.T) {
	mockConfig := NewTestAppConfigService(&model.AppConfig{
		SessionDuration: model.AppConfigVariable{Value: "60"}, // 60 minutes
	})

	t.Run("should generate new key when none exists", func(t *testing.T) {
		db := testutils.NewDatabaseForTest(t)
		mockEnvConfig := newTestEnvConfig()

		// Initialize the JWT service
		service := initJwtService(t, db, mockConfig, mockEnvConfig)

		// Verify the private key was set
		require.NotNil(t, service.privateKey, "Private key should be set")

		// Verify the key has been persisted in the database
		keyProvider, err := jwkutils.GetKeyProvider(db, mockEnvConfig, mockConfig.GetDbConfig().InstanceID.Value)
		require.NoError(t, err, "Failed to init key provider")
		key, err := keyProvider.LoadKey()
		require.NoError(t, err, "Failed to load key from provider")
		require.NotNil(t, key, "Key should be present in the database")

		// Key should have required properties
		keyID, ok := key.KeyID()
		assert.True(t, ok, "Key should have a key ID")
		assert.NotEmpty(t, keyID)

		keyUsage, ok := key.KeyUsage()
		assert.True(t, ok, "Key should have a key usage")
		assert.Equal(t, KeyUsageSigning, keyUsage)
	})

	t.Run("should load existing JWK key", func(t *testing.T) {
		db := testutils.NewDatabaseForTest(t)
		mockEnvConfig := newTestEnvConfig()

		// First create a service to generate a key
		firstService := initJwtService(t, db, mockConfig, mockEnvConfig)

		// Get the key ID of the first service
		origKeyID, ok := firstService.privateKey.KeyID()
		require.True(t, ok)

		// Now create a new service that should load the existing key
		secondService := initJwtService(t, db, mockConfig, mockEnvConfig)

		// Verify the loaded key has the same ID as the original
		loadedKeyID, ok := secondService.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")
	})

	t.Run("should load existing JWK for ECDSA keys", func(t *testing.T) {
		db := testutils.NewDatabaseForTest(t)
		mockEnvConfig := newTestEnvConfig()

		// Create a new JWK and save it to the database
		origKeyID := createECDSAKeyJWK(t, db, mockEnvConfig, mockConfig)

		// Now create a new service that should load the existing key
		svc := initJwtService(t, db, mockConfig, mockEnvConfig)

		// Ensure loaded key has the right algorithm
		alg, ok := svc.privateKey.Algorithm()
		_ = assert.True(t, ok) &&
			assert.Equal(t, jwa.ES256().String(), alg.String(), "Loaded key has the incorrect algorithm")

		// Verify the loaded key has the same ID as the original
		loadedKeyID, ok := svc.privateKey.KeyID()
		_ = assert.True(t, ok) &&
			assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")
	})

	t.Run("should load existing JWK for EdDSA keys", func(t *testing.T) {
		db := testutils.NewDatabaseForTest(t)
		mockEnvConfig := newTestEnvConfig()

		// Create a new JWK and save it to the database
		origKeyID := createEdDSAKeyJWK(t, db, mockEnvConfig, mockConfig)

		// Now create a new service that should load the existing key
		svc := initJwtService(t, db, mockConfig, mockEnvConfig)

		// Ensure loaded key has the right algorithm and curve
		alg, ok := svc.privateKey.Algorithm()
		_ = assert.True(t, ok) &&
			assert.Equal(t, jwa.EdDSA().String(), alg.String(), "Loaded key has the incorrect algorithm")

		var curve jwa.EllipticCurveAlgorithm
		err := svc.privateKey.Get("crv", &curve)
		_ = assert.NoError(t, err, "Failed to get 'crv' claim") &&
			assert.Equal(t, jwa.Ed25519().String(), curve.String(), "Curve does not match expected value")

		// Verify the loaded key has the same ID as the original
		loadedKeyID, ok := svc.privateKey.KeyID()
		_ = assert.True(t, ok) &&
			assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")
	})
}

func TestJwtService_GetPublicJWK(t *testing.T) {
	mockConfig := NewTestAppConfigService(&model.AppConfig{
		SessionDuration: model.AppConfigVariable{Value: "60"}, // 60 minutes
	})

	t.Run("returns public key when private key is initialized", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		// Get the JWK (public key)
		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err, "GetPublicJWK should not return an error when private key is initialized")

		// Verify the returned key is valid
		require.NotNil(t, publicKey, "Public key should not be nil")

		// Validate it's actually a public key
		isPrivate, err := jwk.IsPrivateKey(publicKey)
		require.NoError(t, err)
		assert.False(t, isPrivate, "Returned key should be a public key")

		// Check that key has required properties
		keyID, ok := publicKey.KeyID()
		require.True(t, ok, "Public key should have a key ID")
		assert.NotEmpty(t, keyID, "Key ID should not be empty")

		alg, ok := publicKey.Algorithm()
		require.True(t, ok, "Public key should have an algorithm")
		assert.Equal(t, "RS256", alg.String(), "Algorithm should be RS256")
	})

	t.Run("returns public key when ECDSA private key is initialized", func(t *testing.T) {
		db := testutils.NewDatabaseForTest(t)
		mockEnvConfig := newTestEnvConfig()

		// Create an ECDSA key and save it in the database
		originalKeyID := createECDSAKeyJWK(t, db, mockEnvConfig, mockConfig)

		// Create a JWT service that loads the ECDSA key
		service := initJwtService(t, db, mockConfig, mockEnvConfig)

		// Get the JWK (public key)
		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err, "GetPublicJWK should not return an error when private key is initialized")

		// Verify the returned key is valid
		require.NotNil(t, publicKey, "Public key should not be nil")

		// Validate it's actually a public key
		isPrivate, err := jwk.IsPrivateKey(publicKey)
		require.NoError(t, err)
		assert.False(t, isPrivate, "Returned key should be a public key")

		// Check that key has required properties
		keyID, ok := publicKey.KeyID()
		require.True(t, ok, "Public key should have a key ID")
		assert.Equal(t, originalKeyID, keyID, "Key ID should match the original key ID")

		// Check that the key type is EC
		assert.Equal(t, "EC", publicKey.KeyType().String(), "Key type should be EC")

		// Check that the algorithm is ES256
		alg, ok := publicKey.Algorithm()
		require.True(t, ok, "Public key should have an algorithm")
		assert.Equal(t, "ES256", alg.String(), "Algorithm should be ES256")
	})

	t.Run("returns public key when EdDSA private key is initialized", func(t *testing.T) {
		db := testutils.NewDatabaseForTest(t)
		mockEnvConfig := newTestEnvConfig()

		// Create an EdDSA key and save it in the database
		originalKeyID := createEdDSAKeyJWK(t, db, mockEnvConfig, mockConfig)

		// Create a JWT service that loads the EdDSA key
		service := initJwtService(t, db, mockConfig, mockEnvConfig)

		// Get the JWK (public key)
		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err, "GetPublicJWK should not return an error when private key is initialized")

		// Verify the returned key is valid
		require.NotNil(t, publicKey, "Public key should not be nil")

		// Validate it's actually a public key
		isPrivate, err := jwk.IsPrivateKey(publicKey)
		require.NoError(t, err)
		assert.False(t, isPrivate, "Returned key should be a public key")

		// Check that key has required properties
		keyID, ok := publicKey.KeyID()
		require.True(t, ok, "Public key should have a key ID")
		assert.Equal(t, originalKeyID, keyID, "Key ID should match the original key ID")

		// Check that the key type is OKP
		assert.Equal(t, "OKP", publicKey.KeyType().String(), "Key type should be OKP")

		// Check that the algorithm is EdDSA
		alg, ok := publicKey.Algorithm()
		require.True(t, ok, "Public key should have an algorithm")
		assert.Equal(t, "EdDSA", alg.String(), "Algorithm should be EdDSA")
	})

	t.Run("returns error when private key is not initialized", func(t *testing.T) {
		// Create a service with nil private key
		service := &JwtService{
			privateKey: nil,
		}

		// Try to get the JWK
		publicKey, err := service.GetPublicJWK()

		// Verify it returns an error
		require.Error(t, err, "GetPublicJWK should return an error when private key is nil")
		assert.Contains(t, err.Error(), "key is not initialized", "Error message should indicate key is not initialized")
		assert.Nil(t, publicKey, "Public key should be nil when there's an error")
	})
}

func TestGenerateVerifyAccessToken(t *testing.T) {
	mockConfig := NewTestAppConfigService(&model.AppConfig{
		SessionDuration: model.AppConfigVariable{Value: "60"}, // 60 minutes
	})

	t.Run("generates token for regular user", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		user := model.User{
			Base:    model.Base{ID: "user123"},
			Email:   utils.Ptr("user@example.com"),
			IsAdmin: false,
		}

		tokenString, err := service.GenerateAccessToken(user)
		require.NoError(t, err, "Failed to generate access token")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated token")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, user.ID, subject, "Token subject should match user ID")
		isAdmin, err := GetIsAdmin(claims)
		_ = assert.NoError(t, err, "Failed to get isAdmin claim") &&
			assert.False(t, isAdmin, "isAdmin should be false")
		audience, ok := claims.Audience()
		_ = assert.True(t, ok, "Audience not found in token") &&
			assert.Equal(t, []string{service.envConfig.AppURL}, audience, "Audience should contain the app URL")

		expectedExp := time.Now().Add(1 * time.Hour)
		expiration, ok := claims.Expiration()
		assert.True(t, ok, "Expiration not found in token")
		timeDiff := expectedExp.Sub(expiration).Minutes()
		assert.InDelta(t, 0, timeDiff, 1.0, "Token should expire in approximately 1 hour")
	})

	t.Run("generates token for admin user", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		adminUser := model.User{
			Base:    model.Base{ID: "admin123"},
			Email:   utils.Ptr("admin@example.com"),
			IsAdmin: true,
		}

		tokenString, err := service.GenerateAccessToken(adminUser)
		require.NoError(t, err, "Failed to generate access token")

		claims, err := service.VerifyAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated token")

		isAdmin, err := GetIsAdmin(claims)
		_ = assert.NoError(t, err, "Failed to get isAdmin claim") &&
			assert.True(t, isAdmin, "isAdmin should be true")
		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, adminUser.ID, subject, "Token subject should match user ID")
	})

	t.Run("uses session duration from config", func(t *testing.T) {
		customMockConfig := NewTestAppConfigService(&model.AppConfig{
			SessionDuration: model.AppConfigVariable{Value: "30"}, // 30 minutes
		})
		service, _, _ := setupJwtService(t, customMockConfig)

		user := model.User{
			Base: model.Base{ID: "user456"},
		}

		tokenString, err := service.GenerateAccessToken(user)
		require.NoError(t, err, "Failed to generate access token")

		claims, err := service.VerifyAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated token")

		expectedExp := time.Now().Add(30 * time.Minute)
		expiration, ok := claims.Expiration()
		assert.True(t, ok, "Expiration not found in token")
		timeDiff := expectedExp.Sub(expiration).Minutes()
		assert.InDelta(t, 0, timeDiff, 1.0, "Token should expire in approximately 30 minutes")
	})

	t.Run("works with Ed25519 keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createEdDSAKeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		user := model.User{
			Base:    model.Base{ID: "eddsauser123"},
			Email:   utils.Ptr("eddsauser@example.com"),
			IsAdmin: true,
		}

		tokenString, err := service.GenerateAccessToken(user)
		require.NoError(t, err, "Failed to generate access token with Ed25519 key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated token with Ed25519 key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, user.ID, subject, "Token subject should match user ID")
		isAdmin, err := GetIsAdmin(claims)
		_ = assert.NoError(t, err, "Failed to get isAdmin claim") &&
			assert.True(t, isAdmin, "isAdmin should be true")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err)
		assert.Equal(t, "OKP", publicKey.KeyType().String(), "Key type should be OKP")
		alg, ok := publicKey.Algorithm()
		require.True(t, ok)
		assert.Equal(t, "EdDSA", alg.String(), "Algorithm should be EdDSA")
	})

	t.Run("works with P-256 keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createECDSAKeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		user := model.User{
			Base:    model.Base{ID: "ecdsauser123"},
			Email:   utils.Ptr("ecdsauser@example.com"),
			IsAdmin: true,
		}

		tokenString, err := service.GenerateAccessToken(user)
		require.NoError(t, err, "Failed to generate access token with ECDSA key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated token with ECDSA key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, user.ID, subject, "Token subject should match user ID")
		isAdmin, err := GetIsAdmin(claims)
		_ = assert.NoError(t, err, "Failed to get isAdmin claim") &&
			assert.True(t, isAdmin, "isAdmin should be true")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err)
		assert.Equal(t, "EC", publicKey.KeyType().String(), "Key type should be EC")
		alg, ok := publicKey.Algorithm()
		require.True(t, ok)
		assert.Equal(t, "ES256", alg.String(), "Algorithm should be ES256")
	})

	t.Run("works with RSA-4096 keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createRSA4096KeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		user := model.User{
			Base:    model.Base{ID: "rsauser123"},
			Email:   utils.Ptr("rsauser@example.com"),
			IsAdmin: true,
		}

		tokenString, err := service.GenerateAccessToken(user)
		require.NoError(t, err, "Failed to generate access token with RSA key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated token with RSA key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, user.ID, subject, "Token subject should match user ID")
		isAdmin, err := GetIsAdmin(claims)
		_ = assert.NoError(t, err, "Failed to get isAdmin claim") &&
			assert.True(t, isAdmin, "isAdmin should be true")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err)
		assert.Equal(t, jwa.RSA().String(), publicKey.KeyType().String(), "Key type should be RSA")
		alg, ok := publicKey.Algorithm()
		require.True(t, ok)
		assert.Equal(t, jwa.RS256().String(), alg.String(), "Algorithm should be RS256")
	})
}

func TestGenerateVerifyIdToken(t *testing.T) {
	mockConfig := NewTestAppConfigService(&model.AppConfig{
		SessionDuration: model.AppConfigVariable{Value: "60"}, // 60 minutes
	})

	t.Run("generates and verifies ID token with standard claims", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		userClaims := map[string]interface{}{
			"sub":   "user123",
			"name":  "Test User",
			"email": "user@example.com",
		}
		const clientID = "test-client-123"

		tokenString, err := service.GenerateIDToken(userClaims, clientID, "")
		require.NoError(t, err, "Failed to generate ID token")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyIdToken(tokenString, false)
		require.NoError(t, err, "Failed to verify generated ID token")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, "user123", subject, "Token subject should match user ID")
		audience, ok := claims.Audience()
		_ = assert.True(t, ok, "Audience not found in token") &&
			assert.Equal(t, []string{clientID}, audience, "Audience should contain the client ID")
		issuer, ok := claims.Issuer()
		_ = assert.True(t, ok, "Issuer not found in token") &&
			assert.Equal(t, service.envConfig.AppURL, issuer, "Issuer should match app URL")

		expectedExp := time.Now().Add(1 * time.Hour)
		expiration, ok := claims.Expiration()
		assert.True(t, ok, "Expiration not found in token")
		timeDiff := expectedExp.Sub(expiration).Minutes()
		assert.InDelta(t, 0, timeDiff, 1.0, "Token should expire in approximately 1 hour")
	})

	t.Run("can accept expired tokens if told so", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		userClaims := map[string]interface{}{
			"sub":   "user123",
			"name":  "Test User",
			"email": "user@example.com",
		}
		const clientID = "test-client-123"

		token, err := jwt.NewBuilder().
			Subject(userClaims["sub"].(string)).
			Issuer(service.envConfig.AppURL).
			Audience([]string{clientID}).
			IssuedAt(time.Now().Add(-2 * time.Hour)).
			Expiration(time.Now().Add(-1 * time.Hour)).
			Build()
		require.NoError(t, err, "Failed to build token")

		err = SetTokenType(token, IDTokenJWTType)
		require.NoError(t, err, "Failed to set token type")

		for k, v := range userClaims {
			if k != "sub" {
				err = token.Set(k, v)
				require.NoError(t, err, "Failed to set claim")
			}
		}

		signed, err := jwt.Sign(token, jwt.WithKey(jwa.RS256(), service.privateKey))
		require.NoError(t, err, "Failed to sign token")
		tokenString := string(signed)

		_, err = service.VerifyIdToken(tokenString, false)
		require.Error(t, err, "Verification should fail with expired token when not allowing expired tokens")
		assert.Contains(t, err.Error(), "\"exp\" not satisfied", "Error message should indicate token verification failure")

		claims, err := service.VerifyIdToken(tokenString, true)
		require.NoError(t, err, "Verification should succeed with expired token when allowing expired tokens")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, userClaims["sub"], subject, "Token subject should match user ID")
		issuer, ok := claims.Issuer()
		_ = assert.True(t, ok, "Issuer not found in token") &&
			assert.Equal(t, service.envConfig.AppURL, issuer, "Issuer should match app URL")
	})

	t.Run("generates and verifies ID token with nonce", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		userClaims := map[string]interface{}{
			"sub":  "user456",
			"name": "Another User",
		}
		const clientID = "test-client-456"
		nonce := "random-nonce-value"

		tokenString, err := service.GenerateIDToken(userClaims, clientID, nonce)
		require.NoError(t, err, "Failed to generate ID token with nonce")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err, "Failed to get public key")
		token, err := jwt.Parse([]byte(tokenString), jwt.WithKey(jwa.RS256(), publicKey))
		require.NoError(t, err, "Failed to parse token")

		var tokenNonce string
		err = token.Get("nonce", &tokenNonce)
		require.NoError(t, err, "Failed to get claims")

		assert.Equal(t, nonce, tokenNonce, "Token should contain the correct nonce")
	})

	t.Run("fails verification with incorrect issuer", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		userClaims := map[string]interface{}{
			"sub": "user789",
		}
		tokenString, err := service.GenerateIDToken(userClaims, "client-789", "")
		require.NoError(t, err, "Failed to generate ID token")

		service.envConfig.AppURL = "https://wrong-issuer.com"

		_, err = service.VerifyIdToken(tokenString, false)
		require.Error(t, err, "Verification should fail with incorrect issuer")
		assert.Contains(t, err.Error(), "\"iss\" not satisfied", "Error message should indicate token verification failure")
	})

	t.Run("works with Ed25519 keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createEdDSAKeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		userClaims := map[string]interface{}{
			"sub":   "eddsauser456",
			"name":  "EdDSA User",
			"email": "eddsauser@example.com",
		}
		const clientID = "eddsa-client-123"

		tokenString, err := service.GenerateIDToken(userClaims, clientID, "")
		require.NoError(t, err, "Failed to generate ID token with key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyIdToken(tokenString, false)
		require.NoError(t, err, "Failed to verify generated ID token with key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, "eddsauser456", subject, "Token subject should match user ID")
		issuer, ok := claims.Issuer()
		_ = assert.True(t, ok, "Issuer not found in token") &&
			assert.Equal(t, service.envConfig.AppURL, issuer, "Issuer should match app URL")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err)
		assert.Equal(t, jwa.OKP().String(), publicKey.KeyType().String(), "Key type should be OKP")
		alg, ok := publicKey.Algorithm()
		require.True(t, ok)
		assert.Equal(t, jwa.EdDSA().String(), alg.String(), "Algorithm should be EdDSA")
	})

	t.Run("works with P-256 keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createECDSAKeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		userClaims := map[string]interface{}{
			"sub":   "ecdsauser456",
			"email": "ecdsauser@example.com",
		}
		const clientID = "ecdsa-client-123"

		tokenString, err := service.GenerateIDToken(userClaims, clientID, "")
		require.NoError(t, err, "Failed to generate ID token with key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyIdToken(tokenString, false)
		require.NoError(t, err, "Failed to verify generated ID token with key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, "ecdsauser456", subject, "Token subject should match user ID")
		issuer, ok := claims.Issuer()
		_ = assert.True(t, ok, "Issuer not found in token") &&
			assert.Equal(t, service.envConfig.AppURL, issuer, "Issuer should match app URL")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err)
		assert.Equal(t, jwa.EC().String(), publicKey.KeyType().String(), "Key type should be EC")
		alg, ok := publicKey.Algorithm()
		require.True(t, ok)
		assert.Equal(t, jwa.ES256().String(), alg.String(), "Algorithm should be ES256")
	})

	t.Run("works with RSA-4096 keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createRSA4096KeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		userClaims := map[string]interface{}{
			"sub":   "rsauser456",
			"name":  "RSA User",
			"email": "rsauser@example.com",
		}
		const clientID = "rsa-client-123"

		tokenString, err := service.GenerateIDToken(userClaims, clientID, "")
		require.NoError(t, err, "Failed to generate ID token with key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyIdToken(tokenString, false)
		require.NoError(t, err, "Failed to verify generated ID token with key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, "rsauser456", subject, "Token subject should match user ID")
		issuer, ok := claims.Issuer()
		_ = assert.True(t, ok, "Issuer not found in token") &&
			assert.Equal(t, service.envConfig.AppURL, issuer, "Issuer should match app URL")
	})
}

func TestGenerateVerifyOAuthAccessToken(t *testing.T) {
	mockConfig := NewTestAppConfigService(&model.AppConfig{
		SessionDuration: model.AppConfigVariable{Value: "60"}, // 60 minutes
	})

	t.Run("generates and verifies OAuth access token with standard claims", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		user := model.User{
			Base:  model.Base{ID: "user123"},
			Email: utils.Ptr("user@example.com"),
		}
		const clientID = "test-client-123"

		tokenString, err := service.GenerateOAuthAccessToken(user, clientID)
		require.NoError(t, err, "Failed to generate OAuth access token")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyOAuthAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated OAuth access token")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, user.ID, subject, "Token subject should match user ID")
		audience, ok := claims.Audience()
		_ = assert.True(t, ok, "Audience not found in token") &&
			assert.Equal(t, []string{clientID}, audience, "Audience should contain the client ID")
		issuer, ok := claims.Issuer()
		_ = assert.True(t, ok, "Issuer not found in token") &&
			assert.Equal(t, service.envConfig.AppURL, issuer, "Issuer should match app URL")

		expectedExp := time.Now().Add(1 * time.Hour)
		expiration, ok := claims.Expiration()
		assert.True(t, ok, "Expiration not found in token")
		timeDiff := expectedExp.Sub(expiration).Minutes()
		assert.InDelta(t, 0, timeDiff, 1.0, "Token should expire in approximately 1 hour")
	})

	t.Run("fails verification for expired token", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		user := model.User{Base: model.Base{ID: "user456"}}
		const clientID = "test-client-456"

		token, err := jwt.NewBuilder().
			Subject(user.ID).
			Expiration(time.Now().Add(-1 * time.Hour)).
			IssuedAt(time.Now().Add(-2 * time.Hour)).
			Audience([]string{clientID}).
			Issuer(service.envConfig.AppURL).
			Build()
		require.NoError(t, err, "Failed to build token")

		err = SetTokenType(token, OAuthAccessTokenJWTType)
		require.NoError(t, err, "Failed to set token type")

		signed, err := jwt.Sign(token, jwt.WithKey(jwa.RS256(), service.privateKey))
		require.NoError(t, err, "Failed to sign token")

		_, err = service.VerifyOAuthAccessToken(string(signed))
		require.Error(t, err, "Verification should fail with expired token")
		assert.Contains(t, err.Error(), "\"exp\" not satisfied", "Error message should indicate token verification failure")
	})

	t.Run("fails verification with invalid signature", func(t *testing.T) {
		service1, _, _ := setupJwtService(t, mockConfig)
		service2, _, _ := setupJwtService(t, mockConfig)

		user := model.User{Base: model.Base{ID: "user789"}}
		const clientID = "test-client-789"

		tokenString, err := service1.GenerateOAuthAccessToken(user, clientID)
		require.NoError(t, err, "Failed to generate OAuth access token")

		_, err = service2.VerifyOAuthAccessToken(tokenString)
		require.Error(t, err, "Verification should fail with invalid signature")
		assert.Contains(t, err.Error(), "verification error", "Error message should indicate token verification failure")
	})

	t.Run("works with Ed25519 keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createEdDSAKeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		user := model.User{
			Base:  model.Base{ID: "eddsauser789"},
			Email: utils.Ptr("eddsaoauth@example.com"),
		}
		const clientID = "eddsa-oauth-client"

		tokenString, err := service.GenerateOAuthAccessToken(user, clientID)
		require.NoError(t, err, "Failed to generate OAuth access token with key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyOAuthAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated OAuth access token with key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, user.ID, subject, "Token subject should match user ID")
		audience, ok := claims.Audience()
		_ = assert.True(t, ok, "Audience not found in token") &&
			assert.Equal(t, []string{clientID}, audience, "Audience should contain the client ID")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err)
		assert.Equal(t, jwa.OKP().String(), publicKey.KeyType().String(), "Key type should be OKP")
		alg, ok := publicKey.Algorithm()
		require.True(t, ok)
		assert.Equal(t, jwa.EdDSA().String(), alg.String(), "Algorithm should be EdDSA")
	})

	t.Run("works with ECDSA keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createECDSAKeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		user := model.User{
			Base:  model.Base{ID: "ecdsauser789"},
			Email: utils.Ptr("ecdsaoauth@example.com"),
		}
		const clientID = "ecdsa-oauth-client"

		tokenString, err := service.GenerateOAuthAccessToken(user, clientID)
		require.NoError(t, err, "Failed to generate OAuth access token with key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyOAuthAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated OAuth access token with key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, user.ID, subject, "Token subject should match user ID")
		audience, ok := claims.Audience()
		_ = assert.True(t, ok, "Audience not found in token") &&
			assert.Equal(t, []string{clientID}, audience, "Audience should contain the client ID")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err)
		assert.Equal(t, jwa.EC().String(), publicKey.KeyType().String(), "Key type should be EC")
		alg, ok := publicKey.Algorithm()
		require.True(t, ok)
		assert.Equal(t, jwa.ES256().String(), alg.String(), "Algorithm should be ES256")
	})

	t.Run("works with RSA keys", func(t *testing.T) {
		db, envConfig := newTestDbAndEnv(t)
		origKeyID := createRSA4096KeyJWK(t, db, envConfig, mockConfig)
		service := initJwtService(t, db, mockConfig, envConfig)

		loadedKeyID, ok := service.privateKey.KeyID()
		require.True(t, ok)
		assert.Equal(t, origKeyID, loadedKeyID, "Loaded key should have the same ID as the original")

		user := model.User{
			Base:  model.Base{ID: "rsauser789"},
			Email: utils.Ptr("rsaoauth@example.com"),
		}
		const clientID = "rsa-oauth-client"

		tokenString, err := service.GenerateOAuthAccessToken(user, clientID)
		require.NoError(t, err, "Failed to generate OAuth access token with key")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		claims, err := service.VerifyOAuthAccessToken(tokenString)
		require.NoError(t, err, "Failed to verify generated OAuth access token with key")

		subject, ok := claims.Subject()
		_ = assert.True(t, ok, "User ID not found in token") &&
			assert.Equal(t, user.ID, subject, "Token subject should match user ID")
		audience, ok := claims.Audience()
		_ = assert.True(t, ok, "Audience not found in token") &&
			assert.Equal(t, []string{clientID}, audience, "Audience should contain the client ID")

		publicKey, err := service.GetPublicJWK()
		require.NoError(t, err)
		assert.Equal(t, jwa.RSA().String(), publicKey.KeyType().String(), "Key type should be RSA")
		alg, ok := publicKey.Algorithm()
		require.True(t, ok)
		assert.Equal(t, jwa.RS256().String(), alg.String(), "Algorithm should be RS256")
	})
}

func TestGenerateVerifyOAuthRefreshToken(t *testing.T) {
	mockConfig := NewTestAppConfigService(&model.AppConfig{})

	t.Run("generates and verifies refresh token", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		const (
			userID       = "user123"
			clientID     = "client123"
			refreshToken = "rt-123"
		)

		tokenString, err := service.GenerateOAuthRefreshToken(userID, clientID, refreshToken)
		require.NoError(t, err, "Failed to generate refresh token")
		assert.NotEmpty(t, tokenString, "Token should not be empty")

		resUser, resClient, resRT, err := service.VerifyOAuthRefreshToken(tokenString)
		require.NoError(t, err, "Failed to verify generated token")
		assert.Equal(t, userID, resUser, "Should return correct user ID")
		assert.Equal(t, clientID, resClient, "Should return correct client ID")
		assert.Equal(t, refreshToken, resRT, "Should return correct refresh token")
	})

	t.Run("fails verification for expired token", func(t *testing.T) {
		service, _, _ := setupJwtService(t, mockConfig)

		token, err := jwt.NewBuilder().
			Subject("user789").
			Expiration(time.Now().Add(-1 * time.Hour)).
			IssuedAt(time.Now().Add(-2 * time.Hour)).
			Audience([]string{"client123"}).
			Issuer(service.envConfig.AppURL).
			Build()
		require.NoError(t, err, "Failed to build token")

		signed, err := jwt.Sign(token, jwt.WithKey(jwa.RS256(), service.privateKey))
		require.NoError(t, err, "Failed to sign token")

		_, _, _, err = service.VerifyOAuthRefreshToken(string(signed))
		require.Error(t, err, "Verification should fail with expired token")
		assert.Contains(t, err.Error(), "\"exp\" not satisfied", "Error message should indicate token verification failure")
	})

	t.Run("fails verification with invalid signature", func(t *testing.T) {
		service1, _, _ := setupJwtService(t, mockConfig)
		service2, _, _ := setupJwtService(t, mockConfig)

		tokenString, err := service1.GenerateOAuthRefreshToken("user789", "client123", "my-rt-123")
		require.NoError(t, err, "Failed to generate refresh token")

		_, _, _, err = service2.VerifyOAuthRefreshToken(tokenString)
		require.Error(t, err, "Verification should fail with invalid signature")
		assert.Contains(t, err.Error(), "verification error", "Error message should indicate token verification failure")
	})
}

func TestTokenTypeValidator(t *testing.T) {
	// Create a context for the validator function
	ctx := context.Background()

	t.Run("succeeds when token type matches expected type", func(t *testing.T) {
		// Create a token with the expected type
		token := jwt.New()
		err := token.Set(TokenTypeClaim, AccessTokenJWTType)
		require.NoError(t, err, "Failed to set token type claim")

		// Create a validator function for the expected type
		validator := TokenTypeValidator(AccessTokenJWTType)

		// Validate the token
		err = validator(ctx, token)
		assert.NoError(t, err, "Validator should accept token with matching type")
	})

	t.Run("fails when token type doesn't match expected type", func(t *testing.T) {
		// Create a token with a different type
		token := jwt.New()
		err := token.Set(TokenTypeClaim, OAuthAccessTokenJWTType)
		require.NoError(t, err, "Failed to set token type claim")

		// Create a validator function for a different expected type
		validator := TokenTypeValidator(IDTokenJWTType)

		// Validate the token
		err = validator(ctx, token)
		require.Error(t, err, "Validator should reject token with non-matching type")
		assert.Contains(t, err.Error(), "invalid token type: expected id-token, got oauth-access-token")
	})

	t.Run("fails when token type claim is missing", func(t *testing.T) {
		// Create a token without a type claim
		token := jwt.New()

		// Create a validator function
		validator := TokenTypeValidator(AccessTokenJWTType)

		// Validate the token
		err := validator(ctx, token)
		require.Error(t, err, "Validator should reject token without type claim")
		assert.Contains(t, err.Error(), "failed to get token type claim")
	})
}

func TestGetTokenType(t *testing.T) {
	mockConfig := NewTestAppConfigService(&model.AppConfig{})
	service, _, _ := setupJwtService(t, mockConfig)

	buildTokenForType := func(t *testing.T, typ string, setClaimsFn func(b *jwt.Builder)) string {
		t.Helper()

		b := jwt.NewBuilder()
		b.Subject("user123")
		if setClaimsFn != nil {
			setClaimsFn(b)
		}

		token, err := b.Build()
		require.NoError(t, err, "Failed to build token")

		err = SetTokenType(token, typ)
		require.NoError(t, err, "Failed to set token type")

		alg, _ := service.privateKey.Algorithm()
		signed, err := jwt.Sign(token, jwt.WithKey(alg, service.privateKey))
		require.NoError(t, err, "Failed to sign token")

		return string(signed)
	}

	t.Run("correctly identifies access tokens", func(t *testing.T) {
		tokenString := buildTokenForType(t, AccessTokenJWTType, nil)

		tokenType, _, err := service.GetTokenType(tokenString)
		require.NoError(t, err, "GetTokenType should not return an error")
		assert.Equal(t, AccessTokenJWTType, tokenType, "Should identify access token type")
	})

	t.Run("correctly identifies ID tokens", func(t *testing.T) {
		tokenString := buildTokenForType(t, IDTokenJWTType, nil)

		tokenType, _, err := service.GetTokenType(tokenString)
		require.NoError(t, err, "GetTokenType should not return an error")
		assert.Equal(t, IDTokenJWTType, tokenType, "Should identify ID token type")
	})

	t.Run("fails when token type claim is missing", func(t *testing.T) {
		tokenString := buildTokenForType(t, "", nil)

		_, _, err := service.GetTokenType(tokenString)
		require.Error(t, err, "GetTokenType should return an error for tokens without type claim")
		assert.Contains(t, err.Error(), "failed to get token type claim", "Error message should indicate missing token type claim")
	})
}

func importKey(t *testing.T, db *gorm.DB, envConfig *common.EnvConfigSchema, appConfig *AppConfigService, privateKeyRaw any) string {
	t.Helper()

	privateKey, err := jwkutils.ImportRawKey(privateKeyRaw, "", "")
	require.NoError(t, err, "Failed to import private key")

	return saveKeyToDatabase(t, db, envConfig, appConfig, privateKey)
}

// Because generating a RSA-406 key isn't immediate, we pre-compute one
var (
	rsaKeyPrecomputed    *rsa.PrivateKey
	rsaKeyPrecomputeOnce sync.Once
)

func createRSA4096KeyJWK(t *testing.T, db *gorm.DB, envConfig *common.EnvConfigSchema, appConfig *AppConfigService) string {
	t.Helper()

	rsaKeyPrecomputeOnce.Do(func() {
		var err error
		rsaKeyPrecomputed, err = rsa.GenerateKey(rand.Reader, 4096)
		if err != nil {
			panic("failed to precompute RSA key: " + err.Error())
		}
	})

	// Import as JWK and save it
	return importKey(t, db, envConfig, appConfig, rsaKeyPrecomputed)
}

func createECDSAKeyJWK(t *testing.T, db *gorm.DB, envConfig *common.EnvConfigSchema, appConfig *AppConfigService) string {
	t.Helper()

	// Generate a new P-256 ECDSA key
	privateKeyRaw, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	require.NoError(t, err, "Failed to generate ECDSA key")

	// Import as JWK and save it
	return importKey(t, db, envConfig, appConfig, privateKeyRaw)
}

// Helper function to create an Ed25519 key and save it as JWK
func createEdDSAKeyJWK(t *testing.T, db *gorm.DB, envConfig *common.EnvConfigSchema, appConfig *AppConfigService) string {
	t.Helper()

	// Generate a new Ed25519 key pair
	_, privateKeyRaw, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err, "Failed to generate Ed25519 key")

	// Import as JWK and save it
	return importKey(t, db, envConfig, appConfig, privateKeyRaw)
}