feat: use HS256 for Pocket ID's own session tokens + fixes (#1733)

backend/internal/utils/jwk/key_provider.go

@@ -14,15 +14,28 @@ type KeyProviderOpts struct {
 	EnvConfig *common.EnvConfigSchema
 	DB        *gorm.DB
 	Kek       []byte
+	DBKey     string
 }
 
 type KeyProvider interface {
 	Init(opts KeyProviderOpts) error
 	LoadKey(ctx context.Context) (jwk.Key, error)
 	SaveKey(ctx context.Context, key jwk.Key) error
+	ReplaceKey(ctx context.Context, key jwk.Key) error
 }
 
-func GetKeyProvider(db *gorm.DB, envConfig *common.EnvConfigSchema, instanceID string) (keyProvider KeyProvider, err error) {
+// GetKeyProvider returns the provider for the key used to sign tokens that are consumed externally, such as ID tokens and access tokens for apps
+func GetKeyProvider(db *gorm.DB, envConfig *common.EnvConfigSchema, instanceID string) (KeyProvider, error) {
+	return getKeyProvider(db, envConfig, instanceID, PrivateKeyDBKey)
+}
+
+// GetSessionKeyProvider returns the provider for the symmetric key used to sign session tokens
+// This key is symmetric and separate from the one used to sign tokens for external consumption
+func GetSessionKeyProvider(db *gorm.DB, envConfig *common.EnvConfigSchema, instanceID string) (KeyProvider, error) {
+	return getKeyProvider(db, envConfig, instanceID, SessionKeyDBKey)
+}
+
+func getKeyProvider(db *gorm.DB, envConfig *common.EnvConfigSchema, instanceID string, dbKey string) (keyProvider KeyProvider, err error) {
 	// Load the encryption key (KEK) if present
 	kek, err := LoadKeyEncryptionKey(envConfig, instanceID)
 	if err != nil {
@@ -34,6 +47,7 @@ func GetKeyProvider(db *gorm.DB, envConfig *common.EnvConfigSchema, instanceID s
 		DB:        db,
 		EnvConfig: envConfig,
 		Kek:       kek,
+		DBKey:     dbKey,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to init key provider: %w", err)

backend/internal/utils/jwk/key_provider_database.go

@@ -15,11 +15,21 @@ import (
 	cryptoutils "github.com/pocket-id/pocket-id/backend/internal/utils/crypto"
 )
 
-const PrivateKeyDBKey = "jwt_private_key.json"
+const (
+	// PrivateKeyDBKey is the row in the "kv" table containing the key used to sign tokens that are consumed externally
+	PrivateKeyDBKey = "jwt_private_key.json"
+
+	// SessionKeyDBKey is the row in the "kv" table containing the symmetric key used to sign session tokens
+	SessionKeyDBKey = "session_key.json"
+)
+
+// ErrRetryKeyStorage signals that the caller should reload the key before trying to store it again
+var ErrRetryKeyStorage = errors.New("retry key storage")
 
 type KeyProviderDatabase struct {
-	db  *gorm.DB
-	kek []byte
+	db    *gorm.DB
+	kek   []byte
+	dbKey string
 }
 
 func (f *KeyProviderDatabase) Init(opts KeyProviderOpts) error {
@@ -30,12 +40,18 @@ func (f *KeyProviderDatabase) Init(opts KeyProviderOpts) error {
 	f.db = opts.DB
 	f.kek = opts.Kek
 
+	// Callers that don't ask for a specific row get the token signing key, which is the key most of the codebase deals with
+	f.dbKey = opts.DBKey
+	if f.dbKey == "" {
+		f.dbKey = PrivateKeyDBKey
+	}
+
 	return nil
 }
 
 func (f *KeyProviderDatabase) LoadKey(ctx context.Context) (key jwk.Key, err error) {
 	row := model.KV{
-		Key: PrivateKeyDBKey,
+		Key: f.dbKey,
 	}
 
 	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
@@ -45,7 +61,7 @@ func (f *KeyProviderDatabase) LoadKey(ctx context.Context) (key jwk.Key, err err
 		// Key not present in the database - return nil so a new one can be generated
 		return nil, nil
 	} else if err != nil {
-		return nil, fmt.Errorf("failed to retrieve private key from the database: %w", err)
+		return nil, fmt.Errorf("failed to retrieve key from the database: %w", err)
 	}
 
 	if row.Value == nil || *row.Value == "" {
@@ -56,42 +72,62 @@ func (f *KeyProviderDatabase) LoadKey(ctx context.Context) (key jwk.Key, err err
 	// Decode from base64
 	enc, err := base64.StdEncoding.DecodeString(*row.Value)
 	if err != nil {
-		return nil, fmt.Errorf("failed to read encrypted private key: not a valid base64-encoded value: %w", err)
+		return nil, fmt.Errorf("failed to read encrypted key: not a valid base64-encoded value: %w", err)
 	}
 
 	// Decrypt the data
 	data, err := cryptoutils.Decrypt(f.kek, enc, nil)
 	if err != nil {
-		return nil, fmt.Errorf("failed to decrypt private key: %w", err)
+		return nil, fmt.Errorf("failed to decrypt key: %w", err)
 	}
 
 	// Parse the key
 	key, err = jwk.ParseKey(data)
 	if err != nil {
-		return nil, fmt.Errorf("failed to parse encrypted private key: %w", err)
+		return nil, fmt.Errorf("failed to parse encrypted key: %w", err)
 	}
 
 	return key, nil
 }
 
 func (f *KeyProviderDatabase) SaveKey(ctx context.Context, key jwk.Key) error {
-	// Encode the key to JSON
-	data, err := EncodeJWKBytes(key)
+	// Prepare the encrypted database value before attempting the insert
+	row, err := f.prepareKeyRow(key)
 	if err != nil {
-		return fmt.Errorf("failed to encode key to JSON: %w", err)
+		return err
 	}
 
-	// Encrypt the key then encode to Base64
-	enc, err := cryptoutils.Encrypt(f.kek, data, nil)
-	if err != nil {
-		return fmt.Errorf("failed to encrypt key: %w", err)
-	}
-	// Save to database
-	row := model.KV{
-		Key:   PrivateKeyDBKey,
-		Value: new(base64.StdEncoding.EncodeToString(enc)),
+	// Insert only if the key doesn't exist yet
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	result := f.db.
+		WithContext(ctx).
+		Clauses(clause.OnConflict{
+			Columns:   []clause.Column{{Name: "key"}},
+			DoNothing: true,
+		}).
+		Create(&row)
+	if result.Error != nil {
+		// Preserve ordinary database failures because only an existing row can be resolved by reloading
+		return fmt.Errorf("failed to store key in database: %w", result.Error)
 	}
 
+	// Ask the caller to reload the winning key when another writer created the row first
+	if result.RowsAffected == 0 {
+		return ErrRetryKeyStorage
+	}
+
+	return nil
+}
+
+func (f *KeyProviderDatabase) ReplaceKey(ctx context.Context, key jwk.Key) error {
+	// Prepare the encrypted database value before attempting the replacement
+	row, err := f.prepareKeyRow(key)
+	if err != nil {
+		return err
+	}
+
+	// Upsert explicitly because key rotation must replace an existing key and can also recover a missing row
 	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 	err = f.db.
@@ -103,13 +139,31 @@ func (f *KeyProviderDatabase) SaveKey(ctx context.Context, key jwk.Key) error {
 		Create(&row).
 		Error
 	if err != nil {
-		// There's one scenario where if Pocket ID is started fresh with more than 1 replica, they both could be trying to create the private key in the database at the same time
-		// In this case, only one of the replicas will succeed; the other one(s) will return an error here, which will cascade down and cause the replica(s) to crash and be restarted (at that point they'll load the then-existing key from the database)
-		return fmt.Errorf("failed to store private key in database: %w", err)
+		return fmt.Errorf("failed to replace key in database: %w", err)
 	}
 
 	return nil
 }
 
+func (f *KeyProviderDatabase) prepareKeyRow(key jwk.Key) (model.KV, error) {
+	// Encode the key to JSON
+	data, err := EncodeJWKBytes(key)
+	if err != nil {
+		return model.KV{}, fmt.Errorf("failed to encode key to JSON: %w", err)
+	}
+
+	// Encrypt the key then encode to Base64
+	enc, err := cryptoutils.Encrypt(f.kek, data, nil)
+	if err != nil {
+		return model.KV{}, fmt.Errorf("failed to encrypt key: %w", err)
+	}
+
+	// Build the row once so inserts and explicit replacements encode keys identically
+	return model.KV{
+		Key:   f.dbKey,
+		Value: new(base64.StdEncoding.EncodeToString(enc)),
+	}, nil
+}
+
 // Compile-time interface check
 var _ KeyProvider = (*KeyProviderDatabase)(nil)

backend/internal/utils/jwk/key_provider_database_test.go

@@ -5,11 +5,13 @@ import (
 	"crypto/elliptic"
 	"crypto/rand"
 	"encoding/base64"
+	"errors"
 	"testing"
 
 	"github.com/lestrrat-go/jwx/v4/jwk"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
 
 	"github.com/pocket-id/pocket-id/backend/internal/model"
 	cryptoutils "github.com/pocket-id/pocket-id/backend/internal/utils/crypto"
@@ -257,6 +259,164 @@ func TestKeyProviderDatabase_SaveKey(t *testing.T) {
 
 		assert.Equal(t, keyBytes, parsedKeyBytes, "Expected saved key to match original key")
 	})
+
+	t.Run("SaveKey returns the database error without the retry sentinel", func(t *testing.T) {
+		db := testutils.NewDatabaseForTest(t)
+		kek := generateTestKEK(t)
+		storeErr := errors.New("test database error")
+
+		err = db.Callback().Create().Before("gorm:create").Register("fail_key_storage", func(tx *gorm.DB) {
+			_ = tx.AddError(storeErr)
+		})
+		require.NoError(t, err)
+
+		provider := &KeyProviderDatabase{}
+		err = provider.Init(KeyProviderOpts{
+			DB:  db,
+			Kek: kek,
+		})
+		require.NoError(t, err)
+
+		err = provider.SaveKey(t.Context(), key)
+		require.ErrorIs(t, err, storeErr)
+		require.NotErrorIs(t, err, ErrRetryKeyStorage)
+		require.ErrorContains(t, err, "failed to store key in database")
+	})
+
+	t.Run("SaveKey keeps the first key when the row already exists", func(t *testing.T) {
+		// Initialize a provider backed by an empty database
+		db := testutils.NewDatabaseForTest(t)
+		provider := &KeyProviderDatabase{}
+		err := provider.Init(KeyProviderOpts{
+			DB:  db,
+			Kek: generateTestKEK(t),
+		})
+		require.NoError(t, err)
+
+		// Store the key that should win the conflict
+		err = provider.SaveKey(t.Context(), key)
+		require.NoError(t, err)
+
+		// Try storing a second key and verify the caller is told to reload
+		replacementKey, err := GenerateKey("ES256", "")
+		require.NoError(t, err)
+		err = provider.SaveKey(t.Context(), replacementKey)
+		require.ErrorIs(t, err, ErrRetryKeyStorage)
+
+		// Reload the row and verify the first key was not overwritten
+		loadedKey, err := provider.LoadKey(t.Context())
+		require.NoError(t, err)
+		loadedKeyBytes, err := EncodeJWKBytes(loadedKey)
+		require.NoError(t, err)
+		keyBytes, err := EncodeJWKBytes(key)
+		require.NoError(t, err)
+		assert.Equal(t, keyBytes, loadedKeyBytes)
+	})
+}
+
+func TestKeyProviderDatabase_ReplaceKey(t *testing.T) {
+	// Initialize a provider and store the original key
+	db := testutils.NewDatabaseForTest(t)
+	provider := &KeyProviderDatabase{}
+	err := provider.Init(KeyProviderOpts{
+		DB:  db,
+		Kek: generateTestKEK(t),
+	})
+	require.NoError(t, err)
+
+	originalKey, err := GenerateKey("ES256", "")
+	require.NoError(t, err)
+	err = provider.SaveKey(t.Context(), originalKey)
+	require.NoError(t, err)
+
+	// Explicitly replace the stored key
+	replacementKey, err := GenerateKey("ES256", "")
+	require.NoError(t, err)
+	err = provider.ReplaceKey(t.Context(), replacementKey)
+	require.NoError(t, err)
+
+	// Reload the row and verify the replacement was persisted
+	loadedKey, err := provider.LoadKey(t.Context())
+	require.NoError(t, err)
+	loadedKeyBytes, err := EncodeJWKBytes(loadedKey)
+	require.NoError(t, err)
+	replacementKeyBytes, err := EncodeJWKBytes(replacementKey)
+	require.NoError(t, err)
+	assert.Equal(t, replacementKeyBytes, loadedKeyBytes)
+}
+
+func TestKeyProviderDatabase_DBKey(t *testing.T) {
+	t.Run("keys stored under different rows do not overwrite each other", func(t *testing.T) {
+		db := testutils.NewDatabaseForTest(t)
+		kek := generateTestKEK(t)
+
+		signingKeyProvider := &KeyProviderDatabase{}
+		err := signingKeyProvider.Init(KeyProviderOpts{
+			DB:    db,
+			Kek:   kek,
+			DBKey: PrivateKeyDBKey,
+		})
+		require.NoError(t, err)
+
+		sessionKeyProvider := &KeyProviderDatabase{}
+		err = sessionKeyProvider.Init(KeyProviderOpts{
+			DB:    db,
+			Kek:   kek,
+			DBKey: SessionKeyDBKey,
+		})
+		require.NoError(t, err)
+
+		// Save a different key with each provider
+		signingKey, err := GenerateKey("ES256", "")
+		require.NoError(t, err)
+		err = signingKeyProvider.SaveKey(t.Context(), signingKey)
+		require.NoError(t, err)
+
+		sessionKey, err := GenerateSessionKey()
+		require.NoError(t, err)
+		err = sessionKeyProvider.SaveKey(t.Context(), sessionKey)
+		require.NoError(t, err)
+
+		// Each provider must load back the key it saved
+		loadedSigningKey, err := signingKeyProvider.LoadKey(t.Context())
+		require.NoError(t, err)
+		require.NotNil(t, loadedSigningKey)
+		signingKid, _ := signingKey.KeyID()
+		loadedSigningKid, _ := loadedSigningKey.KeyID()
+		assert.Equal(t, signingKid, loadedSigningKid)
+
+		loadedSessionKey, err := sessionKeyProvider.LoadKey(t.Context())
+		require.NoError(t, err)
+		require.NotNil(t, loadedSessionKey)
+		sessionKid, _ := sessionKey.KeyID()
+		loadedSessionKid, _ := loadedSessionKey.KeyID()
+		assert.Equal(t, sessionKid, loadedSessionKid)
+
+		// Both rows must exist in the database
+		var count int64
+		err = db.Model(&model.KV{}).Where("key IN ?", []string{PrivateKeyDBKey, SessionKeyDBKey}).Count(&count).Error
+		require.NoError(t, err)
+		assert.Equal(t, int64(2), count)
+	})
+
+	t.Run("defaults to the token signing key row when not set", func(t *testing.T) {
+		db := testutils.NewDatabaseForTest(t)
+		kek := generateTestKEK(t)
+
+		provider := &KeyProviderDatabase{}
+		require.NoError(t, provider.Init(KeyProviderOpts{
+			DB:  db,
+			Kek: kek,
+		}))
+
+		key, err := GenerateKey("ES256", "")
+		require.NoError(t, err)
+		require.NoError(t, provider.SaveKey(t.Context(), key))
+
+		var kv model.KV
+		err = db.Where("key = ?", PrivateKeyDBKey).First(&kv).Error
+		require.NoError(t, err, "Expected the key to be stored in the token signing key row")
+	})
 }
 
 func generateTestKEK(t *testing.T) []byte {

backend/internal/utils/jwk/utils.go

@@ -7,6 +7,7 @@ import (
 	"crypto/elliptic"
 	"crypto/rand"
 	"crypto/rsa"
+	"crypto/sha256"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -25,6 +26,12 @@ const (
 	KeyUsageSigning = "sig"
 )
 
+// SessionKeyAlg returns the algorithm used to sign session tokens
+func SessionKeyAlg() jwa.SignatureAlgorithm {
+	// Session tokens are only consumed by Pocket ID itself, so use a faster symmetric key that's kept private
+	return jwa.HS256()
+}
+
 // EncodeJWK encodes a jwk.Key to a writable stream.
 func EncodeJWK(w io.Writer, key jwk.Key) error {
 	enc := json.NewEncoder(w)
@@ -124,6 +131,9 @@ func EnsureAlgInKey(key jwk.Key, alg string, crv string) {
 		// Default to EdDSA and Ed25519 for OKP keys
 		_ = key.Set(jwk.AlgorithmKey, jwa.EdDSA())
 		_ = key.Set(jwk.OKPCrvKey, jwa.Ed25519())
+	case jwa.OctetSeq():
+		// Default to HS256 for symmetric keys
+		_ = key.Set(jwk.AlgorithmKey, jwa.HS256())
 	}
 }
 
@@ -160,3 +170,16 @@ func GenerateKey(alg string, crv string) (key jwk.Key, err error) {
 	// Import the raw key
 	return ImportRawKey(rawKey, alg, crv)
 }
+
+// GenerateSessionKey generates a new symmetric key used to sign session tokens
+func GenerateSessionKey() (jwk.Key, error) {
+	// Use HS256, which is based on SHA256
+	rawKey := make([]byte, sha256.Size)
+	_, err := io.ReadFull(rand.Reader, rawKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate session key: %w", err)
+	}
+
+	// Import the raw key
+	return ImportRawKey(rawKey, SessionKeyAlg().String(), "")
+}

backend/internal/utils/jwk/utils_test.go

@@ -6,6 +6,7 @@ import (
 	"crypto/elliptic"
 	"crypto/rand"
 	"crypto/rsa"
+	"crypto/sha256"
 	"encoding/hex"
 	"testing"
 
@@ -154,6 +155,40 @@ func TestGenerateKey(t *testing.T) {
 	}
 }
 
+func TestGenerateSessionKey(t *testing.T) {
+	key, err := GenerateSessionKey()
+	require.NoError(t, err)
+	require.NotNil(t, key)
+
+	// The session key must be a symmetric key for HS256
+	assert.Equal(t, jwa.OctetSeq(), key.KeyType())
+	alg, ok := key.Algorithm()
+	_ = assert.True(t, ok, "algorithm should be set in the key") &&
+		assert.Equal(t, jwa.HS256().String(), alg.String())
+
+	// Verify other required fields are set
+	kid, ok := key.KeyID()
+	_ = assert.True(t, ok, "key ID should be set") &&
+		assert.NotEmpty(t, kid, "key ID should not be empty")
+
+	usage, ok := key.KeyUsage()
+	_ = assert.True(t, ok, "key usage should be set") &&
+		assert.Equal(t, KeyUsageSigning, usage)
+
+	// Verify the key material has the expected length
+	rawKey, err := jwk.Export[[]byte](key)
+	require.NoError(t, err)
+	assert.Len(t, rawKey, sha256.Size)
+
+	// Each invocation must return a different key
+	otherKey, err := GenerateSessionKey()
+	require.NoError(t, err)
+
+	otherRawKey, err := jwk.Export[[]byte](otherKey)
+	require.NoError(t, err)
+	assert.NotEqual(t, rawKey, otherRawKey, "each generated session key should be different")
+}
+
 func TestEnsureAlgInKey(t *testing.T) {
 	// Generate an RSA-2048 key
 	rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
@@ -282,6 +317,16 @@ func TestEnsureAlgInKey(t *testing.T) {
 				expectedAlg: jwa.EdDSA(),
 				expectedCrv: jwa.Ed25519().String(),
 			},
+			{
+				name: "Symmetric key defaults to HS256",
+				keyGen: func() (any, error) {
+					rawKey := make([]byte, sha256.Size)
+					_, err := rand.Read(rawKey)
+					return rawKey, err
+				},
+				expectedAlg: jwa.HS256(),
+				expectedCrv: "",
+			},
 		}
 
 		for _, tt := range tests {