agentic-coding-workflow/internal/jobs/handlers/secret_reencrypt_pg.go

package handlers

import (
	"context"
	"fmt"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/yan1h/agent-coding-workflow/internal/infra/crypto"
)

// PostgresReencryptStore re-encrypts the platform's encrypted-column tables
// using one transaction per row so the job is idempotent and resumable. Each
// table's blob column(s) are decrypted under the row's current key_version then
// re-sealed under the active version, and the row's key_version is advanced in
// the same tx. Rows with NULL blobs are skipped by the WHERE predicate.
type PostgresReencryptStore struct {
	pool *pgxpool.Pool
}

// NewPostgresReencryptStore builds the store over a pgx pool.
func NewPostgresReencryptStore(pool *pgxpool.Pool) *PostgresReencryptStore {
	return &PostgresReencryptStore{pool: pool}
}

var _ ReencryptStore = (*PostgresReencryptStore)(nil)

// tableSpec describes how to scan + update one encrypted-column table.
type tableSpec struct {
	name   string
	blobs  []string // encrypted column names (1 or 2)
	notNul string   // SQL predicate for "row carries ciphertext"
}

// reencryptSpecs is the canonical list of encrypted-column owners. Order is
// stable so the status endpoint and the handler iterate deterministically.
var reencryptSpecs = []tableSpec{
	{name: "acp_agent_kinds", blobs: []string{"encrypted_env", "encrypted_mcp_servers"},
		notNul: "(encrypted_env IS NOT NULL OR encrypted_mcp_servers IS NOT NULL)"},
	{name: "acp_agent_kind_config_files", blobs: []string{"encrypted_content"},
		notNul: "encrypted_content IS NOT NULL"},
	{name: "workspace_run_profiles", blobs: []string{"encrypted_env"},
		notNul: "encrypted_env IS NOT NULL"},
	{name: "git_credentials", blobs: []string{"encrypted_secret"},
		notNul: "encrypted_secret IS NOT NULL"},
	{name: "llm_endpoints", blobs: []string{"api_key_encrypted"},
		notNul: "api_key_encrypted IS NOT NULL"},
}

func specByName(name string) (tableSpec, bool) {
	for _, s := range reencryptSpecs {
		if s.name == name {
			return s, true
		}
	}
	return tableSpec{}, false
}

// Tables lists the re-encryptable table names.
func (s *PostgresReencryptStore) Tables() []string {
	out := make([]string, 0, len(reencryptSpecs))
	for _, sp := range reencryptSpecs {
		out = append(out, sp.name)
	}
	return out
}

// StaleCount reports how many rows of table remain below activeVersion.
func (s *PostgresReencryptStore) StaleCount(ctx context.Context, table string, active int) (int, error) {
	spec, ok := specByName(table)
	if !ok {
		return 0, fmt.Errorf("unknown reencrypt table %q", table)
	}
	q := fmt.Sprintf(
		`SELECT count(*) FROM %s WHERE key_version < $1 AND %s`, spec.name, spec.notNul)
	var n int
	if err := s.pool.QueryRow(ctx, q, active).Scan(&n); err != nil {
		return 0, fmt.Errorf("stale count %s: %w", table, err)
	}
	return n, nil
}

// ReencryptTable advances up to batch rows below active onto the active version.
// It selects a batch of ids FOR UPDATE SKIP LOCKED, then re-seals each row in its
// own transaction so a crash leaves a clean partial state.
func (s *PostgresReencryptStore) ReencryptTable(ctx context.Context, enc *crypto.KeyedEncryptor, table string, active, batch int) (int, error) {
	spec, ok := specByName(table)
	if !ok {
		return 0, fmt.Errorf("unknown reencrypt table %q", table)
	}

	// Pull a batch of candidate ids (cheap, uses the partial key_version index).
	idQ := fmt.Sprintf(
		`SELECT id FROM %s WHERE key_version < $1 AND %s ORDER BY id LIMIT $2`, spec.name, spec.notNul)
	rows, err := s.pool.Query(ctx, idQ, active, batch)
	if err != nil {
		return 0, fmt.Errorf("select ids %s: %w", table, err)
	}
	var ids []uuid.UUID
	for rows.Next() {
		var id uuid.UUID
		if err := rows.Scan(&id); err != nil {
			rows.Close()
			return 0, fmt.Errorf("scan id %s: %w", table, err)
		}
		ids = append(ids, id)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return 0, err
	}

	done := 0
	for _, id := range ids {
		if err := s.reencryptRow(ctx, enc, spec, active, id); err != nil {
			return done, err
		}
		done++
	}
	return done, nil
}

// reencryptRow re-seals one row's blob(s) in a single transaction.
func (s *PostgresReencryptStore) reencryptRow(ctx context.Context, enc *crypto.KeyedEncryptor, spec tableSpec, active int, id uuid.UUID) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()

	// Lock the row and read its current version + blobs.
	cols := ""
	for _, c := range spec.blobs {
		cols += c + ", "
	}
	selQ := fmt.Sprintf(`SELECT %skey_version FROM %s WHERE id = $1 FOR UPDATE`, cols, spec.name)

	dest := make([]any, 0, len(spec.blobs)+1)
	blobVals := make([]*[]byte, len(spec.blobs))
	for i := range spec.blobs {
		blobVals[i] = new([]byte)
		dest = append(dest, blobVals[i])
	}
	var curVer int
	dest = append(dest, &curVer)

	if err := tx.QueryRow(ctx, selQ, id).Scan(dest...); err != nil {
		if err == pgx.ErrNoRows {
			return nil // row vanished; nothing to do
		}
		return fmt.Errorf("lock row %s/%s: %w", spec.name, id, err)
	}
	if curVer >= active {
		return tx.Commit(ctx) // already advanced by a concurrent run; idempotent
	}

	// Re-seal each non-nil blob: decrypt under curVer, encrypt under active.
	setExpr := ""
	args := []any{id}
	argN := 2
	for i, c := range spec.blobs {
		b := *blobVals[i]
		if len(b) == 0 {
			continue // NULL/empty blob: leave as-is
		}
		plain, derr := enc.DecryptVersion(ctx, b, curVer)
		if derr != nil {
			return fmt.Errorf("decrypt %s/%s.%s v%d: %w", spec.name, id, c, curVer, derr)
		}
		sealed, err := enc.SealWithVersion(ctx, active, plain)
		if err != nil {
			return fmt.Errorf("reseal %s/%s.%s: %w", spec.name, id, c, err)
		}
		if setExpr != "" {
			setExpr += ", "
		}
		setExpr += fmt.Sprintf("%s = $%d", c, argN)
		args = append(args, sealed)
		argN++
	}
	// Always advance key_version (covers the all-NULL-blob defensive case too).
	if setExpr != "" {
		setExpr += ", "
	}
	setExpr += fmt.Sprintf("key_version = $%d", argN)
	args = append(args, active)

	updQ := fmt.Sprintf(`UPDATE %s SET %s WHERE id = $1`, spec.name, setExpr)
	if _, err := tx.Exec(ctx, updQ, args...); err != nil {
		return fmt.Errorf("update %s/%s: %w", spec.name, id, err)
	}
	return tx.Commit(ctx)
}