From 97e8a5b093971e26bfe24659e256fa086dd4f56c Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Wed, 21 Jan 2026 15:18:43 +0200 Subject: [PATCH] fix(coderd): allow agent auth during workspace shutdown (#21538) 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. Modified the auth query to allow `start` build agents to authenticate temporarily during `stop` execution. 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 Auth closes once `stop` completes. Renamed `GetWorkspaceAgentAndLatestBuildByAuthToken` to `GetAuthenticatedWorkspaceAgentAndBuildByAuthToken` since it returns the agent's build (not always latest) during shutdown. Closes coder/internal#1249 Fixes #19467 --- coderd/database/dbauthz/dbauthz.go | 16 +- coderd/database/dbauthz/dbauthz_test.go | 4 +- coderd/database/dbmetrics/querymetrics.go | 16 +- coderd/database/dbmock/dbmock.go | 30 +- coderd/database/querier.go | 7 +- coderd/database/querier_test.go | 405 ++++++++++++++++++++ coderd/database/queries.sql.go | 60 ++- coderd/database/queries/workspaceagents.sql | 52 ++- coderd/httpmw/workspaceagent.go | 2 +- coderd/httpmw/workspaceagent_test.go | 188 +++++++++ coderd/workspaceagents_test.go | 6 +- enterprise/cli/prebuilds_test.go | 2 +- enterprise/coderd/schedule/template_test.go | 2 +- enterprise/coderd/workspaceagents_test.go | 2 +- enterprise/coderd/workspaces_test.go | 6 +- 15 files changed, 726 insertions(+), 72 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index a15ed466ee..53cff0567b 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2244,6 +2244,14 @@ func (q *querier) GetAuditLogsOffset(ctx context.Context, arg database.GetAuditL return q.db.GetAuthorizedAuditLogsOffset(ctx, arg, prep) } +func (q *querier) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) { + // This is a system function + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil { + return database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow{}, err + } + return q.db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, authToken) +} + func (q *querier) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (database.GetAuthorizationUserRolesRow, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil { return database.GetAuthorizationUserRolesRow{}, err @@ -3606,14 +3614,6 @@ func (q *querier) GetWorkspaceACLByID(ctx context.Context, id uuid.UUID) (databa return q.db.GetWorkspaceACLByID(ctx, id) } -func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { - // This is a system function - if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil { - return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, err - } - return q.db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken) -} - func (q *querier) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) { return fetch(q.log, q.auth, q.db.GetWorkspaceAgentAndWorkspaceByID)(ctx, id) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 473b7e43bf..c6ef2d2490 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3472,9 +3472,9 @@ func (s *MethodTestSuite) TestSystemFunctions() { dbm.EXPECT().GetReplicaByID(gomock.Any(), id).Return(database.Replica{}, sql.ErrNoRows).AnyTimes() check.Args(id).Asserts(rbac.ResourceSystem, policy.ActionRead).Errors(sql.ErrNoRows) })) - s.Run("GetWorkspaceAgentAndLatestBuildByAuthToken", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + s.Run("GetAuthenticatedWorkspaceAgentAndBuildByAuthToken", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { tok := uuid.New() - dbm.EXPECT().GetWorkspaceAgentAndLatestBuildByAuthToken(gomock.Any(), tok).Return(database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, sql.ErrNoRows).AnyTimes() + dbm.EXPECT().GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(gomock.Any(), tok).Return(database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow{}, sql.ErrNoRows).AnyTimes() check.Args(tok).Asserts(rbac.ResourceSystem, policy.ActionRead).Errors(sql.ErrNoRows) })) s.Run("GetUserLinksByUserID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 659e3e5e89..38e0a78beb 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -919,6 +919,14 @@ func (m queryMetricsStore) GetAuditLogsOffset(ctx context.Context, arg database. return r0, r1 } +func (m queryMetricsStore) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) { + start := time.Now() + r0, r1 := m.s.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, authToken) + m.queryLatencies.WithLabelValues("GetAuthenticatedWorkspaceAgentAndBuildByAuthToken").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAuthenticatedWorkspaceAgentAndBuildByAuthToken").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (database.GetAuthorizationUserRolesRow, error) { start := time.Now() r0, r1 := m.s.GetAuthorizationUserRoles(ctx, userID) @@ -2207,14 +2215,6 @@ func (m queryMetricsStore) GetWorkspaceACLByID(ctx context.Context, id uuid.UUID return r0, r1 } -func (m queryMetricsStore) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { - start := time.Now() - r0, r1 := m.s.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken) - m.queryLatencies.WithLabelValues("GetWorkspaceAgentAndLatestBuildByAuthToken").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceAgentAndLatestBuildByAuthToken").Inc() - return r0, r1 -} - func (m queryMetricsStore) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) { start := time.Now() r0, r1 := m.s.GetWorkspaceAgentAndWorkspaceByID(ctx, id) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index aad6fc0746..f57a6afb73 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1572,6 +1572,21 @@ func (mr *MockStoreMockRecorder) GetAuditLogsOffset(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAuditLogsOffset", reflect.TypeOf((*MockStore)(nil).GetAuditLogsOffset), ctx, arg) } +// GetAuthenticatedWorkspaceAgentAndBuildByAuthToken mocks base method. +func (m *MockStore) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAuthenticatedWorkspaceAgentAndBuildByAuthToken", ctx, authToken) + ret0, _ := ret[0].(database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAuthenticatedWorkspaceAgentAndBuildByAuthToken indicates an expected call of GetAuthenticatedWorkspaceAgentAndBuildByAuthToken. +func (mr *MockStoreMockRecorder) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, authToken any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAuthenticatedWorkspaceAgentAndBuildByAuthToken", reflect.TypeOf((*MockStore)(nil).GetAuthenticatedWorkspaceAgentAndBuildByAuthToken), ctx, authToken) +} + // GetAuthorizationUserRoles mocks base method. func (m *MockStore) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (database.GetAuthorizationUserRolesRow, error) { m.ctrl.T.Helper() @@ -4122,21 +4137,6 @@ func (mr *MockStoreMockRecorder) GetWorkspaceACLByID(ctx, id any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceACLByID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceACLByID), ctx, id) } -// GetWorkspaceAgentAndLatestBuildByAuthToken mocks base method. -func (m *MockStore) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetWorkspaceAgentAndLatestBuildByAuthToken", ctx, authToken) - ret0, _ := ret[0].(database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// GetWorkspaceAgentAndLatestBuildByAuthToken indicates an expected call of GetWorkspaceAgentAndLatestBuildByAuthToken. -func (mr *MockStoreMockRecorder) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentAndLatestBuildByAuthToken", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentAndLatestBuildByAuthToken), ctx, authToken) -} - // GetWorkspaceAgentAndWorkspaceByID mocks base method. func (m *MockStore) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index e864869b98..a91410676e 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -190,6 +190,12 @@ type sqlcQuerier interface { // GetAuditLogsBefore retrieves `row_limit` number of audit logs before the provided // ID. GetAuditLogsOffset(ctx context.Context, arg GetAuditLogsOffsetParams) ([]GetAuditLogsOffsetRow, error) + // GetAuthenticatedWorkspaceAgentAndBuildByAuthToken returns an authenticated + // workspace agent and its associated build. During normal operation, this is + // the latest build. During shutdown, this may be the previous START build while + // the STOP build is executing, allowing shutdown scripts to authenticate (see + // issue #19467). + GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) // This function returns roles for authorization purposes. Implied member roles // are included. GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (GetAuthorizationUserRolesRow, error) @@ -466,7 +472,6 @@ type sqlcQuerier interface { GetWebpushSubscriptionsByUserID(ctx context.Context, userID uuid.UUID) ([]WebpushSubscription, error) GetWebpushVAPIDKeys(ctx context.Context) (GetWebpushVAPIDKeysRow, error) GetWorkspaceACLByID(ctx context.Context, id uuid.UUID) (GetWorkspaceACLByIDRow, error) - GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (GetWorkspaceAgentAndWorkspaceByIDRow, error) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (WorkspaceAgent, error) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (WorkspaceAgent, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 0ca1fcd465..50863395cf 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -8077,3 +8077,408 @@ func TestDeleteExpiredAPIKeys(t *testing.T) { require.NoError(t, err) require.Len(t, remaining, len(unexpiredTimes)) } + +func TestGetAuthenticatedWorkspaceAgentAndBuildByAuthToken_ShutdownScripts(t *testing.T) { + t.Parallel() + if testing.Short() { + t.SkipNow() + } + + sqlDB := testSQLDB(t) + err := migrations.Up(sqlDB) + require.NoError(t, err) + db := database.New(sqlDB) + + org := dbgen.Organization(t, db, database.Organization{}) + owner := dbgen.User(t, db, database.User{}) + tpl := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + CreatedBy: owner.ID, + }) + ver := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + TemplateID: uuid.NullUUID{ + UUID: tpl.ID, + Valid: true, + }, + OrganizationID: tpl.OrganizationID, + CreatedBy: owner.ID, + }) + + t.Run("DuringStopBuild", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: owner.ID, + OrganizationID: org.ID, + TemplateID: tpl.ID, + }) + + // Create start build with succeeded job (already completed). + startJob := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob) + startJob = dbgen.ProvisionerJob(t, db, nil, startJob) + startResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: startJob.ID, + Transition: database.WorkspaceTransitionStart, + }) + startBuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 1, + Transition: database.WorkspaceTransitionStart, + InitiatorID: owner.ID, + JobID: startJob.ID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: startResource.ID, + }) + + // Create stop build (becomes latest). + stopJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + JobStatus: database.ProvisionerJobStatusRunning, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStop, + InitiatorID: owner.ID, + JobID: stopJob.ID, + }) + + // Agent should still authenticate during stop build execution. + row, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent.AuthToken) + require.NoError(t, err, "agent should authenticate during stop build execution") + require.Equal(t, agent.ID, row.WorkspaceAgent.ID) + require.Equal(t, startBuild.ID, row.WorkspaceBuild.ID, "should return start build, not stop build") + }) + + t.Run("AfterStopJobCompletes", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: owner.ID, + OrganizationID: org.ID, + TemplateID: tpl.ID, + }) + + // Create start build with completed job. + startJob := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob) + startJob = dbgen.ProvisionerJob(t, db, nil, startJob) + + startResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: startJob.ID, + Transition: database.WorkspaceTransitionStart, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 1, + Transition: database.WorkspaceTransitionStart, + InitiatorID: owner.ID, + JobID: startJob.ID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: startResource.ID, + }) + + // Create stop build (becomes latest) with completed job. + stopJob := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusSucceeded, &stopJob) + stopJob = dbgen.ProvisionerJob(t, db, nil, stopJob) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStop, + InitiatorID: owner.ID, + JobID: stopJob.ID, + }) + + // Agent should NOT authenticate after stop job completes. + _, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent.AuthToken) + require.ErrorIs(t, err, sql.ErrNoRows, "agent should not authenticate after stop job completes") + }) + + t.Run("FailedStartBuild", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: owner.ID, + OrganizationID: org.ID, + TemplateID: tpl.ID, + }) + + // Create START build with FAILED job. + startJob := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusFailed, &startJob) + startJob = dbgen.ProvisionerJob(t, db, nil, startJob) + startResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: startJob.ID, + Transition: database.WorkspaceTransitionStart, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 1, + Transition: database.WorkspaceTransitionStart, + InitiatorID: owner.ID, + JobID: startJob.ID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: startResource.ID, + }) + + // Create STOP build with running job. + stopJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + JobStatus: database.ProvisionerJobStatusRunning, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStop, + InitiatorID: owner.ID, + JobID: stopJob.ID, + }) + + // Agent should NOT authenticate (start build failed). + _, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent.AuthToken) + require.ErrorIs(t, err, sql.ErrNoRows, "agent from failed start build should not authenticate") + }) + + t.Run("PendingStopBuild", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: owner.ID, + OrganizationID: org.ID, + TemplateID: tpl.ID, + }) + + // Create start build with succeeded job. + startJob := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob) + startJob = dbgen.ProvisionerJob(t, db, nil, startJob) + startResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: startJob.ID, + Transition: database.WorkspaceTransitionStart, + }) + startBuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 1, + Transition: database.WorkspaceTransitionStart, + InitiatorID: owner.ID, + JobID: startJob.ID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: startResource.ID, + }) + + // Create stop build with pending job (not started yet). + stopJob := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusPending, &stopJob) + stopJob = dbgen.ProvisionerJob(t, db, nil, stopJob) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStop, + InitiatorID: owner.ID, + JobID: stopJob.ID, + }) + + // Agent should authenticate during pending stop build. + row, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent.AuthToken) + require.NoError(t, err, "agent should authenticate during pending stop build") + require.Equal(t, agent.ID, row.WorkspaceAgent.ID) + require.Equal(t, startBuild.ID, row.WorkspaceBuild.ID, "should return start build") + }) + + t.Run("MultipleStartStopCycles", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: owner.ID, + OrganizationID: org.ID, + TemplateID: tpl.ID, + }) + + // Build 1: START (succeeded). + startJob1 := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob1) + startJob1 = dbgen.ProvisionerJob(t, db, nil, startJob1) + startResource1 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: startJob1.ID, + Transition: database.WorkspaceTransitionStart, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 1, + Transition: database.WorkspaceTransitionStart, + InitiatorID: owner.ID, + JobID: startJob1.ID, + }) + agent1 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: startResource1.ID, + }) + + // Build 2: STOP (succeeded). + stopJob1 := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusSucceeded, &stopJob1) + stopJob1 = dbgen.ProvisionerJob(t, db, nil, stopJob1) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStop, + InitiatorID: owner.ID, + JobID: stopJob1.ID, + }) + + // Build 3: START (succeeded). + startJob2 := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob2) + startJob2 = dbgen.ProvisionerJob(t, db, nil, startJob2) + startResource2 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: startJob2.ID, + Transition: database.WorkspaceTransitionStart, + }) + startBuild2 := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 3, + Transition: database.WorkspaceTransitionStart, + InitiatorID: owner.ID, + JobID: startJob2.ID, + }) + agent2 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: startResource2.ID, + }) + + // Build 4: STOP (running). + stopJob2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + JobStatus: database.ProvisionerJobStatusRunning, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 4, + Transition: database.WorkspaceTransitionStop, + InitiatorID: owner.ID, + JobID: stopJob2.ID, + }) + + // Agent from build 3 should authenticate. + row, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent2.AuthToken) + require.NoError(t, err, "agent from most recent start should authenticate during stop") + require.Equal(t, agent2.ID, row.WorkspaceAgent.ID) + require.Equal(t, startBuild2.ID, row.WorkspaceBuild.ID) + + // Agent from build 1 should NOT authenticate. + _, err = db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent1.AuthToken) + require.ErrorIs(t, err, sql.ErrNoRows, "agent from old cycle should not authenticate") + }) + + t.Run("WrongTransitionType", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: owner.ID, + OrganizationID: org.ID, + TemplateID: tpl.ID, + }) + + // Create first start build. + startJob1 := database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + } + setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob1) + startJob1 = dbgen.ProvisionerJob(t, db, nil, startJob1) + startResource1 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: startJob1.ID, + Transition: database.WorkspaceTransitionStart, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 1, + Transition: database.WorkspaceTransitionStart, + InitiatorID: owner.ID, + JobID: startJob1.ID, + }) + agent1 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: startResource1.ID, + }) + + // Create another START build as latest (not STOP). + startJob2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + Type: database.ProvisionerJobTypeWorkspaceBuild, + InitiatorID: owner.ID, + OrganizationID: org.ID, + JobStatus: database.ProvisionerJobStatusRunning, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: ver.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStart, + InitiatorID: owner.ID, + JobID: startJob2.ID, + }) + + // Agent from build 1 should NOT authenticate (latest is not STOP). + _, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent1.AuthToken) + require.ErrorIs(t, err, sql.ErrNoRows, "agent should not authenticate when latest build is not STOP") + }) +} diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 4adba352db..680016f189 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -18024,7 +18024,7 @@ func (q *sqlQuerier) DeleteWorkspaceSubAgentByID(ctx context.Context, id uuid.UU return err } -const getWorkspaceAgentAndLatestBuildByAuthToken = `-- name: GetWorkspaceAgentAndLatestBuildByAuthToken :one +const getAuthenticatedWorkspaceAgentAndBuildByAuthToken = `-- name: GetAuthenticatedWorkspaceAgentAndBuildByAuthToken :one SELECT workspaces.id, workspaces.created_at, workspaces.updated_at, workspaces.owner_id, workspaces.organization_id, workspaces.template_id, workspaces.deleted, workspaces.name, workspaces.autostart_schedule, workspaces.ttl, workspaces.last_used_at, workspaces.dormant_at, workspaces.deleting_at, workspaces.automatic_updates, workspaces.favorite, workspaces.next_start_at, workspaces.group_acl, workspaces.user_acl, workspace_agents.id, workspace_agents.created_at, workspace_agents.updated_at, workspace_agents.name, workspace_agents.first_connected_at, workspace_agents.last_connected_at, workspace_agents.disconnected_at, workspace_agents.resource_id, workspace_agents.auth_token, workspace_agents.auth_instance_id, workspace_agents.architecture, workspace_agents.environment_variables, workspace_agents.operating_system, workspace_agents.instance_metadata, workspace_agents.resource_metadata, workspace_agents.directory, workspace_agents.version, workspace_agents.last_connected_replica_id, workspace_agents.connection_timeout_seconds, workspace_agents.troubleshooting_url, workspace_agents.motd_file, workspace_agents.lifecycle_state, workspace_agents.expanded_directory, workspace_agents.logs_length, workspace_agents.logs_overflowed, workspace_agents.started_at, workspace_agents.ready_at, workspace_agents.subsystems, workspace_agents.display_apps, workspace_agents.api_version, workspace_agents.display_order, workspace_agents.parent_id, workspace_agents.api_key_scope, workspace_agents.deleted, @@ -18054,29 +18054,57 @@ WHERE AND workspaces.deleted = FALSE -- Filter out deleted sub agents. AND workspace_agents.deleted = FALSE - -- Filter out builds that are not the latest. - AND workspace_build_with_user.build_number = ( - -- Select from workspace_builds as it's one less join compared - -- to workspace_build_with_user. - SELECT - MAX(build_number) - FROM - workspace_builds - WHERE - workspace_id = workspace_build_with_user.workspace_id - ) + -- Filter out builds that are not the latest, with exception for shutdown case. + -- Use CASE for short-circuiting: check normal case first (most common), then shutdown case. + AND CASE + -- Normal case: Agent's build is the latest build. + WHEN workspace_build_with_user.build_number = ( + SELECT + MAX(build_number) + FROM + workspace_builds + WHERE + workspace_id = workspace_build_with_user.workspace_id + ) THEN TRUE + -- Shutdown case: Agent from previous START build during STOP build execution. + WHEN workspace_build_with_user.transition = 'start' + -- Agent's START build job succeeded. + AND (SELECT job_status FROM provisioner_jobs WHERE id = workspace_build_with_user.job_id) = 'succeeded' + -- Latest build is a STOP build whose job is still active, + -- and agent's build is immediately previous. + AND EXISTS ( + SELECT 1 + FROM workspace_builds latest + JOIN provisioner_jobs pj ON pj.id = latest.job_id + WHERE latest.workspace_id = workspace_build_with_user.workspace_id + AND latest.build_number = workspace_build_with_user.build_number + 1 + AND latest.build_number = ( + SELECT MAX(build_number) + FROM workspace_builds l2 + WHERE l2.workspace_id = latest.workspace_id + ) + AND latest.transition = 'stop' + AND pj.job_status IN ('pending', 'running') + ) THEN TRUE + ELSE FALSE + END ` -type GetWorkspaceAgentAndLatestBuildByAuthTokenRow struct { +type GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow struct { WorkspaceTable WorkspaceTable `db:"workspace_table" json:"workspace_table"` WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"` WorkspaceBuild WorkspaceBuild `db:"workspace_build" json:"workspace_build"` TaskID uuid.NullUUID `db:"task_id" json:"task_id"` } -func (q *sqlQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { - row := q.db.QueryRowContext(ctx, getWorkspaceAgentAndLatestBuildByAuthToken, authToken) - var i GetWorkspaceAgentAndLatestBuildByAuthTokenRow +// GetAuthenticatedWorkspaceAgentAndBuildByAuthToken returns an authenticated +// workspace agent and its associated build. During normal operation, this is +// the latest build. During shutdown, this may be the previous START build while +// the STOP build is executing, allowing shutdown scripts to authenticate (see +// issue #19467). +func (q *sqlQuerier) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) { + row := q.db.QueryRowContext(ctx, getAuthenticatedWorkspaceAgentAndBuildByAuthToken, authToken) + var i GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow err := row.Scan( &i.WorkspaceTable.ID, &i.WorkspaceTable.CreatedAt, diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql index 0ac15d63a8..3e52e6269f 100644 --- a/coderd/database/queries/workspaceagents.sql +++ b/coderd/database/queries/workspaceagents.sql @@ -281,7 +281,12 @@ WHERE -- Filter out deleted sub agents. AND workspace_agents.deleted = FALSE; --- name: GetWorkspaceAgentAndLatestBuildByAuthToken :one +-- GetAuthenticatedWorkspaceAgentAndBuildByAuthToken returns an authenticated +-- workspace agent and its associated build. During normal operation, this is +-- the latest build. During shutdown, this may be the previous START build while +-- the STOP build is executing, allowing shutdown scripts to authenticate (see +-- issue #19467). +-- name: GetAuthenticatedWorkspaceAgentAndBuildByAuthToken :one SELECT sqlc.embed(workspaces), sqlc.embed(workspace_agents), @@ -311,17 +316,40 @@ WHERE AND workspaces.deleted = FALSE -- Filter out deleted sub agents. AND workspace_agents.deleted = FALSE - -- Filter out builds that are not the latest. - AND workspace_build_with_user.build_number = ( - -- Select from workspace_builds as it's one less join compared - -- to workspace_build_with_user. - SELECT - MAX(build_number) - FROM - workspace_builds - WHERE - workspace_id = workspace_build_with_user.workspace_id - ) + -- Filter out builds that are not the latest, with exception for shutdown case. + -- Use CASE for short-circuiting: check normal case first (most common), then shutdown case. + AND CASE + -- Normal case: Agent's build is the latest build. + WHEN workspace_build_with_user.build_number = ( + SELECT + MAX(build_number) + FROM + workspace_builds + WHERE + workspace_id = workspace_build_with_user.workspace_id + ) THEN TRUE + -- Shutdown case: Agent from previous START build during STOP build execution. + WHEN workspace_build_with_user.transition = 'start' + -- Agent's START build job succeeded. + AND (SELECT job_status FROM provisioner_jobs WHERE id = workspace_build_with_user.job_id) = 'succeeded' + -- Latest build is a STOP build whose job is still active, + -- and agent's build is immediately previous. + AND EXISTS ( + SELECT 1 + FROM workspace_builds latest + JOIN provisioner_jobs pj ON pj.id = latest.job_id + WHERE latest.workspace_id = workspace_build_with_user.workspace_id + AND latest.build_number = workspace_build_with_user.build_number + 1 + AND latest.build_number = ( + SELECT MAX(build_number) + FROM workspace_builds l2 + WHERE l2.workspace_id = latest.workspace_id + ) + AND latest.transition = 'stop' + AND pj.job_status IN ('pending', 'running') + ) THEN TRUE + ELSE FALSE + END ; -- name: InsertWorkspaceAgentScriptTimings :one diff --git a/coderd/httpmw/workspaceagent.go b/coderd/httpmw/workspaceagent.go index d5f4e6fef2..47867e17b2 100644 --- a/coderd/httpmw/workspaceagent.go +++ b/coderd/httpmw/workspaceagent.go @@ -92,7 +92,7 @@ func ExtractWorkspaceAgentAndLatestBuild(opts ExtractWorkspaceAgentAndLatestBuil } //nolint:gocritic // System needs to be able to get workspace agents. - row, err := opts.DB.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), token) + row, err := opts.DB.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), token) if err != nil { if errors.Is(err, sql.ErrNoRows) { optionalWrite(http.StatusUnauthorized, codersdk.Response{ diff --git a/coderd/httpmw/workspaceagent_test.go b/coderd/httpmw/workspaceagent_test.go index 8d79b6ddbd..378d75927c 100644 --- a/coderd/httpmw/workspaceagent_test.go +++ b/coderd/httpmw/workspaceagent_test.go @@ -1,9 +1,11 @@ package httpmw_test import ( + "database/sql" "net/http" "net/http/httptest" "testing" + "time" "github.com/go-chi/chi/v5" "github.com/google/uuid" @@ -12,6 +14,7 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/httpmw" "github.com/coder/coder/v2/codersdk" ) @@ -95,6 +98,180 @@ func TestWorkspaceAgent(t *testing.T) { t.Cleanup(func() { _ = res.Body.Close() }) require.Equal(t, http.StatusUnauthorized, res.StatusCode) }) + + t.Run("DuringShutdown", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + authToken := uuid.New() + req, rtr, ws, tpv := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild( + httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{ + DB: db, + Optional: false, + }), + ) + + // Create a STOP build with running job (becomes latest). + stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{ + OrganizationID: ws.OrganizationID, + JobStatus: database.ProvisionerJobStatusRunning, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: ws.ID, + JobID: stopJob.ID, + TemplateVersionID: tpv.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStop, + }) + + // Agent should still authenticate during shutdown. + rw := httptest.NewRecorder() + req.Header.Set(codersdk.SessionTokenHeader, authToken.String()) + rtr.ServeHTTP(rw, req) + + //nolint:bodyclose // Closed in `t.Cleanup` + res := rw.Result() + t.Cleanup(func() { _ = res.Body.Close() }) + require.Equal(t, http.StatusOK, res.StatusCode, "agent should authenticate during stop build execution") + }) + + t.Run("AfterShutdownCompletes", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + authToken := uuid.New() + req, rtr, ws, tpv := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild( + httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{ + DB: db, + Optional: false, + }), + ) + + // Create a STOP build with completed job (becomes latest). + stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{ + OrganizationID: ws.OrganizationID, + JobStatus: database.ProvisionerJobStatusSucceeded, + CompletedAt: sql.NullTime{ + Time: dbtime.Now(), + Valid: true, + }, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: ws.ID, + JobID: stopJob.ID, + TemplateVersionID: tpv.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStop, + }) + + // Agent should NOT authenticate after stop job completes. + rw := httptest.NewRecorder() + req.Header.Set(codersdk.SessionTokenHeader, authToken.String()) + rtr.ServeHTTP(rw, req) + + //nolint:bodyclose // Closed in `t.Cleanup` + res := rw.Result() + t.Cleanup(func() { _ = res.Body.Close() }) + require.Equal(t, http.StatusUnauthorized, res.StatusCode, "agent should not authenticate after stop job completes") + }) + + t.Run("FailedStartBuild", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + authToken := uuid.New() + + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{ + Status: database.UserStatusActive, + }) + _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + template := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + ActiveVersionID: templateVersion.ID, + CreatedBy: user.ID, + }) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: user.ID, + OrganizationID: org.ID, + TemplateID: template.ID, + }) + + // Create START build with FAILED job status. + startJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{ + OrganizationID: org.ID, + JobStatus: database.ProvisionerJobStatusFailed, + StartedAt: sql.NullTime{ + Time: dbtime.Now().Add(-time.Minute), + Valid: true, + }, + CompletedAt: sql.NullTime{ + Time: dbtime.Now(), + Valid: true, + }, + Error: sql.NullString{ + String: "build failed", + Valid: true, + }, + ErrorCode: sql.NullString{ + String: "FAILED", + Valid: true, + }, + }) + resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: startJob.ID, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + JobID: startJob.ID, + TemplateVersionID: templateVersion.ID, + BuildNumber: 1, + Transition: database.WorkspaceTransitionStart, + }) + _ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: resource.ID, + AuthToken: authToken, + }) + + // Create a STOP build with running job. + stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{ + OrganizationID: org.ID, + JobStatus: database.ProvisionerJobStatusRunning, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + JobID: stopJob.ID, + TemplateVersionID: templateVersion.ID, + BuildNumber: 2, + Transition: database.WorkspaceTransitionStop, + }) + + req := httptest.NewRequest("GET", "/", nil) + rtr := chi.NewRouter() + rtr.Use(httpmw.ExtractWorkspaceAgentAndLatestBuild( + httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{ + DB: db, + Optional: false, + })) + rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) { + _ = httpmw.WorkspaceAgent(r) + rw.WriteHeader(http.StatusOK) + }) + + // Agent should NOT authenticate (start build failed). + rw := httptest.NewRecorder() + req.Header.Set(codersdk.SessionTokenHeader, authToken.String()) + rtr.ServeHTTP(rw, req) + + //nolint:bodyclose // Closed in `t.Cleanup` + res := rw.Result() + t.Cleanup(func() { _ = res.Body.Close() }) + require.Equal(t, http.StatusUnauthorized, res.StatusCode, "agent should not authenticate when start build failed") + }) } func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Handler) http.Handler) (*http.Request, http.Handler, database.WorkspaceTable, database.TemplateVersion) { @@ -123,6 +300,15 @@ func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Ha }) job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ OrganizationID: org.ID, + JobStatus: database.ProvisionerJobStatusSucceeded, + StartedAt: sql.NullTime{ + Time: dbtime.Now().Add(-30 * time.Second), + Valid: true, + }, + CompletedAt: sql.NullTime{ + Time: dbtime.Now(), + Valid: true, + }, }) resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ JobID: job.ID, @@ -131,6 +317,8 @@ func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Ha WorkspaceID: workspace.ID, JobID: job.ID, TemplateVersionID: templateVersion.ID, + BuildNumber: 1, + Transition: database.WorkspaceTransitionStart, }) _ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ ResourceID: resource.ID, diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 497c8c35c6..101b2040e8 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -623,7 +623,7 @@ func TestWorkspaceAgentClientCoordinate_BadVersion(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) agentToken, err := uuid.Parse(r.AgentToken) require.NoError(t, err) - ao, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentToken) + ao, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentToken) require.NoError(t, err) //nolint: bodyclose // closed by ReadBodyAsError @@ -713,7 +713,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) { agentTokenUUID, err := uuid.Parse(r.AgentToken) require.NoError(t, err) ctx := testutil.Context(t, testutil.WaitLong) - agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) + agentAndBuild, err := api.Database.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) require.NoError(t, err) // Connect with no resume token, and ensure that the peer ID is set to a @@ -785,7 +785,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) { agentTokenUUID, err := uuid.Parse(r.AgentToken) require.NoError(t, err) ctx := testutil.Context(t, testutil.WaitLong) - agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) + agentAndBuild, err := api.Database.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID) require.NoError(t, err) // Connect with no resume token, and ensure that the peer ID is set to a diff --git a/enterprise/cli/prebuilds_test.go b/enterprise/cli/prebuilds_test.go index cf0c741050..c5b755c7fc 100644 --- a/enterprise/cli/prebuilds_test.go +++ b/enterprise/cli/prebuilds_test.go @@ -435,7 +435,7 @@ func TestSchedulePrebuilds(t *testing.T) { // Mark the prebuilt workspace's agent as ready so the prebuild can be claimed ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong)) - agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken)) + agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken)) require.NoError(t, err) err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ ID: agent.WorkspaceAgent.ID, diff --git a/enterprise/coderd/schedule/template_test.go b/enterprise/coderd/schedule/template_test.go index 206bac7023..c03a1fcd22 100644 --- a/enterprise/coderd/schedule/template_test.go +++ b/enterprise/coderd/schedule/template_test.go @@ -1363,7 +1363,7 @@ func TestTemplateUpdatePrebuilds(t *testing.T) { // Mark the prebuilt workspace's agent as ready so the prebuild can be claimed // nolint:gocritic agentCtx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong)) - agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(agentCtx, uuid.MustParse(workspaceBuild.AgentToken)) + agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(agentCtx, uuid.MustParse(workspaceBuild.AgentToken)) require.NoError(t, err) err = db.UpdateWorkspaceAgentLifecycleStateByID(agentCtx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ ID: agent.WorkspaceAgent.ID, diff --git a/enterprise/coderd/workspaceagents_test.go b/enterprise/coderd/workspaceagents_test.go index 18e1aa0101..574f2b5be2 100644 --- a/enterprise/coderd/workspaceagents_test.go +++ b/enterprise/coderd/workspaceagents_test.go @@ -187,7 +187,7 @@ func TestReinitializeAgent(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) var prebuildID uuid.UUID require.Eventually(t, func() bool { - agentAndBuild, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, agentToken) + agentAndBuild, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, agentToken) if err != nil { return false } diff --git a/enterprise/coderd/workspaces_test.go b/enterprise/coderd/workspaces_test.go index d9ed713d24..a1d26197a9 100644 --- a/enterprise/coderd/workspaces_test.go +++ b/enterprise/coderd/workspaces_test.go @@ -567,7 +567,7 @@ func TestCreateUserWorkspace(t *testing.T) { }).Do() ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong)) - agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(r.AgentToken)) + agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, uuid.MustParse(r.AgentToken)) require.NoError(t, err) err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ @@ -2788,7 +2788,7 @@ func TestPrebuildUpdateLifecycleParams(t *testing.T) { // Mark the prebuilt workspace's agent as ready so the prebuild can be claimed ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong)) - agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken)) + agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken)) require.NoError(t, err) err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ ID: agent.WorkspaceAgent.ID, @@ -2887,7 +2887,7 @@ func TestPrebuildActivityBump(t *testing.T) { // Mark the prebuilt workspace's agent as ready so the prebuild can be claimed // nolint:gocritic ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong)) - agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(wb.AgentToken)) + agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, uuid.MustParse(wb.AgentToken)) require.NoError(t, err) err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ ID: agent.WorkspaceAgent.ID,