Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 866e676

Browse files
jscottmillerjscottmiller
authored
feat: invalidate provisioner daemon sessions on key deletion (#26532)
## Summary Closes PLAT-305. When a provisioner key is deleted, the associated daemon kept operating on its existing WebSocket connection, because authentication was only checked at connection establishment and deletion was a bare `DELETE` with no session invalidation. This adds four layers of defense so a deleted key promptly stops doing work: 1. **Publish on delete.** `deleteProvisionerKey` publishes to a new per-key pubsub channel (`coderd/pubsub.ProvisionerKeyDeletedChannel`) after a successful delete. Publish errors are logged but still return `204`, since layer 3 is the durable backstop. 2. **Subscribe and tear down.** The daemon serve handler subscribes to its key's channel and terminates the DRPC session on a deletion event. Termination is deferred while a job claimed by the session is active: the daemon may finish and report the in-flight job (`UpdateJob`/`CompleteJob` have no key check), and the last active job's completion performs the cancellation. Because Postgres `LISTEN`/`NOTIFY` does not buffer for non-listeners, the handler also performs a synchronous key-existence re-check immediately after subscribing to close the race between auth and subscription. The subscription uses `SubscribeWithErr` so that an `ErrDroppedMessages` signal (emitted when the pubsub listener reconnects) triggers the same key re-check, closing the listener-outage window in which a deletion notification could be missed. 3. **Backstop on acquire.** `AcquireJob` and `AcquireJobWithCancel` verify the key still exists before waiting for a job, and the `Acquirer` claims jobs in a transaction that first locks the worker's deletable key (`LockProvisionerKeyByIDForShare`, a `FOR KEY SHARE` row lock held until commit) before running the `AcquireProvisionerJob` claim, so a claim cannot commit after the key's deletion. This guards against a missed pubsub message. A missing key row surfaces as its own result rather than overloading the claim query's no-rows response: the acquire terminates with `ErrProvisionerKeyDeleted` (terminating the session, with the same active-job deferral) and hands the consumed wakeup to another waiting daemon in the same domain, rather than silently re-parking and starving peers of job postings. 4. **Heartbeat watchdog.** The per-session heartbeat loop (1m interval) also re-checks the key, so even a session whose deletion notification was silently lost terminates within one heartbeat interval instead of living until the connection breaks (same active-job deferral as layer 2). Reserved keys skip the check. A job that is claimed but never delivered (the session or connection dies between the database claim and the stream send) is marked failed immediately on a fresh context, instead of staying assigned to the worker until the job reaper. Reserved keys (built-in, user-auth, PSK) are exempt throughout, since they are not deletable rows. The acquire-time lookup runs as `dbauthz.AsSystemReadProvisionerDaemons`, because the provisionerd role cannot read provisioner keys and a provisioner key's RBAC object is a provisioner daemon. A single key can back many daemons (and span HA replicas), so the per-key channel fans out to invalidate all of them at once. Per-key channels keep the `LISTEN` count proportional to distinct keys rather than waking every daemon on unrelated deletions. ### Known limitations - **`UpdateJob`/`CompleteJob` intentionally have no key check.** By the time those RPCs arrive the work has already run; rejecting completion would strand a build in "running" (until the job reaper fails it) with real infrastructure left unreconciled. Session termination is deferred while a job is active so the completion can be reported; the daemon may not receive the final RPC response when the deferred termination fires, but the job's outcome is already persisted. - **After termination, the daemon process redials and receives 401s until restarted.** The dial-time exit logic only triggers on 403, and the auth middleware returns 401 for an invalid key; this dial behavior predates this PR and is tracked as a follow-up in [PLAT-452](https://linear.app/codercom/issue/PLAT-452) (return 403 for invalid provisioner keys). ## Tests - `coderd/provisionerdserver`: `TestAcquireJob_ProvisionerKeyDeleted` (both RPC variants), `TestAcquireJob_ReservedProvisionerKey`, `TestHeartbeat_ProvisionerKeyDeleted` (heartbeat watchdog cancels the session after key deletion), `TestAcquirer_ProvisionerKeyDeleted` (a dead-key acquiree exits terminally and its clearance is promoted to a peer in the same domain), and `TestTerminateSession_Deferral` (termination is immediate when idle and deferred until the last active job finishes). - `coderd/database`: `TestAcquireProvisionerJob/ProvisionerKeyLock` covers the lock query against real Postgres: it returns the key ID while the row exists and no rows once it is deleted. The lock-then-claim composition is pinned by `TestAcquirer_ProvisionerKeyDeleted`. - `enterprise/coderd`: `TestProvisionerDaemonServe/KeyDeletionClosesSession` asserts an active session closes after its key is deleted. `KeyDeletedDuringSetupClosesSession` covers the post-subscribe re-check when a key is deleted between auth and subscription, and `DroppedMessageClosesSession` covers the `ErrDroppedMessages` re-check when a deletion is missed during a listener outage. ## Validation - `make` pre-commit (gen/fmt/lint/build) passed via git hooks. - Targeted tests pass; existing acquire tests pass with no regression. - Manual: brought up a dev deployment (coder-in-coder) with a Premium license, created a deletable provisioner key, and started an external daemon with `coder provisionerd start`. Confirmed it authenticated via the key and connected, appearing as `idle` in both `coder provisioner list` (with the key name) and the organization Provisioners UI. - Manual, idle teardown: deleted the key while the daemon was idle. The server logged `provisioner key deleted, terminating session`, the daemon's session closed immediately, and it dropped from `coder provisioner list` (then entered the known 401 redial loop, PLAT-452). - Manual, deferred termination: ran a workspace build (tagged template, `sleep 45` in `local-exec`) pinned to the external daemon and deleted the key mid-build. The server logged `deferring session cancellation until active jobs finish`; the heartbeat watchdog re-checked mid-build and re-deferred rather than force-killing. The build ran to completion (`Apply complete`, workspace `Started`) and only then did `canceling session after job completion` fire. The documented caveat reproduced: the daemon lost the final `CompleteJob` ack, and the build outcome was still persisted correctly. <details> <summary>Implementation plan and design decisions</summary> ### Design - **Per-key vs global channel:** chose per-key (`provisioner_key_deleted:<keyID>`) so daemons do not wake on unrelated deletions. The cost is one `LISTEN` per distinct key per replica on the shared listener connection, which is negligible against Coder's existing channels. - **Missing-key behavior on acquire:** returns an error that tears down the acquire rather than silently returning an empty job. - **Subscribe-startup race:** ordering is `authorize -> UpsertProvisionerDaemon -> Subscribe -> GetProvisionerKeyByID`. The post-subscribe re-check handles a deletion that committed before the `LISTEN` registered (Postgres does not buffer notifications for non-listeners; the in-process buffer only smooths bursts and drops on overflow). - **`NewServer` change:** `KeyID` was added to `provisionerdserver.Options` to avoid a positional signature change across call sites. The in-memory (built-in) daemon leaves it unset and is therefore exempt. ### Files - `coderd/pubsub/provisionerkeydeleted.go` (new) — channel helper. - `enterprise/coderd/provisionerkeys.go` — publish on delete. - `enterprise/coderd/provisionerdaemons.go` — subscribe, re-check, cancel session; pass `KeyID`. - `coderd/provisionerdserver/provisionerdserver.go` — `KeyID` option and acquire-time existence check. </details> --- This pull request was created by Coder Agents on behalf of @jscottmiller.
1 parent e562912 commit 866e676

18 files changed

Lines changed: 1317 additions & 62 deletions

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6921,6 +6921,19 @@ func (q *querier) LockChatByID(ctx context.Context, id uuid.UUID) (uuid.UUID, er
69216921
return q.db.LockChatByID(ctx, id)
69226922
}
69236923

6924+
func (q *querier) LockProvisionerKeyByIDForShare(ctx context.Context, id uuid.UUID) (uuid.UUID, error) {
6925+
// The lock query returns only the key ID, so fetch the key to authorize
6926+
// the read against its RBAC object.
6927+
key, err := q.db.GetProvisionerKeyByID(ctx, id)
6928+
if err != nil {
6929+
return uuid.Nil, err
6930+
}
6931+
if err := q.authorizeContext(ctx, policy.ActionRead, key); err != nil {
6932+
return uuid.Nil, err
6933+
}
6934+
return q.db.LockProvisionerKeyByIDForShare(ctx, id)
6935+
}
6936+
69246937
func (q *querier) MarkAllInboxNotificationsAsRead(ctx context.Context, arg database.MarkAllInboxNotificationsAsReadParams) error {
69256938
resource := rbac.ResourceInboxNotification.WithOwner(arg.UserID.String())
69266939

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4488,6 +4488,13 @@ func (s *MethodTestSuite) TestProvisionerKeys() {
44884488
dbm.EXPECT().GetProvisionerKeyByID(gomock.Any(), pk.ID).Return(pk, nil).AnyTimes()
44894489
check.Args(pk.ID).Asserts(pk, policy.ActionRead).Returns(pk)
44904490
}))
4491+
s.Run("LockProvisionerKeyByIDForShare", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
4492+
org := testutil.Fake(s.T(), faker, database.Organization{})
4493+
pk := testutil.Fake(s.T(), faker, database.ProvisionerKey{OrganizationID: org.ID})
4494+
dbm.EXPECT().GetProvisionerKeyByID(gomock.Any(), pk.ID).Return(pk, nil).AnyTimes()
4495+
dbm.EXPECT().LockProvisionerKeyByIDForShare(gomock.Any(), pk.ID).Return(pk.ID, nil).AnyTimes()
4496+
check.Args(pk.ID).Asserts(pk, policy.ActionRead).Returns(pk.ID)
4497+
}))
44914498
s.Run("GetProvisionerKeyByHashedSecret", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
44924499
org := testutil.Fake(s.T(), faker, database.Organization{})
44934500
pk := testutil.Fake(s.T(), faker, database.ProvisionerKey{OrganizationID: org.ID, HashedSecret: []byte("foo")})

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 8 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go‎

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier_test.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2613,6 +2613,27 @@ func TestAcquireProvisionerJob(t *testing.T) {
26132613
})
26142614
require.ErrorIs(t, err, sql.ErrNoRows)
26152615
})
2616+
2617+
t.Run("ProvisionerKeyLock", func(t *testing.T) {
2618+
t.Parallel()
2619+
var (
2620+
db, _ = dbtestutil.NewDB(t)
2621+
ctx = testutil.Context(t, testutil.WaitMedium)
2622+
org = dbgen.Organization(t, db, database.Organization{})
2623+
key = dbgen.ProvisionerKey(t, db, database.ProvisionerKey{OrganizationID: org.ID})
2624+
)
2625+
2626+
// While the key exists, the lock returns its ID.
2627+
id, err := db.LockProvisionerKeyByIDForShare(ctx, key.ID)
2628+
require.NoError(t, err)
2629+
require.Equal(t, key.ID, id)
2630+
2631+
// Once the key is deleted, the lock reports no rows.
2632+
err = db.DeleteProvisionerKey(ctx, key.ID)
2633+
require.NoError(t, err)
2634+
_, err = db.LockProvisionerKeyByIDForShare(ctx, key.ID)
2635+
require.ErrorIs(t, err, sql.ErrNoRows)
2636+
})
26162637
}
26172638

26182639
func TestUserLastSeenFilter(t *testing.T) {

‎coderd/database/queries.sql.go‎

Lines changed: 21 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/provisionerkeys.sql‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,19 @@ FROM
1919
WHERE
2020
id = $1;
2121

22+
-- name: LockProvisionerKeyByIDForShare :one
23+
-- Locks the provisioner key row with FOR KEY SHARE for the remainder of the
24+
-- current transaction. FOR KEY SHARE conflicts with DELETE, so while the lock
25+
-- is held the key cannot be deleted, and a committed deletion is observed as
26+
-- no rows by later calls.
27+
SELECT
28+
id
29+
FROM
30+
provisioner_keys
31+
WHERE
32+
id = $1
33+
FOR KEY SHARE;
34+
2235
-- name: GetProvisionerKeyByHashedSecret :one
2336
SELECT
2437
*

‎coderd/provisionerdserver/acquirer.go‎

Lines changed: 58 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,11 @@ import (
1515

1616
"cdr.dev/slog/v3"
1717
"github.com/coder/coder/v2/coderd/database"
18+
"github.com/coder/coder/v2/coderd/database/dbauthz"
1819
"github.com/coder/coder/v2/coderd/database/dbtime"
1920
"github.com/coder/coder/v2/coderd/database/provisionerjobs"
2021
"github.com/coder/coder/v2/coderd/database/pubsub"
22+
"github.com/coder/coder/v2/codersdk"
2123
"github.com/coder/quartz"
2224
)
2325

@@ -61,9 +63,12 @@ func WithClock(clock quartz.Clock) AcquirerOption {
6163
}
6264
}
6365

64-
// AcquirerStore is the subset of database.Store that the Acquirer needs
66+
// AcquirerStore is the subset of database.Store that the Acquirer needs. Job
67+
// acquisition runs in a transaction that locks the worker's deletable
68+
// provisioner key (LockProvisionerKeyByIDForShare) before claiming a job
69+
// (AcquireProvisionerJob), so a claim cannot commit after the key's deletion.
6570
type AcquirerStore interface {
66-
AcquireProvisionerJob(context.Context, database.AcquireProvisionerJobParams) (database.ProvisionerJob, error)
71+
InTx(func(database.Store) error, *database.TxOptions) error
6772
}
6873

6974
func NewAcquirer(ctx context.Context, logger slog.Logger, store AcquirerStore, ps pubsub.Pubsub,
@@ -88,11 +93,15 @@ func NewAcquirer(ctx context.Context, logger slog.Logger, store AcquirerStore, p
8893
// tags from the database. The call blocks until a job is acquired, the context is
8994
// done, or the database returns an error _other_ than that no jobs are available.
9095
// If no jobs are available, this method handles retrying as appropriate.
96+
// When keyID is a deletable provisioner key, the claim only succeeds while
97+
// that key row still exists. Reserved keys and the zero value are not
98+
// checked, as they have no row to delete.
9199
func (a *Acquirer) AcquireJob(
92-
ctx context.Context, organization uuid.UUID, worker uuid.UUID, pt []database.ProvisionerType, tags Tags,
100+
ctx context.Context, organization uuid.UUID, worker uuid.UUID, pt []database.ProvisionerType, tags Tags, keyID uuid.UUID,
93101
) (
94102
retJob database.ProvisionerJob, retErr error,
95103
) {
104+
deletableKey := codersdk.IsDeletableProvisionerKey(keyID)
96105
logger := a.logger.With(
97106
slog.F("organization_id", organization),
98107
slog.F("worker_id", worker),
@@ -120,19 +129,52 @@ func (a *Acquirer) AcquireJob(
120129
return database.ProvisionerJob{}, err
121130
case <-clearance:
122131
logger.Debug(ctx, "got clearance to call database")
123-
job, err := a.store.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
124-
OrganizationID: organization,
125-
StartedAt: sql.NullTime{
126-
Time: dbtime.Now(),
127-
Valid: true,
128-
},
129-
WorkerID: uuid.NullUUID{
130-
UUID: worker,
131-
Valid: true,
132-
},
133-
Types: pt,
134-
ProvisionerTags: dbTags,
135-
})
132+
var job database.ProvisionerJob
133+
err := a.store.InTx(func(tx database.Store) error {
134+
if deletableKey {
135+
// Lock the key for the rest of the transaction so the claim
136+
// below cannot commit after the key's deletion. A missing row
137+
// means the key was deleted.
138+
_, err := tx.LockProvisionerKeyByIDForShare(
139+
//nolint:gocritic // The acquire context has no actor that can
140+
// read provisioner keys, so scope the read to this narrow subject.
141+
dbauthz.AsSystemReadProvisionerDaemons(ctx), keyID)
142+
if xerrors.Is(err, sql.ErrNoRows) {
143+
return ErrProvisionerKeyDeleted
144+
}
145+
if err != nil {
146+
return xerrors.Errorf("lock provisioner key: %w", err)
147+
}
148+
}
149+
acquired, err := tx.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
150+
OrganizationID: organization,
151+
StartedAt: sql.NullTime{
152+
Time: dbtime.Now(),
153+
Valid: true,
154+
},
155+
WorkerID: uuid.NullUUID{
156+
UUID: worker,
157+
Valid: true,
158+
},
159+
Types: pt,
160+
ProvisionerTags: dbTags,
161+
})
162+
if err != nil {
163+
return err
164+
}
165+
job = acquired
166+
return nil
167+
}, nil)
168+
if xerrors.Is(err, ErrProvisionerKeyDeleted) {
169+
logger.Debug(ctx, "provisioner key deleted, exiting acquire")
170+
// cancel (not done) hands an in-progress clearance to another
171+
// acquiree in the domain, re-dispatching the wakeup this
172+
// acquiree consumed.
173+
if internalError := a.cancel(dk, clearance); internalError != nil {
174+
return database.ProvisionerJob{}, internalError
175+
}
176+
return database.ProvisionerJob{}, ErrProvisionerKeyDeleted
177+
}
136178
if xerrors.Is(err, sql.ErrNoRows) {
137179
logger.Debug(ctx, "no job available")
138180
continue

0 commit comments

Comments
 (0)