Skip to content

Commit 85001ba

Browse files
mafredrimafredri
committed
feat(coderd): allow agent auth during workspace shutdown
Agents were losing authentication during workspace shutdown, causing shutdown scripts to fail. The auth query required agents to belong to the latest build, but during shutdown a `stop` build becomes latest while the `start` build's agents are still running. Several approaches were considered: - Reusing agent tokens (no unique constraint, would break auth) - Separate validation functions (two queries for agent + latest build) - Grace period with workspace-level tokens (security risk, stale agents) - Agent lifecycle state tracking (requires migration, complex state management) Modified the auth query to allow `start` build agents to authenticate temporarily during `stop` execution. This is simpler than separate functions and avoids migrations. The query allows auth when: - Agent's `start` build job succeeded - Latest build is `stop` with `pending`/`running` job status - Builds are adjacent (`stop` is `build_number + 1`) - Template versions match The version check prevents RBAC issues if someone explicitly changes template versions between builds. Auth closes once `stop` completes. Renamed `GetWorkspaceAgentAndLatestBuildByAuthToken` to `GetAuthenticatedWorkspaceAgentAndBuildByAuthToken` since it returns the agent's build (not always latest) during shutdown. Added tests covering happy path, failure cases (completed stops, failed builds, version mismatches), and edge cases (multiple cycles, wrong transitions) at both database and middleware levels. Fixes coder/internal#1249 Fixes #19467
1 parent 6ef9670 commit 85001ba

15 files changed

Lines changed: 792 additions & 70 deletions

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2244,6 +2244,14 @@ func (q *querier) GetAuditLogsOffset(ctx context.Context, arg database.GetAuditL
22442244
return q.db.GetAuthorizedAuditLogsOffset(ctx, arg, prep)
22452245
}
22462246

2247+
func (q *querier) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) {
2248+
// This is a system function
2249+
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
2250+
return database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow{}, err
2251+
}
2252+
return q.db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, authToken)
2253+
}
2254+
22472255
func (q *querier) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (database.GetAuthorizationUserRolesRow, error) {
22482256
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
22492257
return database.GetAuthorizationUserRolesRow{}, err
@@ -3606,14 +3614,6 @@ func (q *querier) GetWorkspaceACLByID(ctx context.Context, id uuid.UUID) (databa
36063614
return q.db.GetWorkspaceACLByID(ctx, id)
36073615
}
36083616

3609-
func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) {
3610-
// This is a system function
3611-
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
3612-
return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, err
3613-
}
3614-
return q.db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken)
3615-
}
3616-
36173617
func (q *querier) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) {
36183618
return fetch(q.log, q.auth, q.db.GetWorkspaceAgentAndWorkspaceByID)(ctx, id)
36193619
}

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3472,9 +3472,9 @@ func (s *MethodTestSuite) TestSystemFunctions() {
34723472
dbm.EXPECT().GetReplicaByID(gomock.Any(), id).Return(database.Replica{}, sql.ErrNoRows).AnyTimes()
34733473
check.Args(id).Asserts(rbac.ResourceSystem, policy.ActionRead).Errors(sql.ErrNoRows)
34743474
}))
3475-
s.Run("GetWorkspaceAgentAndLatestBuildByAuthToken", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
3475+
s.Run("GetAuthenticatedWorkspaceAgentAndBuildByAuthToken", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
34763476
tok := uuid.New()
3477-
dbm.EXPECT().GetWorkspaceAgentAndLatestBuildByAuthToken(gomock.Any(), tok).Return(database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, sql.ErrNoRows).AnyTimes()
3477+
dbm.EXPECT().GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(gomock.Any(), tok).Return(database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow{}, sql.ErrNoRows).AnyTimes()
34783478
check.Args(tok).Asserts(rbac.ResourceSystem, policy.ActionRead).Errors(sql.ErrNoRows)
34793479
}))
34803480
s.Run("GetUserLinksByUserID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 8 additions & 8 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 & 15 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: 6 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)