From 8cf49147db9cabd9c7b6bee1a68ba9a6b37a8b9b Mon Sep 17 00:00:00 2001 From: jiashuoz Date: Wed, 30 Sep 2026 12:17:11 +0800 Subject: [PATCH 1/2] fix: enforce reconnect expiry after bootstrap row lock waits --- ...026-09-29-guest-reconnect-authorization.md | 10 ++- internal/controld/pgstore/reconnect.go | 21 +++++ internal/controld/pgstore/reconnect_test.go | 90 +++++++++++++++++++ 3 files changed, 117 insertions(+), 4 deletions(-) diff --git a/docs/design/2026-09-29-guest-reconnect-authorization.md b/docs/design/2026-09-29-guest-reconnect-authorization.md index 60bdaaa..8de936d 100644 --- a/docs/design/2026-09-29-guest-reconnect-authorization.md +++ b/docs/design/2026-09-29-guest-reconnect-authorization.md @@ -37,10 +37,12 @@ the connection epoch, consumes the attempt and installs a new bootstrap hash in one transaction. Concurrent consumers produce one success. A lost response cannot be replayed; the caller must start a new attempt. -The PostgreSQL adapter holds runner and session row locks until commit, checks -connected/live state and uses database time after acquiring authorization locks -for expiry. The caller's clock cannot extend a five-second challenge. Database -clock skew can cause a closed refusal; it must never extend validity. No failed +The PostgreSQL adapter holds runner, session and bootstrap row locks until commit, +checks connected/live state and uses database time after acquiring those locks +for expiry. Enrollment and challenge issuance explicitly lock the bootstrap row +before evaluating expiry: an UPDATE predicate alone can be evaluated before a +lock-only transaction releases that row. The caller's clock cannot extend a +five-second challenge. Database clock skew can cause a closed refusal; it must never extend validity. No failed operation returns a token or an epoch. Database errors are fixed sentinels. The connection epoch is scoped to an enrolled boot. It is not the placement diff --git a/internal/controld/pgstore/reconnect.go b/internal/controld/pgstore/reconnect.go index da28b7c..9928d7d 100644 --- a/internal/controld/pgstore/reconnect.go +++ b/internal/controld/pgstore/reconnect.go @@ -46,6 +46,21 @@ func (r pgGuestReconnects) lockScope(ctx context.Context, b control.GuestReconne return nil } +// lockBootstrap must precede expiry evaluation. UPDATE predicates can be +// evaluated before a lock-only transaction releases the tuple, so a timestamp +// predicate in that UPDATE alone does not enforce expiry after a lock wait. +func (r pgGuestReconnects) lockBootstrap(ctx context.Context, b control.GuestReconnectScope) error { + var found int + err := r.s.q(ctx).QueryRow(ctx, `SELECT 1 FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2 FOR UPDATE`, string(b.WorkspaceID), string(b.SessionID)).Scan(&found) + if errors.Is(err, pgx.ErrNoRows) { + return control.ErrReconnectInvalid + } + if err != nil { + return control.ErrUnavailable + } + return nil +} + func (r pgGuestReconnects) EnrollGuest(ctx context.Context, b control.GuestReconnectScope, hash string, id control.GuestReconnectIdentity, now time.Time) error { if hash == "" || len(id.BootEpoch) == 0 || len(id.BootEpoch) > 256 || len(id.PublicKey) == 0 || len(id.PublicKey) > 256 || now.IsZero() { return control.ErrReconnectInvalid @@ -54,6 +69,9 @@ func (r pgGuestReconnects) EnrollGuest(ctx context.Context, b control.GuestRecon if err := r.lockScope(ctx, b); err != nil { return err } + if err := r.lockBootstrap(ctx, b); err != nil { + return err + } ct, err := r.s.q(ctx).Exec(ctx, `UPDATE session_bootstraps SET consumed_at=$1,guest_boot_epoch=$2,guest_public_key=$3 WHERE workspace_id=$4 AND session_id=$5 AND token_hash=$6 AND placement_generation=$7 AND expires_at>$1 AND expires_at>clock_timestamp() AND consumed_at IS NULL AND guest_public_key IS NULL`, now, id.BootEpoch, id.PublicKey, string(b.WorkspaceID), string(b.SessionID), hash, int64(b.PlacementGeneration)) if err != nil { return control.ErrUnavailable @@ -78,6 +96,9 @@ func (r pgGuestReconnects) BeginGuestReconnect(ctx context.Context, b control.Gu if err := r.lockScope(ctx, b); err != nil { return err } + if err := r.lockBootstrap(ctx, b); err != nil { + return err + } var databaseNow time.Time if err := r.s.q(ctx).QueryRow(ctx, `SELECT clock_timestamp()`).Scan(&databaseNow); err != nil { return control.ErrUnavailable diff --git a/internal/controld/pgstore/reconnect_test.go b/internal/controld/pgstore/reconnect_test.go index 626ca37..ec4608f 100644 --- a/internal/controld/pgstore/reconnect_test.go +++ b/internal/controld/pgstore/reconnect_test.go @@ -164,3 +164,93 @@ func TestGuestReconnectExpiresWhileWaitingForAuthorityLock(t *testing.T) { t.Fatal("expired request changed capability or epoch") } } + +// Expiry must be checked after waiting on the capability row itself, not just +// after runner/session locks. A lock-only transaction need not change its tuple. +func TestGuestReconnectExpiryAfterBootstrapRowLock(t *testing.T) { + dsn := startPostgres(t) + for _, operation := range []string{"enroll", "begin"} { + t.Run(operation, func(t *testing.T) { + st := freshStore(t, dsn, t.Name()) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + b := control.GuestReconnectScope{WorkspaceID: "ws_self_hosted", PoolID: "pool_self_hosted", SessionID: "expiry.test", RunnerID: "runner.test", PlacementGeneration: 1, ConnectionGeneration: 1} + if _, err := st.Sessions().CreateSession(ctx, "ws_self_hosted", control.Session{ID: b.SessionID, CreatorID: "actor.test", PoolID: "pool_self_hosted", RunnerID: b.RunnerID, State: control.StateRunning}); err != nil { + t.Fatal(err) + } + if err := st.Fleet().UpsertRunner(ctx, "pool_self_hosted", control.Runner{ID: b.RunnerID, Generation: 1, Connected: true}); err != nil { + t.Fatal(err) + } + now := time.Now().UTC().Truncate(time.Microsecond) + token := control.SessionBootstrap{Hash: "hash.test", PlacementGeneration: 1, ExpiresAt: now.Add(2 * time.Second)} + if err := st.Bootstraps().PutSessionBootstrap(ctx, "ws_self_hosted", b.SessionID, token); err != nil { + t.Fatal(err) + } + id := control.GuestReconnectIdentity{BootEpoch: "boot.test", PublicKey: "key.test"} + r := st.GuestReconnects() + if operation == "begin" { + if err := r.EnrollGuest(ctx, b, token.Hash, id, now); err != nil { + t.Fatal(err) + } + } + tx, err := st.pool.Begin(ctx) + if err != nil { + t.Fatal(err) + } + defer tx.Rollback(context.Background()) + var holder int + if err := tx.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&holder); err != nil { + t.Fatal(err) + } + if _, err := tx.Exec(ctx, `SELECT 1 FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2 FOR UPDATE`, string(b.WorkspaceID), string(b.SessionID)); err != nil { + t.Fatal(err) + } + done := make(chan error, 1) + go func() { + if operation == "enroll" { + done <- r.EnrollGuest(ctx, b, token.Hash, id, now) + return + } + _, err := r.BeginGuestReconnect(ctx, b, control.GuestReconnectAttempt{ID: "attempt.test", Challenge: "nonce.test", ExpiresAt: token.ExpiresAt}, now) + done <- err + }() + deadline := time.Now().Add(3 * time.Second) + for { + var waiting bool + if err := st.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE datname=current_database() AND $1=ANY(pg_blocking_pids(pid)))`, holder).Scan(&waiting); err != nil { + t.Fatal(err) + } + if waiting { + break + } + if time.Now().After(deadline) { + t.Fatal("request did not wait on bootstrap row") + } + time.Sleep(5 * time.Millisecond) + } + if wait := time.Until(token.ExpiresAt.Add(20 * time.Millisecond)); wait > 0 { + time.Sleep(wait) + } + if err := tx.Commit(ctx); err != nil { + t.Fatal(err) + } + err = <-done + want := control.ErrReconnectInvalid + if operation == "begin" { + want = control.ErrReconnectExpired + } + if !errors.Is(err, want) { + t.Fatalf("expired %s: %v", operation, err) + } + var enrolled, pending, consumed bool + var hash string + var epoch int64 + if err := st.pool.QueryRow(ctx, `SELECT guest_public_key IS NOT NULL,reconnect_attempt IS NOT NULL,consumed_at IS NOT NULL,token_hash,guest_connection_epoch FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2`, string(b.WorkspaceID), string(b.SessionID)).Scan(&enrolled, &pending, &consumed, &hash, &epoch); err != nil { + t.Fatal(err) + } + if enrolled != (operation == "begin") || consumed != (operation == "begin") || pending || hash != token.Hash || epoch != 0 { + t.Fatal("expired operation changed authorization") + } + }) + } +} From a8fc41517ef1d1a2a2e4d29c43ed0ae79c36eea5 Mon Sep 17 00:00:00 2001 From: jiashuoz Date: Wed, 30 Sep 2026 12:25:24 +0800 Subject: [PATCH 2/2] fix: acquire bootstrap table write lock before reconnect expiry checks --- ...026-09-29-guest-reconnect-authorization.md | 9 +- internal/controld/pgstore/reconnect.go | 10 +- internal/controld/pgstore/reconnect_test.go | 190 ++++++++++-------- 3 files changed, 124 insertions(+), 85 deletions(-) diff --git a/docs/design/2026-09-29-guest-reconnect-authorization.md b/docs/design/2026-09-29-guest-reconnect-authorization.md index 8de936d..8fa952d 100644 --- a/docs/design/2026-09-29-guest-reconnect-authorization.md +++ b/docs/design/2026-09-29-guest-reconnect-authorization.md @@ -39,9 +39,12 @@ be replayed; the caller must start a new attempt. The PostgreSQL adapter holds runner, session and bootstrap row locks until commit, checks connected/live state and uses database time after acquiring those locks -for expiry. Enrollment and challenge issuance explicitly lock the bootstrap row -before evaluating expiry: an UPDATE predicate alone can be evaluated before a -lock-only transaction releases that row. The caller's clock cannot extend a +for expiry. All three mutations also acquire ROW EXCLUSIVE on the bootstrap +table before evaluating expiry, so a maintenance SHARE lock cannot delay the +subsequent UPDATE past that check. ROW EXCLUSIVE remains compatible with other +ordinary writers. Enrollment and challenge issuance explicitly lock the bootstrap +row: an UPDATE predicate alone can be evaluated before a lock-only transaction +releases that row. The caller's clock cannot extend a five-second challenge. Database clock skew can cause a closed refusal; it must never extend validity. No failed operation returns a token or an epoch. Database errors are fixed sentinels. diff --git a/internal/controld/pgstore/reconnect.go b/internal/controld/pgstore/reconnect.go index 9928d7d..f32ab00 100644 --- a/internal/controld/pgstore/reconnect.go +++ b/internal/controld/pgstore/reconnect.go @@ -46,10 +46,15 @@ func (r pgGuestReconnects) lockScope(ctx context.Context, b control.GuestReconne return nil } -// lockBootstrap must precede expiry evaluation. UPDATE predicates can be +// lockBootstrap obtains the table write lock and row lock before expiry evaluation. UPDATE predicates can be // evaluated before a lock-only transaction releases the tuple, so a timestamp // predicate in that UPDATE alone does not enforce expiry after a lock wait. func (r pgGuestReconnects) lockBootstrap(ctx context.Context, b control.GuestReconnectScope) error { + // SELECT FOR UPDATE takes only ROW SHARE on the relation. A maintenance + // SHARE lock would still make the later UPDATE wait past its expiry check. + if _, err := r.s.q(ctx).Exec(ctx, `LOCK TABLE session_bootstraps IN ROW EXCLUSIVE MODE`); err != nil { + return control.ErrUnavailable + } var found int err := r.s.q(ctx).QueryRow(ctx, `SELECT 1 FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2 FOR UPDATE`, string(b.WorkspaceID), string(b.SessionID)).Scan(&found) if errors.Is(err, pgx.ErrNoRows) { @@ -133,6 +138,9 @@ func (r pgGuestReconnects) ConsumeGuestReconnect(ctx context.Context, b control. if err := r.lockScope(ctx, b); err != nil { return err } + if err := r.lockBootstrap(ctx, b); err != nil { + return err + } var expiry time.Time var current int64 err := r.s.q(ctx).QueryRow(ctx, `SELECT reconnect_expires_at,guest_connection_epoch FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2 AND placement_generation=$3 AND guest_boot_epoch=$4 AND guest_public_key=$5 AND reconnect_attempt=$6 AND reconnect_challenge=$7 AND reconnect_runner_generation=$8 FOR UPDATE`, string(b.WorkspaceID), string(b.SessionID), int64(b.PlacementGeneration), id.BootEpoch, id.PublicKey, a.ID, a.Challenge, int64(b.ConnectionGeneration)).Scan(&expiry, ¤t) diff --git a/internal/controld/pgstore/reconnect_test.go b/internal/controld/pgstore/reconnect_test.go index ec4608f..0e102ac 100644 --- a/internal/controld/pgstore/reconnect_test.go +++ b/internal/controld/pgstore/reconnect_test.go @@ -167,89 +167,117 @@ func TestGuestReconnectExpiresWhileWaitingForAuthorityLock(t *testing.T) { // Expiry must be checked after waiting on the capability row itself, not just // after runner/session locks. A lock-only transaction need not change its tuple. -func TestGuestReconnectExpiryAfterBootstrapRowLock(t *testing.T) { +func TestGuestReconnectExpiryAfterBootstrapLockWait(t *testing.T) { dsn := startPostgres(t) - for _, operation := range []string{"enroll", "begin"} { + for _, operation := range []string{"enroll", "begin", "consume"} { t.Run(operation, func(t *testing.T) { - st := freshStore(t, dsn, t.Name()) - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - b := control.GuestReconnectScope{WorkspaceID: "ws_self_hosted", PoolID: "pool_self_hosted", SessionID: "expiry.test", RunnerID: "runner.test", PlacementGeneration: 1, ConnectionGeneration: 1} - if _, err := st.Sessions().CreateSession(ctx, "ws_self_hosted", control.Session{ID: b.SessionID, CreatorID: "actor.test", PoolID: "pool_self_hosted", RunnerID: b.RunnerID, State: control.StateRunning}); err != nil { - t.Fatal(err) - } - if err := st.Fleet().UpsertRunner(ctx, "pool_self_hosted", control.Runner{ID: b.RunnerID, Generation: 1, Connected: true}); err != nil { - t.Fatal(err) - } - now := time.Now().UTC().Truncate(time.Microsecond) - token := control.SessionBootstrap{Hash: "hash.test", PlacementGeneration: 1, ExpiresAt: now.Add(2 * time.Second)} - if err := st.Bootstraps().PutSessionBootstrap(ctx, "ws_self_hosted", b.SessionID, token); err != nil { - t.Fatal(err) - } - id := control.GuestReconnectIdentity{BootEpoch: "boot.test", PublicKey: "key.test"} - r := st.GuestReconnects() - if operation == "begin" { - if err := r.EnrollGuest(ctx, b, token.Hash, id, now); err != nil { - t.Fatal(err) - } - } - tx, err := st.pool.Begin(ctx) - if err != nil { - t.Fatal(err) - } - defer tx.Rollback(context.Background()) - var holder int - if err := tx.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&holder); err != nil { - t.Fatal(err) - } - if _, err := tx.Exec(ctx, `SELECT 1 FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2 FOR UPDATE`, string(b.WorkspaceID), string(b.SessionID)); err != nil { - t.Fatal(err) - } - done := make(chan error, 1) - go func() { - if operation == "enroll" { - done <- r.EnrollGuest(ctx, b, token.Hash, id, now) - return - } - _, err := r.BeginGuestReconnect(ctx, b, control.GuestReconnectAttempt{ID: "attempt.test", Challenge: "nonce.test", ExpiresAt: token.ExpiresAt}, now) - done <- err - }() - deadline := time.Now().Add(3 * time.Second) - for { - var waiting bool - if err := st.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE datname=current_database() AND $1=ANY(pg_blocking_pids(pid)))`, holder).Scan(&waiting); err != nil { - t.Fatal(err) - } - if waiting { - break - } - if time.Now().After(deadline) { - t.Fatal("request did not wait on bootstrap row") - } - time.Sleep(5 * time.Millisecond) - } - if wait := time.Until(token.ExpiresAt.Add(20 * time.Millisecond)); wait > 0 { - time.Sleep(wait) - } - if err := tx.Commit(ctx); err != nil { - t.Fatal(err) - } - err = <-done - want := control.ErrReconnectInvalid - if operation == "begin" { - want = control.ErrReconnectExpired - } - if !errors.Is(err, want) { - t.Fatalf("expired %s: %v", operation, err) - } - var enrolled, pending, consumed bool - var hash string - var epoch int64 - if err := st.pool.QueryRow(ctx, `SELECT guest_public_key IS NOT NULL,reconnect_attempt IS NOT NULL,consumed_at IS NOT NULL,token_hash,guest_connection_epoch FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2`, string(b.WorkspaceID), string(b.SessionID)).Scan(&enrolled, &pending, &consumed, &hash, &epoch); err != nil { - t.Fatal(err) - } - if enrolled != (operation == "begin") || consumed != (operation == "begin") || pending || hash != token.Hash || epoch != 0 { - t.Fatal("expired operation changed authorization") + for _, lockKind := range []string{"row", "table"} { + t.Run(lockKind, func(t *testing.T) { + st := freshStore(t, dsn, t.Name()) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + b := control.GuestReconnectScope{WorkspaceID: "ws_self_hosted", PoolID: "pool_self_hosted", SessionID: "expiry.test", RunnerID: "runner.test", PlacementGeneration: 1, ConnectionGeneration: 1} + if _, err := st.Sessions().CreateSession(ctx, "ws_self_hosted", control.Session{ID: b.SessionID, CreatorID: "actor.test", PoolID: "pool_self_hosted", RunnerID: b.RunnerID, State: control.StateRunning}); err != nil { + t.Fatal(err) + } + if err := st.Fleet().UpsertRunner(ctx, "pool_self_hosted", control.Runner{ID: b.RunnerID, Generation: 1, Connected: true}); err != nil { + t.Fatal(err) + } + now := time.Now().UTC().Truncate(time.Microsecond) + token := control.SessionBootstrap{Hash: "hash.test", PlacementGeneration: 1, ExpiresAt: now.Add(2 * time.Second)} + if err := st.Bootstraps().PutSessionBootstrap(ctx, "ws_self_hosted", b.SessionID, token); err != nil { + t.Fatal(err) + } + id := control.GuestReconnectIdentity{BootEpoch: "boot.test", PublicKey: "key.test"} + r := st.GuestReconnects() + if operation != "enroll" { + if err := r.EnrollGuest(ctx, b, token.Hash, id, now); err != nil { + t.Fatal(err) + } + } + attempt := control.GuestReconnectAttempt{ID: "attempt.test", Challenge: "nonce.test", ExpiresAt: token.ExpiresAt} + if operation == "consume" { + if _, err := r.BeginGuestReconnect(ctx, b, attempt, now); err != nil { + t.Fatal(err) + } + } + + tx, err := st.pool.Begin(ctx) + if err != nil { + t.Fatal(err) + } + defer tx.Rollback(context.Background()) + var holder int + if err := tx.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&holder); err != nil { + t.Fatal(err) + } + if lockKind == "table" { + if _, err := tx.Exec(ctx, `LOCK TABLE session_bootstraps IN SHARE MODE`); err != nil { + t.Fatal(err) + } + } else { + if _, err := tx.Exec(ctx, `SELECT 1 FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2 FOR UPDATE`, string(b.WorkspaceID), string(b.SessionID)); err != nil { + t.Fatal(err) + } + } + done := make(chan error, 1) + go func() { + if operation == "enroll" { + done <- r.EnrollGuest(ctx, b, token.Hash, id, now) + return + } + if operation == "begin" { + _, err := r.BeginGuestReconnect(ctx, b, attempt, now) + done <- err + return + } + fresh := token + fresh.Hash = "fresh-hash.test" + fresh.ExpiresAt = now.Add(time.Minute) + epoch, err := r.ConsumeGuestReconnect(ctx, b, id, attempt, fresh, now) + if epoch != 0 { + err = fmt.Errorf("unexpected epoch %d", epoch) + } + done <- err + }() + deadline := time.Now().Add(3 * time.Second) + for { + var waiting bool + if err := st.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM pg_stat_activity WHERE datname=current_database() AND $1=ANY(pg_blocking_pids(pid)))`, holder).Scan(&waiting); err != nil { + t.Fatal(err) + } + if waiting { + break + } + if time.Now().After(deadline) { + t.Fatal("request did not wait on bootstrap row") + } + time.Sleep(5 * time.Millisecond) + } + if wait := time.Until(token.ExpiresAt.Add(20 * time.Millisecond)); wait > 0 { + time.Sleep(wait) + } + if err := tx.Commit(ctx); err != nil { + t.Fatal(err) + } + err = <-done + want := control.ErrReconnectInvalid + if operation != "enroll" { + want = control.ErrReconnectExpired + } + if !errors.Is(err, want) { + t.Fatalf("expired %s: %v", operation, err) + } + var enrolled, pending, consumed bool + var hash string + var epoch int64 + if err := st.pool.QueryRow(ctx, `SELECT guest_public_key IS NOT NULL,reconnect_attempt IS NOT NULL,consumed_at IS NOT NULL,token_hash,guest_connection_epoch FROM session_bootstraps WHERE workspace_id=$1 AND session_id=$2`, string(b.WorkspaceID), string(b.SessionID)).Scan(&enrolled, &pending, &consumed, &hash, &epoch); err != nil { + t.Fatal(err) + } + if enrolled != (operation != "enroll") || consumed != (operation != "enroll") || pending != (operation == "consume") || hash != token.Hash || epoch != 0 { + t.Fatal("expired operation changed authorization") + } + }) } }) }