Fix guest reconnect expiry across database lock waits

docs/design/2026-09-29-guest-reconnect-authorization.md

@@ -37,10 +37,15 @@ 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. 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.
 
 The connection epoch is scoped to an enrolled boot. It is not the placement

internal/controld/pgstore/reconnect.go

@@ -46,6 +46,26 @@ func (r pgGuestReconnects) lockScope(ctx context.Context, b control.GuestReconne
 	return nil
 }
 
+// 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) {
+		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 +74,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 +101,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
@@ -112,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, &current)

internal/controld/pgstore/reconnect_test.go

@@ -164,3 +164,121 @@ 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 TestGuestReconnectExpiryAfterBootstrapLockWait(t *testing.T) {
+	dsn := startPostgres(t)
+	for _, operation := range []string{"enroll", "begin", "consume"} {
+		t.Run(operation, func(t *testing.T) {
+			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")
+					}
+				})
+			}
+		})
+	}
+}