Files
2026-06-22 08:55:57 +08:00

107 lines
3.9 KiB
Go

package crypto
import (
"context"
"fmt"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
)
// PostgresKeyStore reads/writes the crypto_keys table that records which key
// versions exist and which is currently active (the write key). The key
// material is NOT stored here — only version metadata.
type PostgresKeyStore struct {
pool poolExec
}
// poolExec is the subset of *pgxpool.Pool used by the store. Declared as an
// interface so tests can inject a fake without a live DB.
type poolExec interface {
QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}
// poolFullExec additionally allows mutations (used by RotateActive/RetireVersion).
// Exec returns the CONCRETE pgconn.CommandTag so *pgxpool.Pool satisfies this
// interface — an interface-typed return would not match pgxpool's method
// signature, making the type assertion in RotateActive/RetireVersion silently
// fail and key rotation always error out.
type poolFullExec interface {
poolExec
Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
}
// NewPostgresKeyStore builds the store over a pgx pool (or any compatible
// querier). It satisfies KeyStore.
func NewPostgresKeyStore(pool poolExec) *PostgresKeyStore {
return &PostgresKeyStore{pool: pool}
}
// ActiveVersion returns the single crypto_keys row whose status='active'. If
// none is found it falls back to version 1 (single-key bootstrap before the
// seed row exists).
func (s *PostgresKeyStore) ActiveVersion(ctx context.Context) (int, error) {
var v int
err := s.pool.QueryRow(ctx,
`SELECT version FROM crypto_keys WHERE status = 'active' ORDER BY version DESC LIMIT 1`,
).Scan(&v)
if err != nil {
if err == pgx.ErrNoRows {
return 1, nil
}
return 0, fmt.Errorf("crypto key active version: %w", err)
}
return v, nil
}
// RotateActive bumps the active write version: it retires the current active
// row and inserts/activates the next version (current+1). The new version's key
// material must already be loadable by the provider (e.g. APP_MASTER_KEY_V<n>
// staged) before this is called. Returns the new active version.
//
// pool must support Exec (a *pgxpool.Pool does). Both rows are written in one
// transaction by the caller's pool if it is transactional; here we issue two
// statements and rely on the unique active invariant being eventually consistent
// within the rotation window (both versions remain loadable, so decrypts never
// break).
func (s *PostgresKeyStore) RotateActive(ctx context.Context, provider string, keyRef string) (int, error) {
fe, ok := s.pool.(poolFullExec)
if !ok {
return 0, fmt.Errorf("crypto key store: pool does not support writes")
}
cur, err := s.ActiveVersion(ctx)
if err != nil {
return 0, err
}
next := cur + 1
if _, err := fe.Exec(ctx,
`UPDATE crypto_keys SET status = 'retiring' WHERE status = 'active'`); err != nil {
return 0, fmt.Errorf("retire current key: %w", err)
}
if _, err := fe.Exec(ctx,
`INSERT INTO crypto_keys (version, provider, key_ref, status)
VALUES ($1, $2, NULLIF($3, ''), 'active')
ON CONFLICT (version) DO UPDATE SET status = 'active', provider = EXCLUDED.provider, key_ref = EXCLUDED.key_ref`,
next, provider, keyRef); err != nil {
return 0, fmt.Errorf("insert next key: %w", err)
}
return next, nil
}
// RetireVersion marks a fully-superseded version as retired (no longer needed
// after the re-encrypt job has moved every row off it). The provider may then
// safely drop that key material.
func (s *PostgresKeyStore) RetireVersion(ctx context.Context, version int) error {
fe, ok := s.pool.(poolFullExec)
if !ok {
return fmt.Errorf("crypto key store: pool does not support writes")
}
if _, err := fe.Exec(ctx,
`UPDATE crypto_keys SET status = 'retired' WHERE version = $1`, version); err != nil {
return fmt.Errorf("retire key v%d: %w", version, err)
}
return nil
}
var _ KeyStore = (*PostgresKeyStore)(nil)