bugfix

2026-06-22 08:55:57 +08:00
parent dbb87823e8
commit 6ade6e8fa9
325 changed files with 41131 additions and 855 deletions
@@ -0,0 +1,195 @@
+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)
+}