From c4264dc4455c3f3f96e22eeb4dc4dc0ba045771c Mon Sep 17 00:00:00 2001 From: Aikins Laryea Date: Sun, 13 Sep 2026 10:54:50 +0000 Subject: [PATCH 1/5] store: keep PostgreSQL auth saves consistent Stage credentials until PostgreSQL accepts the save, then publish the local file. Serialize participating writers through publication and compensation with a session advisory lock based on relation identity. Require direct or session-pooled connections and coordinated writer upgrades. Cover filesystem failures, unchanged-local retries, and concurrent writes with unit and PostgreSQL integration tests. --- .env.example | 3 + internal/store/postgresstore.go | 286 ++++++- .../store/postgresstore_integration_test.go | 157 ++++ internal/store/postgresstore_test.go | 726 ++++++++++++++++++ 4 files changed, 1131 insertions(+), 41 deletions(-) create mode 100644 internal/store/postgresstore_integration_test.go create mode 100644 internal/store/postgresstore_test.go diff --git a/.env.example b/.env.example index 5b0546f4c59..fecad9abf17 100644 --- a/.env.example +++ b/.env.example @@ -12,6 +12,9 @@ # ------------------------------------------------------------------------------ # Postgres Token Store (optional) # ------------------------------------------------------------------------------ +# Use a direct PostgreSQL connection or session pooling; transaction pooling is unsupported. +# Stop older instances before upgrading writers sharing an auth table: save compensation +# requires every writer to participate in the same session-level advisory locking protocol. # PGSTORE_DSN=postgresql://user:pass@localhost:5432/cliproxy # PGSTORE_SCHEMA=public # PGSTORE_LOCAL_PATH=/var/lib/cliproxy diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index aeaac3e1fca..2b9690eb2f8 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -3,6 +3,7 @@ package store import ( "context" "database/sql" + "database/sql/driver" "encoding/json" "errors" "fmt" @@ -46,6 +47,7 @@ type PostgresStore struct { authDir string cooldownStore *postgresCooldownStateStore mu sync.Mutex + renameFile func(string, string) error } // NewPostgresStore establishes a connection to PostgreSQL and prepares the local workspace. @@ -240,6 +242,25 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri return "", fmt.Errorf("postgres store: create auth directory: %w", err) } + localPrevious, errReadPrevious := os.ReadFile(path) + localExists := errReadPrevious == nil + if errReadPrevious != nil && !errors.Is(errReadPrevious, fs.ErrNotExist) { + return "", fmt.Errorf("postgres store: read existing metadata: %w", errReadPrevious) + } + relID, err := s.relativeAuthID(path) + if err != nil { + return "", err + } + tmp := path + ".tmp" + if errRemove := os.Remove(tmp); errRemove != nil && !errors.Is(errRemove, fs.ErrNotExist) { + return "", fmt.Errorf("postgres store: remove stale temp auth file: %w", errRemove) + } + defer func() { + if errRemove := os.Remove(tmp); errRemove != nil && !errors.Is(errRemove, fs.ErrNotExist) { + log.WithError(errRemove).Warn("postgres store: remove temporary auth file") + } + }() + switch { case auth.Storage != nil: if auth.Metadata == nil { @@ -249,7 +270,7 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { setter.SetMetadata(auth.Metadata) } - if err = auth.Storage.SaveTokenToFile(path); err != nil { + if err = auth.Storage.SaveTokenToFile(tmp); err != nil { return "", err } case auth.Metadata != nil: @@ -258,42 +279,77 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri if errMarshal != nil { return "", fmt.Errorf("postgres store: marshal metadata: %w", errMarshal) } - if existing, errRead := os.ReadFile(path); errRead == nil { - if jsonEqual(existing, raw) { - return path, nil - } - } else if errRead != nil && !errors.Is(errRead, fs.ErrNotExist) { - return "", fmt.Errorf("postgres store: read existing metadata: %w", errRead) - } - tmp := path + ".tmp" if errWrite := os.WriteFile(tmp, raw, 0o600); errWrite != nil { return "", fmt.Errorf("postgres store: write temp auth file: %w", errWrite) } - if errRename := os.Rename(tmp, path); errRename != nil { - return "", fmt.Errorf("postgres store: rename auth file: %w", errRename) - } default: return "", fmt.Errorf("postgres store: nothing to persist for %s", auth.ID) } + if errChmod := os.Chmod(tmp, 0o600); errChmod != nil && !errors.Is(errChmod, fs.ErrNotExist) { + return "", fmt.Errorf("postgres store: secure temp auth file: %w", errChmod) + } + + candidate, errReadCandidate := os.ReadFile(tmp) + if errReadCandidate != nil { + return "", fmt.Errorf("postgres store: read temp auth file: %w", errReadCandidate) + } + + err = s.withAuthLock(ctx, relID, func(conn *sql.Conn) error { + var ( + durablePrevious []byte + durablePreviousExists bool + durableChanged bool + ) + if len(candidate) == 0 { + durablePrevious, durablePreviousExists, err = s.deleteAuthRecordReturning(ctx, conn, relID) + durableChanged = durablePreviousExists + } else { + durablePrevious, durablePreviousExists, err = s.replaceAuthRecord(ctx, conn, relID, candidate) + durableChanged = true + } + if err != nil { + return err + } + if localExists && jsonEqual(localPrevious, candidate) { + return nil + } + if errRename := s.renameAuthFile(tmp, path); errRename != nil { + if durableChanged { + errRollback := s.rollbackAuthRecord(context.WithoutCancel(ctx), conn, relID, candidate, durablePrevious, durablePreviousExists) + if errRollback != nil { + return errors.Join( + fmt.Errorf("postgres store: publish auth file: %w", errRename), + fmt.Errorf("postgres store: database rollback failed: %w", errRollback), + ) + } + } + return fmt.Errorf("postgres store: publish auth file: %w", errRename) + } + return nil + }) + if err != nil { + return "", err + } + normalizeSavedPostgresAuth(auth, path) + return path, nil +} + +func (s *PostgresStore) renameAuthFile(oldPath, newPath string) error { + if s.renameFile != nil { + return s.renameFile(oldPath, newPath) + } + return os.Rename(oldPath, newPath) +} +func normalizeSavedPostgresAuth(auth *cliproxyauth.Auth, path string) { if auth.Attributes == nil { auth.Attributes = make(map[string]string) } auth.Attributes[cliproxyauth.AttributePath] = path auth.Attributes[cliproxyauth.AttributeSourceBackend] = cliproxyauth.AuthSourcePostgres - if strings.TrimSpace(auth.FileName) == "" { auth.FileName = auth.ID } - - relID, err := s.relativeAuthID(path) - if err != nil { - return "", err - } - if err = s.upsertAuthRecord(ctx, relID, path); err != nil { - return "", err - } - return path, nil } // List enumerates all auth records stored in PostgreSQL. @@ -539,35 +595,183 @@ func (s *PostgresStore) syncAuthFile(ctx context.Context, relID, path string) er return s.persistAuth(ctx, relID, data) } -func (s *PostgresStore) upsertAuthRecord(ctx context.Context, relID, path string) error { - data, err := os.ReadFile(path) +func (s *PostgresStore) persistAuth(ctx context.Context, relID string, data []byte) error { + return s.withAuthLock(ctx, relID, func(conn *sql.Conn) error { + jsonPayload := json.RawMessage(data) + query := fmt.Sprintf(` + INSERT INTO %s (id, content, created_at, updated_at) + VALUES ($1, $2, NOW(), NOW()) + ON CONFLICT (id) + DO UPDATE SET content = EXCLUDED.content, updated_at = NOW() + `, s.fullTableName(s.cfg.AuthTable)) + if _, err := conn.ExecContext(ctx, query, relID, jsonPayload); err != nil { + return fmt.Errorf("postgres store: upsert auth record: %w", err) + } + return nil + }) +} + +func (s *PostgresStore) deleteAuthRecord(ctx context.Context, relID string) error { + return s.withAuthLock(ctx, relID, func(conn *sql.Conn) error { + query := fmt.Sprintf("DELETE FROM %s WHERE id = $1", s.fullTableName(s.cfg.AuthTable)) + if _, err := conn.ExecContext(ctx, query, relID); err != nil { + return fmt.Errorf("postgres store: delete auth record: %w", err) + } + return nil + }) +} + +func (s *PostgresStore) withAuthLock(ctx context.Context, relID string, save func(*sql.Conn) error) (err error) { + conn, err := s.db.Conn(ctx) if err != nil { - return fmt.Errorf("postgres store: read auth file: %w", err) + return fmt.Errorf("postgres store: acquire auth connection: %w", err) } - if len(data) == 0 { - return s.deleteAuthRecord(ctx, relID) + defer func() { + if errClose := conn.Close(); errClose != nil && !errors.Is(errClose, sql.ErrConnDone) { + err = errors.Join(err, fmt.Errorf("postgres store: close auth connection: %w", errClose)) + } + }() + var key int64 + if errKey := conn.QueryRowContext(ctx, "SELECT hashtextextended($1::regclass::oid::text || ':' || $2::text, 0)", s.fullTableName(s.cfg.AuthTable), relID).Scan(&key); errKey != nil { + return fmt.Errorf("postgres store: resolve auth lock: %w", errKey) } - return s.persistAuth(ctx, relID, data) + locked := false + // The same session must own the lock through commit, publication, and compensation. + defer func() { + var unlocked bool + errUnlock := conn.QueryRowContext(context.WithoutCancel(ctx), "SELECT pg_advisory_unlock($1)", key).Scan(&unlocked) + if errUnlock == nil && locked && !unlocked { + errUnlock = errors.New("auth lock ownership lost") + } + if errUnlock != nil { + err = errors.Join(err, fmt.Errorf("postgres store: unlock auth record: %w", errUnlock)) + _ = conn.Raw(func(any) error { return driver.ErrBadConn }) + } + }() + if _, errLock := conn.ExecContext(ctx, "SELECT pg_advisory_lock($1)", key); errLock != nil { + return fmt.Errorf("postgres store: lock auth record: %w", errLock) + } + locked = true + return save(conn) } -func (s *PostgresStore) persistAuth(ctx context.Context, relID string, data []byte) error { - jsonPayload := json.RawMessage(data) - query := fmt.Sprintf(` +func (s *PostgresStore) replaceAuthRecord(ctx context.Context, conn *sql.Conn, relID string, candidate []byte) (previous []byte, previousExists bool, err error) { + // READ COMMITTED lets the retry observe a row that won ON CONFLICT. + tx, errBegin := conn.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelReadCommitted}) + if errBegin != nil { + return nil, false, fmt.Errorf("postgres store: begin auth replacement: %w", errBegin) + } + defer func() { + if err == nil { + return + } + if errRollback := tx.Rollback(); errRollback != nil && !errors.Is(errRollback, sql.ErrTxDone) { + err = errors.Join(err, fmt.Errorf("postgres store: rollback auth replacement: %w", errRollback)) + } + }() + + table := s.fullTableName(s.cfg.AuthTable) + selectQuery := fmt.Sprintf("SELECT content FROM %s WHERE id = $1 FOR UPDATE", table) + updateQuery := fmt.Sprintf("UPDATE %s SET content = $2, updated_at = NOW() WHERE id = $1", table) + insertQuery := fmt.Sprintf(` INSERT INTO %s (id, content, created_at, updated_at) VALUES ($1, $2, NOW(), NOW()) - ON CONFLICT (id) - DO UPDATE SET content = EXCLUDED.content, updated_at = NOW() - `, s.fullTableName(s.cfg.AuthTable)) - if _, err := s.db.ExecContext(ctx, query, relID, jsonPayload); err != nil { - return fmt.Errorf("postgres store: upsert auth record: %w", err) + ON CONFLICT (id) DO NOTHING + `, table) + for { + var content string + errScan := tx.QueryRowContext(ctx, selectQuery, relID).Scan(&content) + switch { + case errScan == nil: + result, errUpdate := tx.ExecContext(ctx, updateQuery, relID, json.RawMessage(candidate)) + if errUpdate != nil { + err = fmt.Errorf("postgres store: replace auth record: %w", errUpdate) + return nil, false, err + } + rows, errRows := result.RowsAffected() + if errRows != nil { + err = fmt.Errorf("postgres store: inspect auth replacement: %w", errRows) + return nil, false, err + } + if rows != 1 { + err = fmt.Errorf("postgres store: locked auth record disappeared before replacement") + return nil, false, err + } + if errCommit := tx.Commit(); errCommit != nil { + err = fmt.Errorf("postgres store: commit auth replacement: %w", errCommit) + return nil, false, err + } + return []byte(content), true, nil + case !errors.Is(errScan, sql.ErrNoRows): + err = fmt.Errorf("postgres store: lock auth record: %w", errScan) + return nil, false, err + } + + result, errInsert := tx.ExecContext(ctx, insertQuery, relID, json.RawMessage(candidate)) + if errInsert != nil { + err = fmt.Errorf("postgres store: insert auth record: %w", errInsert) + return nil, false, err + } + rows, errRows := result.RowsAffected() + if errRows != nil { + err = fmt.Errorf("postgres store: inspect auth insert: %w", errRows) + return nil, false, err + } + if rows == 0 { + continue + } + if rows != 1 { + err = fmt.Errorf("postgres store: inserted %d auth records, want 1", rows) + return nil, false, err + } + if errCommit := tx.Commit(); errCommit != nil { + err = fmt.Errorf("postgres store: commit auth insert: %w", errCommit) + return nil, false, err + } + return nil, false, nil } - return nil } -func (s *PostgresStore) deleteAuthRecord(ctx context.Context, relID string) error { - query := fmt.Sprintf("DELETE FROM %s WHERE id = $1", s.fullTableName(s.cfg.AuthTable)) - if _, err := s.db.ExecContext(ctx, query, relID); err != nil { - return fmt.Errorf("postgres store: delete auth record: %w", err) +func (s *PostgresStore) deleteAuthRecordReturning(ctx context.Context, conn *sql.Conn, relID string) ([]byte, bool, error) { + query := fmt.Sprintf("DELETE FROM %s WHERE id = $1 RETURNING content", s.fullTableName(s.cfg.AuthTable)) + var content string + if err := conn.QueryRowContext(ctx, query, relID).Scan(&content); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return nil, false, nil + } + return nil, false, fmt.Errorf("postgres store: delete auth record: %w", err) + } + return []byte(content), true, nil +} + +func (s *PostgresStore) rollbackAuthRecord(ctx context.Context, conn *sql.Conn, relID string, candidate, previous []byte, previousExists bool) error { + var ( + result sql.Result + err error + ) + if previousExists && len(candidate) == 0 { + query := fmt.Sprintf(` + INSERT INTO %s (id, content, created_at, updated_at) + VALUES ($1, $2, NOW(), NOW()) + ON CONFLICT (id) DO NOTHING + `, s.fullTableName(s.cfg.AuthTable)) + result, err = conn.ExecContext(ctx, query, relID, json.RawMessage(previous)) + } else if previousExists { + query := fmt.Sprintf("UPDATE %s SET content = $2, updated_at = NOW() WHERE id = $1 AND content = $3", s.fullTableName(s.cfg.AuthTable)) + result, err = conn.ExecContext(ctx, query, relID, json.RawMessage(previous), json.RawMessage(candidate)) + } else { + query := fmt.Sprintf("DELETE FROM %s WHERE id = $1 AND content = $2", s.fullTableName(s.cfg.AuthTable)) + result, err = conn.ExecContext(ctx, query, relID, json.RawMessage(candidate)) + } + if err != nil { + return fmt.Errorf("postgres store: rollback auth record: %w", err) + } + rows, errRows := result.RowsAffected() + if errRows != nil { + return fmt.Errorf("postgres store: inspect auth rollback: %w", errRows) + } + if rows != 1 { + return fmt.Errorf("postgres store: auth record changed before rollback") } return nil } diff --git a/internal/store/postgresstore_integration_test.go b/internal/store/postgresstore_integration_test.go new file mode 100644 index 00000000000..2df1fc57ba8 --- /dev/null +++ b/internal/store/postgresstore_integration_test.go @@ -0,0 +1,157 @@ +package store + +import ( + "context" + "database/sql" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { + dsn := os.Getenv("TEST_POSTGRES_DSN") + if dsn == "" { + t.Skip("TEST_POSTGRES_DSN must point to a disposable PostgreSQL database") + } + for _, test := range []struct { + name string + previous string + candidate string + newer string + operation string + }{ + {"same-content-update", `{"value":"old"}`, `{"value":"new"}`, `{"value":"new"}`, "save"}, + {"same-content-insert", "", `{"value":"new"}`, `{"value":"new"}`, "save"}, + {"jsonb-equivalent", `{"value":"old"}`, `{"a":1,"b":2}`, `{"b":2, "a":1}`, "save"}, + {"empty-save", `{"value":"old"}`, "", "", "save"}, + {"delete", `{"value":"old"}`, "", "", "delete"}, + {"watcher-update", `{"value":"old"}`, `{"value":"new"}`, `{"value":"new"}`, "watcher"}, + {"watcher-delete", `{"value":"old"}`, "", "", "watcher"}, + {"qualified-table", `{"value":"old"}`, `{"value":"new"}`, `{"value":"new"}`, "save"}, + {"cancelled-publication", `{"value":"old"}`, `{"value":"new"}`, `{"value":"new"}`, "save"}, + } { + t.Run(test.name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + db, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + defer func() { _ = db.Close() }() + table := fmt.Sprintf("auth_test_%d", time.Now().UnixNano()) + if _, err = db.ExecContext(ctx, "CREATE TABLE "+table+" (id TEXT PRIMARY KEY, content JSONB NOT NULL, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW())"); err != nil { + t.Fatal(err) + } + defer func() { + if _, errDrop := db.ExecContext(context.Background(), "DROP TABLE "+table); errDrop != nil { + t.Error(errDrop) + } + }() + newStore := func() *PostgresStore { + t.Helper() + conn, errOpen := sql.Open("pgx", dsn) + if errOpen != nil { + t.Fatal(errOpen) + } + conn.SetMaxOpenConns(1) + t.Cleanup(func() { _ = conn.Close() }) + return &PostgresStore{db: conn, cfg: PostgresStoreConfig{AuthTable: table}, authDir: t.TempDir()} + } + a, b := newStore(), newStore() + if test.name == "qualified-table" { + if err = db.QueryRowContext(ctx, "SELECT current_schema()").Scan(&b.cfg.Schema); err != nil { + t.Fatal(err) + } + } + saveCtx, cancelSave := context.WithCancel(ctx) + defer cancelSave() + const id = "credential.json" + if test.previous != "" { + if _, err = db.ExecContext(ctx, "INSERT INTO "+table+" (id, content) VALUES ($1, $2)", id, test.previous); err != nil { + t.Fatal(err) + } + for _, store := range []*PostgresStore{a, b} { + if err = os.WriteFile(filepath.Join(store.authDir, id), []byte(test.previous), 0o600); err != nil { + t.Fatal(err) + } + } + } + var pid int + if err = b.db.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&pid); err != nil { + t.Fatal(err) + } + completed := make(chan error, 1) + a.renameFile = func(string, string) error { + go func() { + var errWrite error + switch test.operation { + case "save": + _, errWrite = b.Save(ctx, &cliproxyauth.Auth{ID: id, Storage: &postgresAuthTestStorage{data: []byte(test.newer)}}) + case "delete": + errWrite = b.Delete(ctx, id) + case "watcher": + path := filepath.Join(b.authDir, id) + if errWrite = os.WriteFile(path, []byte(test.newer), 0o600); errWrite == nil { + errWrite = b.PersistAuthFiles(ctx, "", path) + } + } + completed <- errWrite + }() + for { + var waiting bool + if errWait := db.QueryRowContext(ctx, "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE pid = $1 AND locktype = 'advisory' AND NOT granted)", pid).Scan(&waiting); errWait != nil { + t.Error(errWait) + return errors.New("publish rejected") + } + if waiting { + break + } + select { + case errEarly := <-completed: + t.Errorf("concurrent writer completed before compensation: %v", errEarly) + completed <- errEarly + return errors.New("publish rejected") + default: + } + } + if test.name == "cancelled-publication" { + cancelSave() + } + return errors.New("publish rejected") + } + _, err = a.Save(saveCtx, &cliproxyauth.Auth{ID: id, Storage: &postgresAuthTestStorage{data: []byte(test.candidate)}}) + if err == nil || !strings.Contains(err.Error(), "publish rejected") || strings.Contains(err.Error(), "rollback failed") { + t.Fatalf("Save() error = %v", err) + } + select { + case err = <-completed: + if err != nil { + t.Fatal(err) + } + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + var durable string + err = db.QueryRowContext(ctx, "SELECT content FROM "+table+" WHERE id = $1", id).Scan(&durable) + if test.newer == "" { + if !errors.Is(err, sql.ErrNoRows) { + t.Fatalf("record = %q, error = %v, want missing", durable, err) + } + } else if err != nil || !jsonEqual([]byte(durable), []byte(test.newer)) { + t.Fatalf("record = %q, error = %v, want %s", durable, err, test.newer) + } + var previous []byte + if test.previous != "" { + previous = []byte(test.previous) + } + assertPostgresAuthLocal(t, filepath.Join(a.authDir, id), previous) + assertPostgresAuthNoTemp(t, filepath.Join(a.authDir, id)) + }) + } +} diff --git a/internal/store/postgresstore_test.go b/internal/store/postgresstore_test.go new file mode 100644 index 00000000000..80b04f827f3 --- /dev/null +++ b/internal/store/postgresstore_test.go @@ -0,0 +1,726 @@ +package store + +import ( + "bytes" + "context" + "database/sql" + "database/sql/driver" + "errors" + "io" + "io/fs" + "os" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/empty" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +type postgresAuthTestCall struct { + query string + args []driver.NamedValue +} + +type postgresAuthTestBackend struct { + mu sync.Mutex + calls []postgresAuthTestCall + failAt int + inspect func(postgresAuthTestCall) + content []byte + hasContent bool + unlockMissing bool + unlockErr error +} + +func (b *postgresAuthTestBackend) call(query string, args []driver.NamedValue) error { + if strings.Contains(query, "pg_advisory_") { + return nil + } + call := postgresAuthTestCall{query: query, args: clonePostgresAuthTestArgs(args)} + b.mu.Lock() + b.calls = append(b.calls, call) + callIndex := len(b.calls) + inspect := b.inspect + fail := b.failAt == callIndex + b.mu.Unlock() + if inspect != nil { + inspect(call) + } + if fail { + return errors.New("database rejected operation") + } + return nil +} + +func (b *postgresAuthTestBackend) exec(query string, args []driver.NamedValue) (driver.Result, error) { + if err := b.call(query, args); err != nil { + return nil, err + } + b.mu.Lock() + defer b.mu.Unlock() + return postgresAuthTestExec(query, args, &b.content, &b.hasContent) +} + +func (b *postgresAuthTestBackend) query(query string, args []driver.NamedValue) (driver.Rows, error) { + if strings.Contains(query, "hashtextextended") { + return &postgresAuthTestRows{columns: []string{"key"}, values: [][]driver.Value{{int64(1)}}}, nil + } + if strings.Contains(query, "pg_advisory_unlock") { + if b.unlockErr != nil { + return nil, b.unlockErr + } + return &postgresAuthTestRows{columns: []string{"unlocked"}, values: [][]driver.Value{{!b.unlockMissing}}}, nil + } + if err := b.call(query, args); err != nil { + return nil, err + } + b.mu.Lock() + defer b.mu.Unlock() + return postgresAuthTestQuery(query, &b.content, &b.hasContent) +} + +func (b *postgresAuthTestBackend) setContent(data []byte) { + b.mu.Lock() + b.content = append([]byte(nil), data...) + b.hasContent = true + b.mu.Unlock() +} + +func (b *postgresAuthTestBackend) durableSnapshot() ([]byte, bool) { + b.mu.Lock() + defer b.mu.Unlock() + return append([]byte(nil), b.content...), b.hasContent +} + +func (b *postgresAuthTestBackend) snapshotCalls() []postgresAuthTestCall { + b.mu.Lock() + defer b.mu.Unlock() + calls := make([]postgresAuthTestCall, len(b.calls)) + copy(calls, b.calls) + return calls +} + +func postgresAuthTestExec(query string, args []driver.NamedValue, content *[]byte, hasContent *bool) (driver.Result, error) { + trimmed := strings.TrimSpace(query) + switch { + case strings.HasPrefix(trimmed, "INSERT INTO"): + if strings.Contains(query, "DO NOTHING") && *hasContent { + return driver.RowsAffected(0), nil + } + *content = postgresAuthTestValue(args, 1) + *hasContent = true + return driver.RowsAffected(1), nil + case strings.HasPrefix(trimmed, "UPDATE"): + if !*hasContent { + return driver.RowsAffected(0), nil + } + if len(args) > 2 && !bytes.Equal(*content, postgresAuthTestValue(args, 2)) { + return driver.RowsAffected(0), nil + } + *content = postgresAuthTestValue(args, 1) + return driver.RowsAffected(1), nil + case strings.HasPrefix(trimmed, "DELETE") && len(args) > 1: + if !*hasContent || !bytes.Equal(*content, postgresAuthTestValue(args, 1)) { + return driver.RowsAffected(0), nil + } + *content = nil + *hasContent = false + return driver.RowsAffected(1), nil + case strings.HasPrefix(trimmed, "DELETE"): + *content = nil + *hasContent = false + return driver.RowsAffected(1), nil + default: + return driver.RowsAffected(1), nil + } +} + +func postgresAuthTestQuery(query string, content *[]byte, hasContent *bool) (driver.Rows, error) { + trimmed := strings.TrimSpace(query) + rows := &postgresAuthTestRows{columns: []string{"content"}} + switch { + case strings.HasPrefix(trimmed, "SELECT") && strings.Contains(query, "FOR UPDATE"): + if *hasContent { + rows.values = [][]driver.Value{{string(*content)}} + } + case strings.HasPrefix(trimmed, "DELETE") && strings.Contains(query, "RETURNING content"): + if *hasContent { + rows.values = [][]driver.Value{{string(*content)}} + *content = nil + *hasContent = false + } + default: + return nil, errors.New("query unsupported") + } + return rows, nil +} + +type postgresAuthTestRows struct { + columns []string + values [][]driver.Value + index int +} + +func (r *postgresAuthTestRows) Columns() []string { return r.columns } + +func (*postgresAuthTestRows) Close() error { return nil } + +func (r *postgresAuthTestRows) Next(dest []driver.Value) error { + if r.index >= len(r.values) { + return io.EOF + } + copy(dest, r.values[r.index]) + r.index++ + return nil +} + +type postgresAuthTestConnector struct { + backend *postgresAuthTestBackend +} + +func (c *postgresAuthTestConnector) Connect(context.Context) (driver.Conn, error) { + return &postgresAuthTestConn{backend: c.backend}, nil +} + +func (*postgresAuthTestConnector) Driver() driver.Driver { return postgresAuthTestDriver{} } + +type postgresAuthTestDriver struct{} + +func (postgresAuthTestDriver) Open(string) (driver.Conn, error) { + return nil, errors.New("use connector") +} + +type postgresAuthTestConn struct { + backend *postgresAuthTestBackend + tx *postgresAuthTestTx +} + +func (*postgresAuthTestConn) Prepare(string) (driver.Stmt, error) { + return nil, errors.New("prepare unsupported") +} + +func (*postgresAuthTestConn) Close() error { return nil } + +func (c *postgresAuthTestConn) Begin() (driver.Tx, error) { + return c.beginTx() +} + +func (c *postgresAuthTestConn) ExecContext(_ context.Context, query string, args []driver.NamedValue) (driver.Result, error) { + if c.tx != nil { + return c.tx.exec(query, args) + } + return c.backend.exec(query, args) +} + +func (c *postgresAuthTestConn) QueryContext(_ context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { + if c.tx != nil { + return c.tx.query(query, args) + } + return c.backend.query(query, args) +} + +func (c *postgresAuthTestConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) { + return c.beginTx() +} + +func (c *postgresAuthTestConn) beginTx() (driver.Tx, error) { + if c.tx != nil { + return nil, errors.New("transaction already active") + } + c.backend.mu.Lock() + tx := &postgresAuthTestTx{ + conn: c, + content: append([]byte(nil), c.backend.content...), + hasContent: c.backend.hasContent, + } + c.backend.mu.Unlock() + c.tx = tx + return tx, nil +} + +type postgresAuthTestTx struct { + conn *postgresAuthTestConn + content []byte + hasContent bool + done bool +} + +func (tx *postgresAuthTestTx) exec(query string, args []driver.NamedValue) (driver.Result, error) { + if err := tx.conn.backend.call(query, args); err != nil { + return nil, err + } + if strings.HasPrefix(strings.TrimSpace(query), "INSERT INTO") && strings.Contains(query, "DO NOTHING") && !tx.hasContent { + tx.conn.backend.mu.Lock() + if tx.conn.backend.hasContent { + tx.content = append([]byte(nil), tx.conn.backend.content...) + tx.hasContent = true + tx.conn.backend.mu.Unlock() + return driver.RowsAffected(0), nil + } + tx.conn.backend.mu.Unlock() + } + return postgresAuthTestExec(query, args, &tx.content, &tx.hasContent) +} + +func (tx *postgresAuthTestTx) query(query string, args []driver.NamedValue) (driver.Rows, error) { + if err := tx.conn.backend.call(query, args); err != nil { + return nil, err + } + return postgresAuthTestQuery(query, &tx.content, &tx.hasContent) +} + +func (tx *postgresAuthTestTx) Commit() error { + if tx.done { + return errors.New("transaction already done") + } + tx.conn.backend.mu.Lock() + tx.conn.backend.content = append([]byte(nil), tx.content...) + tx.conn.backend.hasContent = tx.hasContent + tx.conn.backend.mu.Unlock() + tx.done = true + tx.conn.tx = nil + return nil +} + +func (tx *postgresAuthTestTx) Rollback() error { + if tx.done { + return sql.ErrTxDone + } + tx.done = true + tx.conn.tx = nil + return nil +} + +type postgresAuthTestStorage struct { + path string + data []byte + mode fs.FileMode +} + +func (s *postgresAuthTestStorage) SaveTokenToFile(path string) error { + s.path = path + mode := s.mode + if mode == 0 { + mode = 0o600 + } + return os.WriteFile(path, s.data, mode) +} + +func TestPostgresStoreSavePersistsBeforeLocalPublication(t *testing.T) { + previous := []byte(`{"value":"previous"}`) + backend := &postgresAuthTestBackend{} + backend.setContent(previous) + store := newPostgresAuthStoreForTest(t, backend) + path := filepath.Join(store.authDir, "credential.json") + if err := os.WriteFile(path, previous, 0o600); err != nil { + t.Fatalf("write previous auth: %v", err) + } + store.renameFile = func(oldPath, newPath string) error { + assertPostgresAuthLocal(t, path, previous) + durable, exists := backend.durableSnapshot() + if !exists { + t.Fatal("database record missing before local publication") + } + assertPostgresAuthJSONValue(t, durable, "candidate") + return os.Rename(oldPath, newPath) + } + + _, err := store.Save(context.Background(), &cliproxyauth.Auth{ + ID: "credential.json", + Metadata: map[string]any{"value": "candidate"}, + }) + if err != nil { + t.Fatalf("Save() error = %v", err) + } + published, errRead := os.ReadFile(path) + if errRead != nil { + t.Fatalf("read published auth: %v", errRead) + } + assertPostgresAuthJSONValue(t, published, "candidate") + assertPostgresAuthNoTemp(t, path) +} + +func TestPostgresStoreSaveDatabaseFailureDoesNotPublishLocalFile(t *testing.T) { + for _, test := range []struct { + name string + previous []byte + }{ + {name: "existing", previous: []byte(`{"value":"previous"}`)}, + {name: "new"}, + } { + t.Run(test.name, func(t *testing.T) { + backend := &postgresAuthTestBackend{failAt: 1} + if test.previous != nil { + backend.setContent(test.previous) + } + store := newPostgresAuthStoreForTest(t, backend) + path := filepath.Join(store.authDir, "credential.json") + if test.previous != nil { + if err := os.WriteFile(path, test.previous, 0o600); err != nil { + t.Fatalf("write previous auth: %v", err) + } + } + + _, err := store.Save(context.Background(), &cliproxyauth.Auth{ + ID: "credential.json", + Metadata: map[string]any{"value": "candidate"}, + }) + if err == nil { + t.Fatal("Save() succeeded, want database error") + } + assertPostgresAuthLocal(t, path, test.previous) + assertPostgresAuthNoTemp(t, path) + }) + } +} + +func TestPostgresStoreSavePublishFailureRestoresDatabase(t *testing.T) { + t.Run("existing", func(t *testing.T) { + localPrevious := []byte(`{"value":"local"}`) + durablePrevious := []byte(`{"value":"durable"}`) + backend := &postgresAuthTestBackend{} + backend.setContent(durablePrevious) + store := newPostgresAuthStoreForTest(t, backend) + path := filepath.Join(store.authDir, "credential.json") + if err := os.WriteFile(path, localPrevious, 0o600); err != nil { + t.Fatalf("write previous auth: %v", err) + } + store.renameFile = func(string, string) error { return errors.New("publish rejected") } + + _, err := store.Save(context.Background(), &cliproxyauth.Auth{ + ID: "credential.json", + Metadata: map[string]any{"value": "candidate"}, + }) + if err == nil || !strings.Contains(err.Error(), "publish rejected") { + t.Fatalf("Save() error = %v, want publish error", err) + } + durable, exists := backend.durableSnapshot() + if !exists || !bytes.Equal(durable, durablePrevious) { + t.Fatalf("database record = (%q, %t), want (%q, true)", durable, exists, durablePrevious) + } + calls := backend.snapshotCalls() + if len(calls) != 3 || !strings.Contains(calls[2].query, "content = $3") { + t.Fatalf("database calls = %#v, want conditional restore", calls) + } + assertPostgresAuthLocal(t, path, localPrevious) + assertPostgresAuthNoTemp(t, path) + }) + + t.Run("new", func(t *testing.T) { + backend := &postgresAuthTestBackend{} + store := newPostgresAuthStoreForTest(t, backend) + path := filepath.Join(store.authDir, "credential.json") + store.renameFile = func(string, string) error { return errors.New("publish rejected") } + + _, err := store.Save(context.Background(), &cliproxyauth.Auth{ + ID: "credential.json", + Metadata: map[string]any{"value": "candidate"}, + }) + if err == nil || !strings.Contains(err.Error(), "publish rejected") { + t.Fatalf("Save() error = %v, want publish error", err) + } + if durable, exists := backend.durableSnapshot(); exists { + t.Fatalf("database record = %q, want missing", durable) + } + calls := backend.snapshotCalls() + if len(calls) != 3 || !strings.Contains(calls[2].query, "content = $2") { + t.Fatalf("database calls = %#v, want candidate-bound delete", calls) + } + assertPostgresAuthLocal(t, path, nil) + assertPostgresAuthNoTemp(t, path) + }) +} + +func TestPostgresStoreSaveRollbackConflictPreservesNewerDatabaseRecord(t *testing.T) { + previous := []byte(`{"value":"previous"}`) + newer := []byte(`{"value":"newer"}`) + backend := &postgresAuthTestBackend{} + backend.setContent(previous) + backend.inspect = func(call postgresAuthTestCall) { + if strings.HasPrefix(strings.TrimSpace(call.query), "UPDATE") && strings.Contains(call.query, "content = $3") { + backend.setContent(newer) + } + } + store := newPostgresAuthStoreForTest(t, backend) + path := filepath.Join(store.authDir, "credential.json") + if err := os.WriteFile(path, previous, 0o600); err != nil { + t.Fatalf("write previous auth: %v", err) + } + store.renameFile = func(string, string) error { return errors.New("publish rejected") } + + _, err := store.Save(context.Background(), &cliproxyauth.Auth{ + ID: "credential.json", + Metadata: map[string]any{"value": "candidate"}, + }) + if err == nil || !strings.Contains(err.Error(), "database rollback failed") { + t.Fatalf("Save() error = %v, want rollback conflict", err) + } + durable, exists := backend.durableSnapshot() + if !exists || !bytes.Equal(durable, newer) { + t.Fatalf("database record = (%q, %t), want (%q, true)", durable, exists, newer) + } + assertPostgresAuthLocal(t, path, previous) + assertPostgresAuthNoTemp(t, path) +} + +func TestPostgresStoreSaveInsertConflictRestoresConcurrentRecord(t *testing.T) { + localPrevious := []byte(`{"value":"local"}`) + concurrent := []byte(`{"value":"concurrent"}`) + backend := &postgresAuthTestBackend{} + backend.inspect = func(call postgresAuthTestCall) { + if strings.HasPrefix(strings.TrimSpace(call.query), "INSERT INTO") { + backend.setContent(concurrent) + } + } + store := newPostgresAuthStoreForTest(t, backend) + path := filepath.Join(store.authDir, "credential.json") + if err := os.WriteFile(path, localPrevious, 0o600); err != nil { + t.Fatalf("write previous auth: %v", err) + } + store.renameFile = func(string, string) error { return errors.New("publish rejected") } + + _, err := store.Save(context.Background(), &cliproxyauth.Auth{ + ID: "credential.json", + Metadata: map[string]any{"value": "candidate"}, + }) + if err == nil { + t.Fatal("Save() succeeded, want publish error") + } + durable, exists := backend.durableSnapshot() + if !exists || !bytes.Equal(durable, concurrent) { + t.Fatalf("database record = (%q, %t), want (%q, true)", durable, exists, concurrent) + } + calls := backend.snapshotCalls() + if len(calls) != 5 || !strings.Contains(calls[0].query, "FOR UPDATE") || !strings.Contains(calls[1].query, "DO NOTHING") || !strings.Contains(calls[2].query, "FOR UPDATE") || !strings.HasPrefix(strings.TrimSpace(calls[3].query), "UPDATE") || !strings.Contains(calls[4].query, "content = $3") { + t.Fatalf("database calls = %#v, want insert conflict retry and conditional restore", calls) + } + assertPostgresAuthLocal(t, path, localPrevious) + assertPostgresAuthNoTemp(t, path) +} + +func TestPostgresStoreUnlockFailureDiscardsConnection(t *testing.T) { + for _, test := range []struct { + name string + missing bool + err error + }{ + {name: "missing", missing: true}, + {name: "query-error", err: errors.New("unlock rejected")}, + } { + t.Run(test.name, func(t *testing.T) { + backend := &postgresAuthTestBackend{unlockMissing: test.missing, unlockErr: test.err} + store := newPostgresAuthStoreForTest(t, backend) + err := store.withAuthLock(context.Background(), "credential.json", func(*sql.Conn) error { return nil }) + if err == nil || !strings.Contains(err.Error(), "unlock auth record") { + t.Fatalf("withAuthLock() error = %v", err) + } + if got := store.db.Stats().OpenConnections; got != 0 { + t.Fatalf("open connections = %d, want failed session discarded", got) + } + }) + } +} + +func TestPostgresStoreSaveUnchangedLocalPersistsDurableRecord(t *testing.T) { + for _, test := range []struct { + name string + previous string + metadataOnly bool + }{ + {name: "storage-missing"}, + {name: "storage-stale", previous: `{"value":"old"}`}, + {name: "metadata-missing", metadataOnly: true}, + {name: "metadata-stale", previous: `{"value":"old"}`, metadataOnly: true}, + } { + t.Run(test.name, func(t *testing.T) { + backend := &postgresAuthTestBackend{} + if test.previous != "" { + backend.setContent([]byte(test.previous)) + } + store := newPostgresAuthStoreForTest(t, backend) + candidate := []byte(`{"disabled":false,"value":"candidate"}`) + path := filepath.Join(store.authDir, "credential.json") + if err := os.WriteFile(path, candidate, 0o600); err != nil { + t.Fatal(err) + } + store.renameFile = func(string, string) error { + t.Fatal("unchanged local file should not be replaced") + return nil + } + auth := &cliproxyauth.Auth{ID: "credential.json", Storage: &postgresAuthTestStorage{data: candidate}} + if test.metadataOnly { + auth.Storage = nil + auth.Metadata = map[string]any{"value": "candidate"} + } + if _, err := store.Save(context.Background(), auth); err != nil { + t.Fatal(err) + } + durable, exists := backend.durableSnapshot() + if !exists || !bytes.Equal(durable, candidate) { + t.Fatalf("durable record = %q, exists = %t", durable, exists) + } + assertPostgresAuthNoTemp(t, path) + }) + } +} + +func TestPostgresStoreSaveMissingOutputFails(t *testing.T) { + backend := &postgresAuthTestBackend{} + store := newPostgresAuthStoreForTest(t, backend) + auth := &cliproxyauth.Auth{ID: "credential.json", Storage: &empty.EmptyStorage{}} + path, err := store.Save(context.Background(), auth) + if err == nil || path != "" || auth.Attributes[cliproxyauth.AttributeSourceBackend] != "" { + t.Fatalf("Save() = (%q, %v), attributes = %v", path, err, auth.Attributes) + } + if len(backend.snapshotCalls()) != 0 { + t.Fatal("missing output must not reach the database") + } + assertPostgresAuthNoTemp(t, filepath.Join(store.authDir, auth.ID)) +} + +func TestPostgresStoreSaveUsesTemporaryPathForTokenStorage(t *testing.T) { + backend := &postgresAuthTestBackend{} + store := newPostgresAuthStoreForTest(t, backend) + storage := &postgresAuthTestStorage{data: []byte(`{"type":"codex","token":"value"}`), mode: 0o644} + store.renameFile = func(oldPath, newPath string) error { + if got, want := oldPath, newPath+".tmp"; got != want { + t.Fatalf("temporary path = %q, want %q", got, want) + } + info, errStat := os.Stat(oldPath) + if errStat != nil { + t.Fatalf("stat temporary auth: %v", errStat) + } + if got := info.Mode().Perm(); got != 0o600 { + t.Fatalf("temporary auth mode = %04o, want 0600", got) + } + return os.Rename(oldPath, newPath) + } + auth := &cliproxyauth.Auth{ + ID: "credential.json", + Metadata: map[string]any{"type": "codex"}, + Storage: storage, + } + + path, err := store.Save(context.Background(), auth) + if err != nil { + t.Fatalf("Save() error = %v", err) + } + if storage.path != path+".tmp" { + t.Fatalf("storage path = %q, want %q", storage.path, path+".tmp") + } + assertPostgresAuthLocal(t, path, storage.data) + if got := auth.Attributes[cliproxyauth.AttributeSourceBackend]; got != cliproxyauth.AuthSourcePostgres { + t.Fatalf("source backend = %q, want %q", got, cliproxyauth.AuthSourcePostgres) + } + assertPostgresAuthNoTemp(t, path) +} + +func TestPostgresStoreSaveEmptyPayloadDeletesDurableRecord(t *testing.T) { + previous := []byte(`{"value":"previous"}`) + backend := &postgresAuthTestBackend{} + backend.setContent(previous) + store := newPostgresAuthStoreForTest(t, backend) + path := filepath.Join(store.authDir, "credential.json") + if err := os.WriteFile(path, previous, 0o600); err != nil { + t.Fatalf("write previous auth: %v", err) + } + + gotPath, err := store.Save(context.Background(), &cliproxyauth.Auth{ + ID: "credential.json", + Storage: &postgresAuthTestStorage{}, + }) + if err != nil { + t.Fatalf("Save() error = %v", err) + } + if gotPath != path { + t.Fatalf("Save() path = %q, want %q", gotPath, path) + } + if durable, exists := backend.durableSnapshot(); exists { + t.Fatalf("database record = %q, want missing", durable) + } + assertPostgresAuthLocal(t, path, []byte{}) + assertPostgresAuthNoTemp(t, path) +} + +func newPostgresAuthStoreForTest(t *testing.T, backend *postgresAuthTestBackend) *PostgresStore { + t.Helper() + authDir := filepath.Join(t.TempDir(), "auths") + if err := os.MkdirAll(authDir, 0o700); err != nil { + t.Fatalf("create auth dir: %v", err) + } + db := sql.OpenDB(&postgresAuthTestConnector{backend: backend}) + db.SetMaxOpenConns(1) + t.Cleanup(func() { + if errClose := db.Close(); errClose != nil { + t.Errorf("close database: %v", errClose) + } + }) + return &PostgresStore{ + db: db, + cfg: PostgresStoreConfig{AuthTable: defaultAuthTable}, + authDir: authDir, + } +} + +func assertPostgresAuthJSONValue(t *testing.T, data []byte, want string) { + t.Helper() + if got := string(data); !strings.Contains(got, `"value":"`+want+`"`) { + t.Fatalf("auth JSON = %q, want value %q", got, want) + } +} + +func assertPostgresAuthLocal(t *testing.T, path string, want []byte) { + t.Helper() + got, err := os.ReadFile(path) + if want == nil { + if !errors.Is(err, os.ErrNotExist) { + t.Fatalf("read local auth error = %v, want not exist", err) + } + return + } + if err != nil { + t.Fatalf("read local auth: %v", err) + } + if !bytes.Equal(got, want) { + t.Fatalf("local auth bytes = %q, want %q", got, want) + } +} + +func assertPostgresAuthNoTemp(t *testing.T, path string) { + t.Helper() + if _, err := os.Stat(path + ".tmp"); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("stat temp auth error = %v, want not exist", err) + } +} + +func clonePostgresAuthTestArgs(args []driver.NamedValue) []driver.NamedValue { + cloned := make([]driver.NamedValue, len(args)) + copy(cloned, args) + for i := range cloned { + if value, ok := cloned[i].Value.([]byte); ok { + cloned[i].Value = append([]byte(nil), value...) + } + } + return cloned +} + +func postgresAuthTestValue(args []driver.NamedValue, index int) []byte { + if index >= len(args) { + return nil + } + switch value := args[index].Value.(type) { + case []byte: + return append([]byte(nil), value...) + case string: + return []byte(value) + default: + return nil + } +} + +var _ driver.ExecerContext = (*postgresAuthTestConn)(nil) +var _ driver.QueryerContext = (*postgresAuthTestConn)(nil) +var _ driver.ConnBeginTx = (*postgresAuthTestConn)(nil) +var _ driver.Connector = (*postgresAuthTestConnector)(nil) From 58f93157a349deba0275cd70b8b4caf421a339d8 Mon Sep 17 00:00:00 2001 From: Aikins Laryea Date: Sun, 13 Sep 2026 15:59:20 +0000 Subject: [PATCH 2/5] store: restore auth timestamps after failed publication --- internal/store/postgresstore.go | 63 +++++----- .../store/postgresstore_integration_test.go | 91 +++++++++++++++ internal/store/postgresstore_test.go | 109 ++++++++++++------ 3 files changed, 202 insertions(+), 61 deletions(-) diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index 2b9690eb2f8..ebbabb92a94 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -296,7 +296,7 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri err = s.withAuthLock(ctx, relID, func(conn *sql.Conn) error { var ( - durablePrevious []byte + durablePrevious postgresAuthRecord durablePreviousExists bool durableChanged bool ) @@ -655,11 +655,17 @@ func (s *PostgresStore) withAuthLock(ctx context.Context, relID string, save fun return save(conn) } -func (s *PostgresStore) replaceAuthRecord(ctx context.Context, conn *sql.Conn, relID string, candidate []byte) (previous []byte, previousExists bool, err error) { +type postgresAuthRecord struct { + content []byte + createdAt time.Time + updatedAt time.Time +} + +func (s *PostgresStore) replaceAuthRecord(ctx context.Context, conn *sql.Conn, relID string, candidate []byte) (previous postgresAuthRecord, previousExists bool, err error) { // READ COMMITTED lets the retry observe a row that won ON CONFLICT. tx, errBegin := conn.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelReadCommitted}) if errBegin != nil { - return nil, false, fmt.Errorf("postgres store: begin auth replacement: %w", errBegin) + return previous, false, fmt.Errorf("postgres store: begin auth replacement: %w", errBegin) } defer func() { if err == nil { @@ -671,7 +677,7 @@ func (s *PostgresStore) replaceAuthRecord(ctx context.Context, conn *sql.Conn, r }() table := s.fullTableName(s.cfg.AuthTable) - selectQuery := fmt.Sprintf("SELECT content FROM %s WHERE id = $1 FOR UPDATE", table) + selectQuery := fmt.Sprintf("SELECT content, created_at, updated_at FROM %s WHERE id = $1 FOR UPDATE", table) updateQuery := fmt.Sprintf("UPDATE %s SET content = $2, updated_at = NOW() WHERE id = $1", table) insertQuery := fmt.Sprintf(` INSERT INTO %s (id, content, created_at, updated_at) @@ -680,71 +686,74 @@ func (s *PostgresStore) replaceAuthRecord(ctx context.Context, conn *sql.Conn, r `, table) for { var content string - errScan := tx.QueryRowContext(ctx, selectQuery, relID).Scan(&content) + errScan := tx.QueryRowContext(ctx, selectQuery, relID).Scan(&content, &previous.createdAt, &previous.updatedAt) switch { case errScan == nil: + previous.content = []byte(content) result, errUpdate := tx.ExecContext(ctx, updateQuery, relID, json.RawMessage(candidate)) if errUpdate != nil { err = fmt.Errorf("postgres store: replace auth record: %w", errUpdate) - return nil, false, err + return previous, false, err } rows, errRows := result.RowsAffected() if errRows != nil { err = fmt.Errorf("postgres store: inspect auth replacement: %w", errRows) - return nil, false, err + return previous, false, err } if rows != 1 { err = fmt.Errorf("postgres store: locked auth record disappeared before replacement") - return nil, false, err + return previous, false, err } if errCommit := tx.Commit(); errCommit != nil { err = fmt.Errorf("postgres store: commit auth replacement: %w", errCommit) - return nil, false, err + return previous, false, err } - return []byte(content), true, nil + return previous, true, nil case !errors.Is(errScan, sql.ErrNoRows): err = fmt.Errorf("postgres store: lock auth record: %w", errScan) - return nil, false, err + return previous, false, err } result, errInsert := tx.ExecContext(ctx, insertQuery, relID, json.RawMessage(candidate)) if errInsert != nil { err = fmt.Errorf("postgres store: insert auth record: %w", errInsert) - return nil, false, err + return previous, false, err } rows, errRows := result.RowsAffected() if errRows != nil { err = fmt.Errorf("postgres store: inspect auth insert: %w", errRows) - return nil, false, err + return previous, false, err } if rows == 0 { continue } if rows != 1 { err = fmt.Errorf("postgres store: inserted %d auth records, want 1", rows) - return nil, false, err + return previous, false, err } if errCommit := tx.Commit(); errCommit != nil { err = fmt.Errorf("postgres store: commit auth insert: %w", errCommit) - return nil, false, err + return previous, false, err } - return nil, false, nil + return previous, false, nil } } -func (s *PostgresStore) deleteAuthRecordReturning(ctx context.Context, conn *sql.Conn, relID string) ([]byte, bool, error) { - query := fmt.Sprintf("DELETE FROM %s WHERE id = $1 RETURNING content", s.fullTableName(s.cfg.AuthTable)) +func (s *PostgresStore) deleteAuthRecordReturning(ctx context.Context, conn *sql.Conn, relID string) (postgresAuthRecord, bool, error) { + query := fmt.Sprintf("DELETE FROM %s WHERE id = $1 RETURNING content, created_at, updated_at", s.fullTableName(s.cfg.AuthTable)) + var previous postgresAuthRecord var content string - if err := conn.QueryRowContext(ctx, query, relID).Scan(&content); err != nil { + if err := conn.QueryRowContext(ctx, query, relID).Scan(&content, &previous.createdAt, &previous.updatedAt); err != nil { if errors.Is(err, sql.ErrNoRows) { - return nil, false, nil + return previous, false, nil } - return nil, false, fmt.Errorf("postgres store: delete auth record: %w", err) + return previous, false, fmt.Errorf("postgres store: delete auth record: %w", err) } - return []byte(content), true, nil + previous.content = []byte(content) + return previous, true, nil } -func (s *PostgresStore) rollbackAuthRecord(ctx context.Context, conn *sql.Conn, relID string, candidate, previous []byte, previousExists bool) error { +func (s *PostgresStore) rollbackAuthRecord(ctx context.Context, conn *sql.Conn, relID string, candidate []byte, previous postgresAuthRecord, previousExists bool) error { var ( result sql.Result err error @@ -752,13 +761,13 @@ func (s *PostgresStore) rollbackAuthRecord(ctx context.Context, conn *sql.Conn, if previousExists && len(candidate) == 0 { query := fmt.Sprintf(` INSERT INTO %s (id, content, created_at, updated_at) - VALUES ($1, $2, NOW(), NOW()) + VALUES ($1, $2, $3, $4) ON CONFLICT (id) DO NOTHING `, s.fullTableName(s.cfg.AuthTable)) - result, err = conn.ExecContext(ctx, query, relID, json.RawMessage(previous)) + result, err = conn.ExecContext(ctx, query, relID, json.RawMessage(previous.content), previous.createdAt, previous.updatedAt) } else if previousExists { - query := fmt.Sprintf("UPDATE %s SET content = $2, updated_at = NOW() WHERE id = $1 AND content = $3", s.fullTableName(s.cfg.AuthTable)) - result, err = conn.ExecContext(ctx, query, relID, json.RawMessage(previous), json.RawMessage(candidate)) + query := fmt.Sprintf("UPDATE %s SET content = $2, created_at = $4, updated_at = $5 WHERE id = $1 AND content = $3", s.fullTableName(s.cfg.AuthTable)) + result, err = conn.ExecContext(ctx, query, relID, json.RawMessage(previous.content), json.RawMessage(candidate), previous.createdAt, previous.updatedAt) } else { query := fmt.Sprintf("DELETE FROM %s WHERE id = $1 AND content = $2", s.fullTableName(s.cfg.AuthTable)) result, err = conn.ExecContext(ctx, query, relID, json.RawMessage(candidate)) diff --git a/internal/store/postgresstore_integration_test.go b/internal/store/postgresstore_integration_test.go index 2df1fc57ba8..277dac85501 100644 --- a/internal/store/postgresstore_integration_test.go +++ b/internal/store/postgresstore_integration_test.go @@ -155,3 +155,94 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { }) } } + +func TestPostgresStorePublicationRestoresCompleteRecord(t *testing.T) { + dsn := os.Getenv("TEST_POSTGRES_DSN") + if dsn == "" { + t.Skip("TEST_POSTGRES_DSN must point to a disposable PostgreSQL database") + } + for _, test := range []struct { + name string + existing bool + candidate string + fail bool + cancel bool + }{ + {"update-failure", true, `{"value":"candidate"}`, true, false}, + {"delete-failure", true, "", true, false}, + {"cancelled-update", true, `{"value":"candidate"}`, true, true}, + {"cancelled-delete", true, "", true, true}, + {"insert-failure", false, `{"value":"candidate"}`, true, false}, + {"successful-update", true, `{"value":"candidate"}`, false, false}, + } { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + db, err := sql.Open("pgx", dsn) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if errClose := db.Close(); errClose != nil { + t.Error(errClose) + } + }) + table := fmt.Sprintf("auth_record_test_%d", time.Now().UnixNano()) + if _, err = db.ExecContext(ctx, "CREATE TABLE "+table+" (id TEXT PRIMARY KEY, content JSONB NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW())"); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if _, errDrop := db.ExecContext(ctx, "DROP TABLE "+table); errDrop != nil { + t.Error(errDrop) + } + }) + store := &PostgresStore{db: db, cfg: PostgresStoreConfig{AuthTable: table}, authDir: t.TempDir()} + const id = "credential.json" + const previous = `{"type":"codex","value":"before"}` + created := time.Date(2001, 2, 3, 4, 5, 6, 123456000, time.UTC) + updated := time.Date(2002, 3, 4, 5, 6, 7, 654321000, time.UTC) + if test.existing { + if _, err = db.ExecContext(ctx, "INSERT INTO "+table+" (id, content, created_at, updated_at) VALUES ($1, $2, $3, $4)", id, previous, created, updated); err != nil { + t.Fatal(err) + } + if err = os.WriteFile(filepath.Join(store.authDir, id), []byte(previous), 0o600); err != nil { + t.Fatal(err) + } + } + saveCtx, cancel := context.WithCancel(ctx) + defer cancel() + publishErr := errors.New("publication rejected") + if test.fail { + store.renameFile = func(string, string) error { + if test.cancel { + cancel() + } + return publishErr + } + } + _, err = store.Save(saveCtx, &cliproxyauth.Auth{ID: id, Storage: &postgresAuthTestStorage{data: []byte(test.candidate)}}) + if test.fail && (!errors.Is(err, publishErr) || strings.Contains(err.Error(), "rollback failed")) || !test.fail && err != nil { + t.Fatalf("Save() error = %v", err) + } + var content string + var gotCreated, gotUpdated time.Time + err = db.QueryRowContext(ctx, "SELECT content, created_at, updated_at FROM "+table+" WHERE id = $1", id).Scan(&content, &gotCreated, &gotUpdated) + if !test.existing { + if !errors.Is(err, sql.ErrNoRows) { + t.Fatalf("failed insertion left a record: %v", err) + } + return + } + wantContent := test.candidate + if test.fail { + wantContent = previous + } + if err != nil || !jsonEqual([]byte(content), []byte(wantContent)) || !gotCreated.Equal(created) || test.fail && !gotUpdated.Equal(updated) || !test.fail && !gotUpdated.After(updated) { + t.Fatalf("record = (%s, %v, %v), error = %v", content, gotCreated, gotUpdated, err) + } + listed, errList := store.List(ctx) + if errList != nil || len(listed) != 1 || !listed[0].CreatedAt.Equal(gotCreated) || !listed[0].UpdatedAt.Equal(gotUpdated) { + t.Fatalf("List() did not preserve record timestamps: %+v, %v", listed, errList) + } + }) + } +} diff --git a/internal/store/postgresstore_test.go b/internal/store/postgresstore_test.go index 80b04f827f3..29d15a2b320 100644 --- a/internal/store/postgresstore_test.go +++ b/internal/store/postgresstore_test.go @@ -6,6 +6,7 @@ import ( "database/sql" "database/sql/driver" "errors" + "fmt" "io" "io/fs" "os" @@ -13,6 +14,7 @@ import ( "strings" "sync" "testing" + "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/empty" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -30,6 +32,8 @@ type postgresAuthTestBackend struct { inspect func(postgresAuthTestCall) content []byte hasContent bool + createdAt time.Time + updatedAt time.Time unlockMissing bool unlockErr error } @@ -60,7 +64,7 @@ func (b *postgresAuthTestBackend) exec(query string, args []driver.NamedValue) ( } b.mu.Lock() defer b.mu.Unlock() - return postgresAuthTestExec(query, args, &b.content, &b.hasContent) + return postgresAuthTestExec(query, args, &b.content, &b.hasContent, &b.createdAt, &b.updatedAt) } func (b *postgresAuthTestBackend) query(query string, args []driver.NamedValue) (driver.Rows, error) { @@ -78,13 +82,15 @@ func (b *postgresAuthTestBackend) query(query string, args []driver.NamedValue) } b.mu.Lock() defer b.mu.Unlock() - return postgresAuthTestQuery(query, &b.content, &b.hasContent) + return postgresAuthTestQuery(query, &b.content, &b.hasContent, b.createdAt, b.updatedAt) } func (b *postgresAuthTestBackend) setContent(data []byte) { b.mu.Lock() b.content = append([]byte(nil), data...) b.hasContent = true + b.createdAt = time.Unix(100, 0).UTC() + b.updatedAt = time.Unix(200, 0).UTC() b.mu.Unlock() } @@ -102,7 +108,7 @@ func (b *postgresAuthTestBackend) snapshotCalls() []postgresAuthTestCall { return calls } -func postgresAuthTestExec(query string, args []driver.NamedValue, content *[]byte, hasContent *bool) (driver.Result, error) { +func postgresAuthTestExec(query string, args []driver.NamedValue, content *[]byte, hasContent *bool, createdAt, updatedAt *time.Time) (driver.Result, error) { trimmed := strings.TrimSpace(query) switch { case strings.HasPrefix(trimmed, "INSERT INTO"): @@ -111,6 +117,10 @@ func postgresAuthTestExec(query string, args []driver.NamedValue, content *[]byt } *content = postgresAuthTestValue(args, 1) *hasContent = true + *createdAt, *updatedAt = time.Unix(300, 0).UTC(), time.Unix(300, 0).UTC() + if len(args) == 4 { + *createdAt, *updatedAt = args[2].Value.(time.Time), args[3].Value.(time.Time) + } return driver.RowsAffected(1), nil case strings.HasPrefix(trimmed, "UPDATE"): if !*hasContent { @@ -120,6 +130,10 @@ func postgresAuthTestExec(query string, args []driver.NamedValue, content *[]byt return driver.RowsAffected(0), nil } *content = postgresAuthTestValue(args, 1) + *updatedAt = time.Unix(300, 0).UTC() + if len(args) == 5 { + *createdAt, *updatedAt = args[3].Value.(time.Time), args[4].Value.(time.Time) + } return driver.RowsAffected(1), nil case strings.HasPrefix(trimmed, "DELETE") && len(args) > 1: if !*hasContent || !bytes.Equal(*content, postgresAuthTestValue(args, 1)) { @@ -137,7 +151,7 @@ func postgresAuthTestExec(query string, args []driver.NamedValue, content *[]byt } } -func postgresAuthTestQuery(query string, content *[]byte, hasContent *bool) (driver.Rows, error) { +func postgresAuthTestQuery(query string, content *[]byte, hasContent *bool, createdAt, updatedAt time.Time) (driver.Rows, error) { trimmed := strings.TrimSpace(query) rows := &postgresAuthTestRows{columns: []string{"content"}} switch { @@ -154,6 +168,12 @@ func postgresAuthTestQuery(query string, content *[]byte, hasContent *bool) (dri default: return nil, errors.New("query unsupported") } + if strings.Contains(query, "created_at") { + rows.columns = append(rows.columns, "created_at", "updated_at") + for i := range rows.values { + rows.values[i] = append(rows.values[i], createdAt, updatedAt) + } + } return rows, nil } @@ -234,6 +254,8 @@ func (c *postgresAuthTestConn) beginTx() (driver.Tx, error) { conn: c, content: append([]byte(nil), c.backend.content...), hasContent: c.backend.hasContent, + createdAt: c.backend.createdAt, + updatedAt: c.backend.updatedAt, } c.backend.mu.Unlock() c.tx = tx @@ -244,6 +266,8 @@ type postgresAuthTestTx struct { conn *postgresAuthTestConn content []byte hasContent bool + createdAt time.Time + updatedAt time.Time done bool } @@ -256,19 +280,21 @@ func (tx *postgresAuthTestTx) exec(query string, args []driver.NamedValue) (driv if tx.conn.backend.hasContent { tx.content = append([]byte(nil), tx.conn.backend.content...) tx.hasContent = true + tx.createdAt = tx.conn.backend.createdAt + tx.updatedAt = tx.conn.backend.updatedAt tx.conn.backend.mu.Unlock() return driver.RowsAffected(0), nil } tx.conn.backend.mu.Unlock() } - return postgresAuthTestExec(query, args, &tx.content, &tx.hasContent) + return postgresAuthTestExec(query, args, &tx.content, &tx.hasContent, &tx.createdAt, &tx.updatedAt) } func (tx *postgresAuthTestTx) query(query string, args []driver.NamedValue) (driver.Rows, error) { if err := tx.conn.backend.call(query, args); err != nil { return nil, err } - return postgresAuthTestQuery(query, &tx.content, &tx.hasContent) + return postgresAuthTestQuery(query, &tx.content, &tx.hasContent, tx.createdAt, tx.updatedAt) } func (tx *postgresAuthTestTx) Commit() error { @@ -278,6 +304,8 @@ func (tx *postgresAuthTestTx) Commit() error { tx.conn.backend.mu.Lock() tx.conn.backend.content = append([]byte(nil), tx.content...) tx.conn.backend.hasContent = tx.hasContent + tx.conn.backend.createdAt = tx.createdAt + tx.conn.backend.updatedAt = tx.updatedAt tx.conn.backend.mu.Unlock() tx.done = true tx.conn.tx = nil @@ -377,36 +405,49 @@ func TestPostgresStoreSaveDatabaseFailureDoesNotPublishLocalFile(t *testing.T) { } func TestPostgresStoreSavePublishFailureRestoresDatabase(t *testing.T) { - t.Run("existing", func(t *testing.T) { - localPrevious := []byte(`{"value":"local"}`) - durablePrevious := []byte(`{"value":"durable"}`) - backend := &postgresAuthTestBackend{} - backend.setContent(durablePrevious) - store := newPostgresAuthStoreForTest(t, backend) - path := filepath.Join(store.authDir, "credential.json") - if err := os.WriteFile(path, localPrevious, 0o600); err != nil { - t.Fatalf("write previous auth: %v", err) - } - store.renameFile = func(string, string) error { return errors.New("publish rejected") } + for _, empty := range []bool{false, true} { + t.Run(fmt.Sprintf("existing/empty=%t", empty), func(t *testing.T) { + localPrevious := []byte(`{"value":"local"}`) + durablePrevious := []byte(`{"value":"durable"}`) + backend := &postgresAuthTestBackend{} + backend.setContent(durablePrevious) + store := newPostgresAuthStoreForTest(t, backend) + path := filepath.Join(store.authDir, "credential.json") + if err := os.WriteFile(path, localPrevious, 0o600); err != nil { + t.Fatalf("write previous auth: %v", err) + } + store.renameFile = func(string, string) error { return errors.New("publish rejected") } - _, err := store.Save(context.Background(), &cliproxyauth.Auth{ - ID: "credential.json", - Metadata: map[string]any{"value": "candidate"}, + auth := &cliproxyauth.Auth{ + ID: "credential.json", + Metadata: map[string]any{"value": "candidate"}, + } + if empty { + auth.Storage = &postgresAuthTestStorage{} + } + _, err := store.Save(context.Background(), auth) + if err == nil || !strings.Contains(err.Error(), "publish rejected") { + t.Fatalf("Save() error = %v, want publish error", err) + } + durable, exists := backend.durableSnapshot() + if !exists || !bytes.Equal(durable, durablePrevious) { + t.Fatalf("database record = (%q, %t), want (%q, true)", durable, exists, durablePrevious) + } + if !backend.createdAt.Equal(time.Unix(100, 0)) || !backend.updatedAt.Equal(time.Unix(200, 0)) { + t.Fatalf("database timestamps = (%v, %v), want original creation and update times", backend.createdAt, backend.updatedAt) + } + calls := backend.snapshotCalls() + if empty { + if len(calls) != 2 || !strings.Contains(calls[1].query, "DO NOTHING") { + t.Fatalf("database calls = %#v, want conditional reinsert", calls) + } + } else if len(calls) != 3 || !strings.Contains(calls[2].query, "content = $3") { + t.Fatalf("database calls = %#v, want conditional restore", calls) + } + assertPostgresAuthLocal(t, path, localPrevious) + assertPostgresAuthNoTemp(t, path) }) - if err == nil || !strings.Contains(err.Error(), "publish rejected") { - t.Fatalf("Save() error = %v, want publish error", err) - } - durable, exists := backend.durableSnapshot() - if !exists || !bytes.Equal(durable, durablePrevious) { - t.Fatalf("database record = (%q, %t), want (%q, true)", durable, exists, durablePrevious) - } - calls := backend.snapshotCalls() - if len(calls) != 3 || !strings.Contains(calls[2].query, "content = $3") { - t.Fatalf("database calls = %#v, want conditional restore", calls) - } - assertPostgresAuthLocal(t, path, localPrevious) - assertPostgresAuthNoTemp(t, path) - }) + } t.Run("new", func(t *testing.T) { backend := &postgresAuthTestBackend{} From 355e8c28b9c75fbd19bebcd6b547fdc25f5eade8 Mon Sep 17 00:00:00 2001 From: Aikins Laryea Date: Sun, 13 Sep 2026 16:32:29 +0000 Subject: [PATCH 3/5] store: isolate concurrent auth save staging --- internal/store/postgresstore.go | 21 ++++++++-------- .../store/postgresstore_integration_test.go | 24 +++++++++++++++++-- internal/store/postgresstore_test.go | 18 +++++++++----- 3 files changed, 45 insertions(+), 18 deletions(-) diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index ebbabb92a94..eead6bbe316 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -242,22 +242,18 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri return "", fmt.Errorf("postgres store: create auth directory: %w", err) } - localPrevious, errReadPrevious := os.ReadFile(path) - localExists := errReadPrevious == nil - if errReadPrevious != nil && !errors.Is(errReadPrevious, fs.ErrNotExist) { - return "", fmt.Errorf("postgres store: read existing metadata: %w", errReadPrevious) - } relID, err := s.relativeAuthID(path) if err != nil { return "", err } - tmp := path + ".tmp" - if errRemove := os.Remove(tmp); errRemove != nil && !errors.Is(errRemove, fs.ErrNotExist) { - return "", fmt.Errorf("postgres store: remove stale temp auth file: %w", errRemove) + stagedDir, errStage := os.MkdirTemp(filepath.Dir(path), filepath.Base(path)+".tmp-*") + if errStage != nil { + return "", fmt.Errorf("postgres store: create temp auth directory: %w", errStage) } + tmp := filepath.Join(stagedDir, "auth") defer func() { - if errRemove := os.Remove(tmp); errRemove != nil && !errors.Is(errRemove, fs.ErrNotExist) { - log.WithError(errRemove).Warn("postgres store: remove temporary auth file") + if errRemove := os.RemoveAll(stagedDir); errRemove != nil { + log.WithError(errRemove).Warn("postgres store: remove temporary auth directory") } }() @@ -295,6 +291,11 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri } err = s.withAuthLock(ctx, relID, func(conn *sql.Conn) error { + localPrevious, errReadPrevious := os.ReadFile(path) + localExists := errReadPrevious == nil + if errReadPrevious != nil && !errors.Is(errReadPrevious, fs.ErrNotExist) { + return fmt.Errorf("postgres store: read existing metadata: %w", errReadPrevious) + } var ( durablePrevious postgresAuthRecord durablePreviousExists bool diff --git a/internal/store/postgresstore_integration_test.go b/internal/store/postgresstore_integration_test.go index 277dac85501..2a2b103b557 100644 --- a/internal/store/postgresstore_integration_test.go +++ b/internal/store/postgresstore_integration_test.go @@ -35,6 +35,9 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { {"watcher-delete", `{"value":"old"}`, "", "", "watcher"}, {"qualified-table", `{"value":"old"}`, `{"value":"new"}`, `{"value":"new"}`, "save"}, {"cancelled-publication", `{"value":"old"}`, `{"value":"new"}`, `{"value":"new"}`, "save"}, + {"shared-failure", `{"value":"old"}`, `{"value":"first"}`, `{"value":"second"}`, "save"}, + {"shared-success", `{"value":"old"}`, `{"value":"first"}`, `{"value":"second"}`, "save"}, + {"shared-stale-read", `{"value":"old"}`, `{"value":"first"}`, `{"value":"old"}`, "save"}, } { t.Run(test.name, func(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -64,6 +67,11 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { return &PostgresStore{db: conn, cfg: PostgresStoreConfig{AuthTable: table}, authDir: t.TempDir()} } a, b := newStore(), newStore() + shared := strings.HasPrefix(test.name, "shared-") + publish := shared && test.name != "shared-failure" + if shared { + b.authDir = a.authDir + } if test.name == "qualified-table" { if err = db.QueryRowContext(ctx, "SELECT current_schema()").Scan(&b.cfg.Schema); err != nil { t.Fatal(err) @@ -87,7 +95,7 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { t.Fatal(err) } completed := make(chan error, 1) - a.renameFile = func(string, string) error { + a.renameFile = func(oldPath, newPath string) error { go func() { var errWrite error switch test.operation { @@ -123,10 +131,19 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { if test.name == "cancelled-publication" { cancelSave() } + if shared { + staged, errRead := os.ReadFile(oldPath) + if errRead != nil || !jsonEqual(staged, []byte(test.candidate)) { + t.Errorf("first staged payload = %s, error = %v, want %s", staged, errRead, test.candidate) + } + } + if publish { + return os.Rename(oldPath, newPath) + } return errors.New("publish rejected") } _, err = a.Save(saveCtx, &cliproxyauth.Auth{ID: id, Storage: &postgresAuthTestStorage{data: []byte(test.candidate)}}) - if err == nil || !strings.Contains(err.Error(), "publish rejected") || strings.Contains(err.Error(), "rollback failed") { + if publish && err != nil || !publish && (err == nil || !strings.Contains(err.Error(), "publish rejected") || strings.Contains(err.Error(), "rollback failed")) { t.Fatalf("Save() error = %v", err) } select { @@ -150,6 +167,9 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { if test.previous != "" { previous = []byte(test.previous) } + if shared { + previous = []byte(test.newer) + } assertPostgresAuthLocal(t, filepath.Join(a.authDir, id), previous) assertPostgresAuthNoTemp(t, filepath.Join(a.authDir, id)) }) diff --git a/internal/store/postgresstore_test.go b/internal/store/postgresstore_test.go index 29d15a2b320..36097d64caf 100644 --- a/internal/store/postgresstore_test.go +++ b/internal/store/postgresstore_test.go @@ -626,8 +626,8 @@ func TestPostgresStoreSaveUsesTemporaryPathForTokenStorage(t *testing.T) { store := newPostgresAuthStoreForTest(t, backend) storage := &postgresAuthTestStorage{data: []byte(`{"type":"codex","token":"value"}`), mode: 0o644} store.renameFile = func(oldPath, newPath string) error { - if got, want := oldPath, newPath+".tmp"; got != want { - t.Fatalf("temporary path = %q, want %q", got, want) + if !strings.HasPrefix(oldPath, newPath+".tmp-") || oldPath != storage.path { + t.Fatalf("temporary path = %q, want unique storage path beside %q", oldPath, newPath) } info, errStat := os.Stat(oldPath) if errStat != nil { @@ -648,8 +648,8 @@ func TestPostgresStoreSaveUsesTemporaryPathForTokenStorage(t *testing.T) { if err != nil { t.Fatalf("Save() error = %v", err) } - if storage.path != path+".tmp" { - t.Fatalf("storage path = %q, want %q", storage.path, path+".tmp") + if !strings.HasPrefix(storage.path, path+".tmp-") { + t.Fatalf("storage path = %q, want unique temporary path beside %q", storage.path, path) } assertPostgresAuthLocal(t, path, storage.data) if got := auth.Attributes[cliproxyauth.AttributeSourceBackend]; got != cliproxyauth.AuthSourcePostgres { @@ -731,8 +731,14 @@ func assertPostgresAuthLocal(t *testing.T, path string, want []byte) { func assertPostgresAuthNoTemp(t *testing.T, path string) { t.Helper() - if _, err := os.Stat(path + ".tmp"); !errors.Is(err, os.ErrNotExist) { - t.Fatalf("stat temp auth error = %v, want not exist", err) + entries, err := os.ReadDir(filepath.Dir(path)) + if err != nil { + t.Fatal(err) + } + for _, entry := range entries { + if strings.HasPrefix(entry.Name(), filepath.Base(path)+".tmp") { + t.Errorf("temporary auth file remains: %s", entry.Name()) + } } } From 54f56fc51f77ee63e59b0d1daac50d5f066d9864 Mon Sep 17 00:00:00 2001 From: Aikins Laryea Date: Sun, 13 Sep 2026 16:50:48 +0000 Subject: [PATCH 4/5] store: preserve auth timestamps for unchanged content --- internal/store/postgresstore.go | 2 +- .../store/postgresstore_integration_test.go | 24 ++++++++++++++++--- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index eead6bbe316..cee50f7787f 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -679,7 +679,7 @@ func (s *PostgresStore) replaceAuthRecord(ctx context.Context, conn *sql.Conn, r table := s.fullTableName(s.cfg.AuthTable) selectQuery := fmt.Sprintf("SELECT content, created_at, updated_at FROM %s WHERE id = $1 FOR UPDATE", table) - updateQuery := fmt.Sprintf("UPDATE %s SET content = $2, updated_at = NOW() WHERE id = $1", table) + updateQuery := fmt.Sprintf("UPDATE %s SET content = $2, updated_at = CASE WHEN content = $2 THEN updated_at ELSE NOW() END WHERE id = $1", table) insertQuery := fmt.Sprintf(` INSERT INTO %s (id, content, created_at, updated_at) VALUES ($1, $2, NOW(), NOW()) diff --git a/internal/store/postgresstore_integration_test.go b/internal/store/postgresstore_integration_test.go index 2a2b103b557..51030d811df 100644 --- a/internal/store/postgresstore_integration_test.go +++ b/internal/store/postgresstore_integration_test.go @@ -194,6 +194,9 @@ func TestPostgresStorePublicationRestoresCompleteRecord(t *testing.T) { {"cancelled-delete", true, "", true, true}, {"insert-failure", false, `{"value":"candidate"}`, true, false}, {"successful-update", true, `{"value":"candidate"}`, false, false}, + {"unchanged-storage", true, `{"type":"codex","value":"before","disabled":false}`, false, false}, + {"unchanged-metadata", true, `{"disabled":false,"value":"before","type":"codex"}`, false, false}, + {"unchanged-local-repair", true, `{"type":"codex","value":"before","disabled":false}`, false, false}, } { t.Run(test.name, func(t *testing.T) { ctx := context.Background() @@ -217,7 +220,7 @@ func TestPostgresStorePublicationRestoresCompleteRecord(t *testing.T) { }) store := &PostgresStore{db: db, cfg: PostgresStoreConfig{AuthTable: table}, authDir: t.TempDir()} const id = "credential.json" - const previous = `{"type":"codex","value":"before"}` + const previous = `{"type":"codex","value":"before","disabled":false}` created := time.Date(2001, 2, 3, 4, 5, 6, 123456000, time.UTC) updated := time.Date(2002, 3, 4, 5, 6, 7, 654321000, time.UTC) if test.existing { @@ -230,6 +233,11 @@ func TestPostgresStorePublicationRestoresCompleteRecord(t *testing.T) { } saveCtx, cancel := context.WithCancel(ctx) defer cancel() + if test.name == "unchanged-local-repair" { + if err = os.WriteFile(filepath.Join(store.authDir, id), []byte(`{"value":"stale"}`), 0o600); err != nil { + t.Fatal(err) + } + } publishErr := errors.New("publication rejected") if test.fail { store.renameFile = func(string, string) error { @@ -239,7 +247,12 @@ func TestPostgresStorePublicationRestoresCompleteRecord(t *testing.T) { return publishErr } } - _, err = store.Save(saveCtx, &cliproxyauth.Auth{ID: id, Storage: &postgresAuthTestStorage{data: []byte(test.candidate)}}) + auth := &cliproxyauth.Auth{ID: id, Storage: &postgresAuthTestStorage{data: []byte(test.candidate)}} + if test.name == "unchanged-metadata" { + auth.Storage = nil + auth.Metadata = map[string]any{"type": "codex", "value": "before"} + } + _, err = store.Save(saveCtx, auth) if test.fail && (!errors.Is(err, publishErr) || strings.Contains(err.Error(), "rollback failed")) || !test.fail && err != nil { t.Fatalf("Save() error = %v", err) } @@ -256,9 +269,14 @@ func TestPostgresStorePublicationRestoresCompleteRecord(t *testing.T) { if test.fail { wantContent = previous } - if err != nil || !jsonEqual([]byte(content), []byte(wantContent)) || !gotCreated.Equal(created) || test.fail && !gotUpdated.Equal(updated) || !test.fail && !gotUpdated.After(updated) { + preserveUpdated := test.fail || strings.HasPrefix(test.name, "unchanged-") + if err != nil || !jsonEqual([]byte(content), []byte(wantContent)) || !gotCreated.Equal(created) || preserveUpdated && !gotUpdated.Equal(updated) || !preserveUpdated && !gotUpdated.After(updated) { t.Fatalf("record = (%s, %v, %v), error = %v", content, gotCreated, gotUpdated, err) } + local, errLocal := os.ReadFile(filepath.Join(store.authDir, id)) + if errLocal != nil || !jsonEqual(local, []byte(wantContent)) { + t.Fatalf("local content = %s, error = %v, want %s", local, errLocal, wantContent) + } listed, errList := store.List(ctx) if errList != nil || len(listed) != 1 || !listed[0].CreatedAt.Equal(gotCreated) || !listed[0].UpdatedAt.Equal(gotUpdated) { t.Fatalf("List() did not preserve record timestamps: %+v, %v", listed, errList) From f8880a27a950dbdf7ea5f600227ba826fb6aeb50 Mon Sep 17 00:00:00 2001 From: Aikins Laryea Date: Sun, 13 Sep 2026 17:10:01 +0000 Subject: [PATCH 5/5] store: serialize auth deletion and watcher reads --- internal/store/postgresstore.go | 48 +++++++++---------- .../store/postgresstore_integration_test.go | 8 ++++ 2 files changed, 30 insertions(+), 26 deletions(-) diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index cee50f7787f..15877ad3ba7 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -439,14 +439,16 @@ func (s *PostgresStore) Delete(ctx context.Context, id string) error { s.mu.Lock() defer s.mu.Unlock() - if err = os.Remove(path); err != nil && !errors.Is(err, fs.ErrNotExist) { - return fmt.Errorf("postgres store: delete auth file: %w", err) - } relID, err := s.relativeAuthID(path) if err != nil { return err } - return s.deleteAuthRecord(ctx, relID) + return s.withAuthLock(ctx, relID, func(conn *sql.Conn) error { + if errRemove := os.Remove(path); errRemove != nil && !errors.Is(errRemove, fs.ErrNotExist) { + return fmt.Errorf("postgres store: delete auth file: %w", errRemove) + } + return s.deleteAuthRecord(ctx, conn, relID) + }) } // PersistAuthFiles stores the provided auth file changes in PostgreSQL. @@ -583,21 +585,17 @@ func (s *PostgresStore) syncAuthFromDatabase(ctx context.Context) error { } func (s *PostgresStore) syncAuthFile(ctx context.Context, relID, path string) error { - data, err := os.ReadFile(path) - if err != nil { - if errors.Is(err, fs.ErrNotExist) { - return s.deleteAuthRecord(ctx, relID) - } - return fmt.Errorf("postgres store: read auth file: %w", err) - } - if len(data) == 0 { - return s.deleteAuthRecord(ctx, relID) - } - return s.persistAuth(ctx, relID, data) -} - -func (s *PostgresStore) persistAuth(ctx context.Context, relID string, data []byte) error { return s.withAuthLock(ctx, relID, func(conn *sql.Conn) error { + data, err := os.ReadFile(path) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + return s.deleteAuthRecord(ctx, conn, relID) + } + return fmt.Errorf("postgres store: read auth file: %w", err) + } + if len(data) == 0 { + return s.deleteAuthRecord(ctx, conn, relID) + } jsonPayload := json.RawMessage(data) query := fmt.Sprintf(` INSERT INTO %s (id, content, created_at, updated_at) @@ -612,14 +610,12 @@ func (s *PostgresStore) persistAuth(ctx context.Context, relID string, data []by }) } -func (s *PostgresStore) deleteAuthRecord(ctx context.Context, relID string) error { - return s.withAuthLock(ctx, relID, func(conn *sql.Conn) error { - query := fmt.Sprintf("DELETE FROM %s WHERE id = $1", s.fullTableName(s.cfg.AuthTable)) - if _, err := conn.ExecContext(ctx, query, relID); err != nil { - return fmt.Errorf("postgres store: delete auth record: %w", err) - } - return nil - }) +func (s *PostgresStore) deleteAuthRecord(ctx context.Context, conn *sql.Conn, relID string) error { + query := fmt.Sprintf("DELETE FROM %s WHERE id = $1", s.fullTableName(s.cfg.AuthTable)) + if _, err := conn.ExecContext(ctx, query, relID); err != nil { + return fmt.Errorf("postgres store: delete auth record: %w", err) + } + return nil } func (s *PostgresStore) withAuthLock(ctx context.Context, relID string, save func(*sql.Conn) error) (err error) { diff --git a/internal/store/postgresstore_integration_test.go b/internal/store/postgresstore_integration_test.go index 51030d811df..0babda3c594 100644 --- a/internal/store/postgresstore_integration_test.go +++ b/internal/store/postgresstore_integration_test.go @@ -38,6 +38,9 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { {"shared-failure", `{"value":"old"}`, `{"value":"first"}`, `{"value":"second"}`, "save"}, {"shared-success", `{"value":"old"}`, `{"value":"first"}`, `{"value":"second"}`, "save"}, {"shared-stale-read", `{"value":"old"}`, `{"value":"first"}`, `{"value":"old"}`, "save"}, + {"shared-delete", `{"value":"old"}`, `{"value":"first"}`, "", "delete"}, + {"shared-observe-existing", `{"value":"old"}`, `{"value":"first"}`, `{"value":"first"}`, "observe"}, + {"shared-observe-missing", "", `{"value":"first"}`, `{"value":"first"}`, "observe"}, } { t.Run(test.name, func(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -103,6 +106,8 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { _, errWrite = b.Save(ctx, &cliproxyauth.Auth{ID: id, Storage: &postgresAuthTestStorage{data: []byte(test.newer)}}) case "delete": errWrite = b.Delete(ctx, id) + case "observe": + errWrite = b.PersistAuthFiles(ctx, "", filepath.Join(b.authDir, id)) case "watcher": path := filepath.Join(b.authDir, id) if errWrite = os.WriteFile(path, []byte(test.newer), 0o600); errWrite == nil { @@ -169,6 +174,9 @@ func TestPostgresStoreConcurrentPublicationFailure(t *testing.T) { } if shared { previous = []byte(test.newer) + if test.operation == "delete" { + previous = nil + } } assertPostgresAuthLocal(t, filepath.Join(a.authDir, id), previous) assertPostgresAuthNoTemp(t, filepath.Join(a.authDir, id))