diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index a1473f48f0..ee43dc228b 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4873,6 +4873,14 @@ func (q *querier) GetWorkspaceAppsCreatedAfter(ctx context.Context, createdAt ti return q.db.GetWorkspaceAppsCreatedAfter(ctx, createdAt) } +func (q *querier) GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.GetWorkspaceBuildAgentsByInstanceIDRow, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err == nil { + return q.db.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID) + } + + return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetWorkspaceBuildAgentsByInstanceID)(ctx, authInstanceID) +} + func (q *querier) GetWorkspaceBuildByID(ctx context.Context, buildID uuid.UUID) (database.WorkspaceBuild, error) { build, err := q.db.GetWorkspaceBuildByID(ctx, buildID) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 913a26d6fa..8226234644 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3257,6 +3257,19 @@ func (s *MethodTestSuite) TestWorkspace() { Returns([]database.WorkspaceAgent{agt}). FailSystemObjectChecks() })) + s.Run("GetWorkspaceBuildAgentsByInstanceID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + w := testutil.Fake(s.T(), faker, database.WorkspaceTable{}) + agt := testutil.Fake(s.T(), faker, database.WorkspaceAgent{}) + row := testutil.Fake(s.T(), faker, database.GetWorkspaceBuildAgentsByInstanceIDRow{}) + row.WorkspaceAgent = agt + row.WorkspaceTable = w + authInstanceID := "instance-id" + dbm.EXPECT().GetWorkspaceBuildAgentsByInstanceID(gomock.Any(), authInstanceID).Return([]database.GetWorkspaceBuildAgentsByInstanceIDRow{row}, nil).AnyTimes() + check.Args(authInstanceID). + Asserts(rbac.ResourceSystem, policy.ActionRead, w, policy.ActionRead). + Returns([]database.GetWorkspaceBuildAgentsByInstanceIDRow{row}). + FailSystemObjectChecks() + })) s.Run("UpdateWorkspaceAgentLifecycleStateByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { w := testutil.Fake(s.T(), faker, database.Workspace{}) agt := testutil.Fake(s.T(), faker, database.WorkspaceAgent{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 1bca8d357d..95697f97fa 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3249,6 +3249,14 @@ func (m queryMetricsStore) GetWorkspaceAppsCreatedAfter(ctx context.Context, cre return r0, r1 } +func (m queryMetricsStore) GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.GetWorkspaceBuildAgentsByInstanceIDRow, error) { + start := time.Now() + r0, r1 := m.s.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID) + m.queryLatencies.WithLabelValues("GetWorkspaceBuildAgentsByInstanceID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceBuildAgentsByInstanceID").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetWorkspaceBuildByID(ctx context.Context, id uuid.UUID) (database.WorkspaceBuild, error) { start := time.Now() r0, r1 := m.s.GetWorkspaceBuildByID(ctx, id) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index eca308d5d0..dab6b4edc3 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -6078,6 +6078,21 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAppsCreatedAfter(ctx, createdAt any return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAppsCreatedAfter", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAppsCreatedAfter), ctx, createdAt) } +// GetWorkspaceBuildAgentsByInstanceID mocks base method. +func (m *MockStore) GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.GetWorkspaceBuildAgentsByInstanceIDRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetWorkspaceBuildAgentsByInstanceID", ctx, authInstanceID) + ret0, _ := ret[0].([]database.GetWorkspaceBuildAgentsByInstanceIDRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetWorkspaceBuildAgentsByInstanceID indicates an expected call of GetWorkspaceBuildAgentsByInstanceID. +func (mr *MockStoreMockRecorder) GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceBuildAgentsByInstanceID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceBuildAgentsByInstanceID), ctx, authInstanceID) +} + // GetWorkspaceBuildByID mocks base method. func (m *MockStore) GetWorkspaceBuildByID(ctx context.Context, id uuid.UUID) (database.WorkspaceBuild, error) { m.ctrl.T.Helper() diff --git a/coderd/database/modelmethods.go b/coderd/database/modelmethods.go index 18c651ce9b..fde2d38c00 100644 --- a/coderd/database/modelmethods.go +++ b/coderd/database/modelmethods.go @@ -932,6 +932,11 @@ func (r GetWorkspaceAgentAndWorkspaceByIDRow) RBACObject() rbac.Object { return r.WorkspaceTable.RBACObject() } +// A workspace agent belongs to the owner of the associated workspace. +func (r GetWorkspaceBuildAgentsByInstanceIDRow) RBACObject() rbac.Object { + return r.WorkspaceTable.RBACObject() +} + // UpsertConnectionLogParams contains the parameters for upserting a // connection log entry. This struct is hand-maintained (not generated // by sqlc) because the single-row UpsertConnectionLog query was diff --git a/coderd/database/querier.go b/coderd/database/querier.go index cc8a499a8b..03737e00aa 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -802,6 +802,7 @@ type sqlcQuerier interface { GetWorkspaceAppsByAgentID(ctx context.Context, agentID uuid.UUID) ([]WorkspaceApp, error) GetWorkspaceAppsByAgentIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceApp, error) GetWorkspaceAppsCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceApp, error) + GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]GetWorkspaceBuildAgentsByInstanceIDRow, error) GetWorkspaceBuildByID(ctx context.Context, id uuid.UUID) (WorkspaceBuild, error) GetWorkspaceBuildByJobID(ctx context.Context, jobID uuid.UUID) (WorkspaceBuild, error) GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx context.Context, arg GetWorkspaceBuildByWorkspaceIDAndBuildNumberParams) (WorkspaceBuild, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 41134ca12e..2d6981dbbf 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -7187,6 +7187,103 @@ func markWorkspaceAgentDeleted(ctx context.Context, t *testing.T, sqlDB *sql.DB, require.NoError(t, err) } +type workspaceBuildAgentQueryFixture struct { + Workspace database.WorkspaceTable + Build database.WorkspaceBuild + Agent database.WorkspaceAgent +} + +func setupWorkspaceBuildAgentQueryWorkspace(t testing.TB, db database.Store, deleted bool) database.WorkspaceTable { + t.Helper() + + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + template := dbgen.Template(t, db, database.Template{ + CreatedBy: user.ID, + OrganizationID: org.ID, + }) + return dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: user.ID, + OrganizationID: org.ID, + TemplateID: template.ID, + Deleted: deleted, + }) +} + +func setupWorkspaceBuildAgentQueryFixture( + t testing.TB, + db database.Store, + authInstanceID string, + name string, + createdAt time.Time, + workspace database.WorkspaceTable, +) workspaceBuildAgentQueryFixture { + t.Helper() + + if workspace.ID == uuid.Nil { + workspace = setupWorkspaceBuildAgentQueryWorkspace(t, db, false) + } + templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + TemplateID: uuid.NullUUID{UUID: workspace.TemplateID, Valid: true}, + OrganizationID: workspace.OrganizationID, + CreatedBy: workspace.OwnerID, + }) + job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: workspace.OrganizationID, + Type: database.ProvisionerJobTypeWorkspaceBuild, + }) + build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: templateVersion.ID, + JobID: job.ID, + }) + resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: job.ID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + Name: name, + ResourceID: resource.ID, + CreatedAt: createdAt, + AuthInstanceID: sql.NullString{ + String: authInstanceID, + Valid: true, + }, + }) + + return workspaceBuildAgentQueryFixture{ + Workspace: workspace, + Build: build, + Agent: agent, + } +} + +func setupProvisionerJobAgentQueryFixture( + t testing.TB, + db database.Store, + authInstanceID string, + name string, + createdAt time.Time, + jobType database.ProvisionerJobType, +) database.WorkspaceAgent { + t.Helper() + + job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + Type: jobType, + }) + resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: job.ID, + }) + return dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + Name: name, + ResourceID: resource.ID, + CreatedAt: createdAt, + AuthInstanceID: sql.NullString{ + String: authInstanceID, + Valid: true, + }, + }) +} + func TestGetWorkspaceAgentsByInstanceID(t *testing.T) { t.Parallel() @@ -7304,6 +7401,90 @@ func TestGetWorkspaceAgentsByInstanceID(t *testing.T) { }) } +func TestGetWorkspaceBuildAgentsByInstanceID(t *testing.T) { + t.Parallel() + + t.Run("ReturnsWorkspaceBuildRootAgentsNewestFirst", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano()) + olderCreatedAt := dbtime.Now().Add(-time.Hour) + newerCreatedAt := dbtime.Now() + + older := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "older", olderCreatedAt, database.WorkspaceTable{}) + newer := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "newer", newerCreatedAt, database.WorkspaceTable{}) + + ctx := testutil.Context(t, testutil.WaitShort) + + agents, err := db.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID) + require.NoError(t, err) + require.Len(t, agents, 2) + assert.Equal(t, []uuid.UUID{newer.Agent.ID, older.Agent.ID}, []uuid.UUID{agents[0].WorkspaceAgent.ID, agents[1].WorkspaceAgent.ID}) + assert.Equal(t, []uuid.UUID{newer.Build.ID, older.Build.ID}, []uuid.UUID{agents[0].WorkspaceBuildID, agents[1].WorkspaceBuildID}) + assert.Equal(t, newer.Workspace.ID, agents[0].WorkspaceTable.ID) + assert.Equal(t, older.Workspace.ID, agents[1].WorkspaceTable.ID) + assert.Equal(t, newer.Workspace.OwnerID, agents[0].WorkspaceTable.OwnerID) + assert.Equal(t, older.Workspace.OwnerID, agents[1].WorkspaceTable.OwnerID) + assert.Equal(t, newer.Workspace.OrganizationID, agents[0].WorkspaceTable.OrganizationID) + assert.Equal(t, older.Workspace.OrganizationID, agents[1].WorkspaceTable.OrganizationID) + assert.False(t, agents[0].WorkspaceTable.Deleted) + assert.False(t, agents[1].WorkspaceTable.Deleted) + }) + + t.Run("ExcludesDeletedAgentsSubAgentsAndNonWorkspaceBuildJobs", func(t *testing.T) { + t.Parallel() + + db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano()) + baseCreatedAt := dbtime.Now() + + root := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "root", baseCreatedAt.Add(-time.Hour), database.WorkspaceTable{}) + _ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ParentID: uuid.NullUUID{UUID: root.Agent.ID, Valid: true}, + Name: "sub", + ResourceID: root.Agent.ResourceID, + CreatedAt: baseCreatedAt.Add(time.Minute), + AuthInstanceID: sql.NullString{ + String: authInstanceID, + Valid: true, + }, + }) + deletedAgent := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "deleted", baseCreatedAt.Add(2*time.Minute), database.WorkspaceTable{}) + _ = setupProvisionerJobAgentQueryFixture(t, db, authInstanceID, "template-import", baseCreatedAt.Add(3*time.Minute), database.ProvisionerJobTypeTemplateVersionImport) + _ = setupProvisionerJobAgentQueryFixture(t, db, authInstanceID, "dry-run", baseCreatedAt.Add(4*time.Minute), database.ProvisionerJobTypeTemplateVersionDryRun) + + ctx := testutil.Context(t, testutil.WaitShort) + markWorkspaceAgentDeleted(ctx, t, sqlDB, deletedAgent.Agent.ID) + + agents, err := db.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID) + require.NoError(t, err) + require.Len(t, agents, 1) + assert.Equal(t, root.Agent.ID, agents[0].WorkspaceAgent.ID) + assert.False(t, agents[0].WorkspaceAgent.ParentID.Valid) + assert.Equal(t, root.Build.ID, agents[0].WorkspaceBuildID) + }) + + t.Run("ExcludesDeletedWorkspaces", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano()) + baseCreatedAt := dbtime.Now() + active := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "active", baseCreatedAt, database.WorkspaceTable{}) + deletedWorkspace := setupWorkspaceBuildAgentQueryWorkspace(t, db, true) + _ = setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "deleted-workspace", baseCreatedAt.Add(time.Minute), deletedWorkspace) + + ctx := testutil.Context(t, testutil.WaitShort) + + agents, err := db.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID) + require.NoError(t, err) + require.Len(t, agents, 1) + assert.Equal(t, active.Agent.ID, agents[0].WorkspaceAgent.ID) + assert.Equal(t, active.Workspace.ID, agents[0].WorkspaceTable.ID) + }) +} + func requireUsersMatch(t testing.TB, expected []database.User, found []database.GetUsersRow, msg string) { t.Helper() require.ElementsMatch(t, expected, database.ConvertUserRows(found), msg) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index fc6df826cd..8995b8e4ed 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -28766,6 +28766,122 @@ func (q *sqlQuerier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Co return items, nil } +const getWorkspaceBuildAgentsByInstanceID = `-- name: GetWorkspaceBuildAgentsByInstanceID :many +SELECT + 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, + workspace_builds.id AS workspace_build_id, + 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 +FROM + workspace_agents +JOIN + workspace_resources +ON + workspace_resources.id = workspace_agents.resource_id +JOIN + workspace_builds +ON + workspace_builds.job_id = workspace_resources.job_id +JOIN + provisioner_jobs +ON + provisioner_jobs.id = workspace_builds.job_id +JOIN + workspaces +ON + workspaces.id = workspace_builds.workspace_id +WHERE + workspace_agents.auth_instance_id = $1 :: TEXT + AND workspace_agents.deleted = FALSE + AND workspace_agents.parent_id IS NULL + AND provisioner_jobs.type = 'workspace_build'::provisioner_job_type + AND workspaces.deleted = FALSE +ORDER BY + workspace_agents.created_at DESC +` + +type GetWorkspaceBuildAgentsByInstanceIDRow struct { + WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"` + WorkspaceBuildID uuid.UUID `db:"workspace_build_id" json:"workspace_build_id"` + WorkspaceTable WorkspaceTable `db:"workspace_table" json:"workspace_table"` +} + +func (q *sqlQuerier) GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]GetWorkspaceBuildAgentsByInstanceIDRow, error) { + rows, err := q.db.QueryContext(ctx, getWorkspaceBuildAgentsByInstanceID, authInstanceID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetWorkspaceBuildAgentsByInstanceIDRow + for rows.Next() { + var i GetWorkspaceBuildAgentsByInstanceIDRow + if err := rows.Scan( + &i.WorkspaceAgent.ID, + &i.WorkspaceAgent.CreatedAt, + &i.WorkspaceAgent.UpdatedAt, + &i.WorkspaceAgent.Name, + &i.WorkspaceAgent.FirstConnectedAt, + &i.WorkspaceAgent.LastConnectedAt, + &i.WorkspaceAgent.DisconnectedAt, + &i.WorkspaceAgent.ResourceID, + &i.WorkspaceAgent.AuthToken, + &i.WorkspaceAgent.AuthInstanceID, + &i.WorkspaceAgent.Architecture, + &i.WorkspaceAgent.EnvironmentVariables, + &i.WorkspaceAgent.OperatingSystem, + &i.WorkspaceAgent.InstanceMetadata, + &i.WorkspaceAgent.ResourceMetadata, + &i.WorkspaceAgent.Directory, + &i.WorkspaceAgent.Version, + &i.WorkspaceAgent.LastConnectedReplicaID, + &i.WorkspaceAgent.ConnectionTimeoutSeconds, + &i.WorkspaceAgent.TroubleshootingURL, + &i.WorkspaceAgent.MOTDFile, + &i.WorkspaceAgent.LifecycleState, + &i.WorkspaceAgent.ExpandedDirectory, + &i.WorkspaceAgent.LogsLength, + &i.WorkspaceAgent.LogsOverflowed, + &i.WorkspaceAgent.StartedAt, + &i.WorkspaceAgent.ReadyAt, + pq.Array(&i.WorkspaceAgent.Subsystems), + pq.Array(&i.WorkspaceAgent.DisplayApps), + &i.WorkspaceAgent.APIVersion, + &i.WorkspaceAgent.DisplayOrder, + &i.WorkspaceAgent.ParentID, + &i.WorkspaceAgent.APIKeyScope, + &i.WorkspaceAgent.Deleted, + &i.WorkspaceBuildID, + &i.WorkspaceTable.ID, + &i.WorkspaceTable.CreatedAt, + &i.WorkspaceTable.UpdatedAt, + &i.WorkspaceTable.OwnerID, + &i.WorkspaceTable.OrganizationID, + &i.WorkspaceTable.TemplateID, + &i.WorkspaceTable.Deleted, + &i.WorkspaceTable.Name, + &i.WorkspaceTable.AutostartSchedule, + &i.WorkspaceTable.Ttl, + &i.WorkspaceTable.LastUsedAt, + &i.WorkspaceTable.DormantAt, + &i.WorkspaceTable.DeletingAt, + &i.WorkspaceTable.AutomaticUpdates, + &i.WorkspaceTable.Favorite, + &i.WorkspaceTable.NextStartAt, + &i.WorkspaceTable.GroupACL, + &i.WorkspaceTable.UserACL, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const insertWorkspaceAgent = `-- name: InsertWorkspaceAgent :one INSERT INTO workspace_agents ( diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql index b75fb61b15..4ca48d0dbe 100644 --- a/coderd/database/queries/workspaceagents.sql +++ b/coderd/database/queries/workspaceagents.sql @@ -22,6 +22,38 @@ WHERE ORDER BY created_at DESC; +-- name: GetWorkspaceBuildAgentsByInstanceID :many +SELECT + sqlc.embed(workspace_agents), + workspace_builds.id AS workspace_build_id, + sqlc.embed(workspaces) +FROM + workspace_agents +JOIN + workspace_resources +ON + workspace_resources.id = workspace_agents.resource_id +JOIN + workspace_builds +ON + workspace_builds.job_id = workspace_resources.job_id +JOIN + provisioner_jobs +ON + provisioner_jobs.id = workspace_builds.job_id +JOIN + workspaces +ON + workspaces.id = workspace_builds.workspace_id +WHERE + workspace_agents.auth_instance_id = @auth_instance_id :: TEXT + AND workspace_agents.deleted = FALSE + AND workspace_agents.parent_id IS NULL + AND provisioner_jobs.type = 'workspace_build'::provisioner_job_type + AND workspaces.deleted = FALSE +ORDER BY + workspace_agents.created_at DESC; + -- name: GetWorkspaceAgentsByResourceIDs :many SELECT * diff --git a/coderd/workspaceresourceauth.go b/coderd/workspaceresourceauth.go index 8371dfb693..a9bf320c95 100644 --- a/coderd/workspaceresourceauth.go +++ b/coderd/workspaceresourceauth.go @@ -1,21 +1,17 @@ package coderd import ( - "encoding/json" "fmt" "net/http" "sort" "strings" - "github.com/google/uuid" "github.com/mitchellh/mapstructure" "github.com/coder/coder/v2/coderd/awsidentity" "github.com/coder/coder/v2/coderd/azureidentity" - "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/httpapi" - "github.com/coder/coder/v2/coderd/provisionerdserver" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" ) @@ -136,120 +132,61 @@ func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, in systemCtx := dbauthz.AsSystemRestricted(ctx) agentName = strings.TrimSpace(agentName) - agents, err := api.Database.GetWorkspaceAgentsByInstanceID(systemCtx, instanceID) + agents, err := api.Database.GetWorkspaceBuildAgentsByInstanceID(systemCtx, instanceID) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Internal error fetching provisioner job agent.", + Message: "Internal error fetching workspace agent.", Detail: err.Error(), }) return } - - // Template version agents can share an instance ID with workspace build - // agents. Keep only workspace build agents before resolving ambiguity so - // template version agents do not force CODER_AGENT_NAME. - // - // We attach the provisioner job to each candidate during the filter - // loop so the post-selection code below can read it directly from the - // chosen candidate instead of re-querying. The previous code re-fetched - // the resource and job for the surviving agent, firing the - // resource->job->build->workspace dbauthz cascade twice and saturating - // the pgx pool under load. - type instanceCandidate struct { - agent database.WorkspaceAgent - job database.ProvisionerJob - } - buildCandidates := make([]instanceCandidate, 0, len(agents)) - for _, candidate := range agents { - resource, err := api.Database.GetWorkspaceResourceByID(systemCtx, candidate.ResourceID) - if err != nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Internal error fetching provisioner job resource.", - Detail: err.Error(), - }) - return - } - job, err := api.Database.GetProvisionerJobByID(systemCtx, resource.JobID) - if err != nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Internal error fetching provisioner job.", - Detail: err.Error(), - }) - return - } - if job.Type == database.ProvisionerJobTypeWorkspaceBuild { - buildCandidates = append(buildCandidates, instanceCandidate{ - agent: candidate, - job: job, - }) - } - } - if len(buildCandidates) == 0 { + if len(agents) == 0 { httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ Message: fmt.Sprintf("Instance with id %q not found.", instanceID), }) return } - var selected instanceCandidate + selected := agents[0] if agentName != "" { - for _, candidate := range buildCandidates { - if candidate.agent.Name == agentName { + found := false + for _, candidate := range agents { + if candidate.WorkspaceAgent.Name == agentName { selected = candidate + found = true break } } - if selected.agent.ID == uuid.Nil { + if !found { httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ Message: fmt.Sprintf("No agent found with instance ID %q and name %q.", instanceID, agentName), }) return } - } else { - if len(buildCandidates) != 1 { - // Include agent names in the error message to help operators - // configure CODER_AGENT_NAME. The caller has already proven - // cloud instance identity, so agent names are not sensitive - // here. - names := make([]string, len(buildCandidates)) - for i, candidate := range buildCandidates { - names[i] = candidate.agent.Name - } - sort.Strings(names) - httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ - Message: fmt.Sprintf( - "Multiple agents found with instance ID %q. Set CODER_AGENT_NAME to one of: %s", - instanceID, - strings.Join(names, ", "), - ), - }) - return + } else if len(agents) != 1 { + // Include agent names in the error message to help operators + // configure CODER_AGENT_NAME. The caller has already proven + // cloud instance identity, so agent names are not sensitive + // here. + names := make([]string, len(agents)) + for i, candidate := range agents { + names[i] = candidate.WorkspaceAgent.Name } - selected = buildCandidates[0] - } - agent := selected.agent - job := selected.job - var jobData provisionerdserver.WorkspaceProvisionJob - err = json.Unmarshal(job.Input, &jobData) - if err != nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Internal error extracting job data.", - Detail: err.Error(), - }) - return - } - resourceHistory, err := api.Database.GetWorkspaceBuildByID(systemCtx, jobData.WorkspaceBuildID) - if err != nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Internal error fetching workspace build.", - Detail: err.Error(), + sort.Strings(names) + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: fmt.Sprintf( + "Multiple agents found with instance ID %q. Set CODER_AGENT_NAME to one of: %s", + instanceID, + strings.Join(names, ", "), + ), }) return } + agent := selected.WorkspaceAgent // This token should only be exchanged if the instance ID is valid // for the latest history. If an instance ID is recycled by a cloud, // we'd hate to leak access to a user's workspace. - latestHistory, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(systemCtx, resourceHistory.WorkspaceID) + latestHistory, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(systemCtx, selected.WorkspaceTable.ID) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error fetching the latest workspace build.", @@ -257,7 +194,7 @@ func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, in }) return } - if latestHistory.ID != resourceHistory.ID { + if latestHistory.ID != selected.WorkspaceBuildID { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: fmt.Sprintf("Resource found for id %q, but isn't registered on the latest history.", instanceID), }) diff --git a/coderd/workspaceresourceauth_test.go b/coderd/workspaceresourceauth_test.go index 0b95b267a0..1f8004d1b3 100644 --- a/coderd/workspaceresourceauth_test.go +++ b/coderd/workspaceresourceauth_test.go @@ -91,6 +91,50 @@ func TestPostWorkspaceAuthAWSInstanceIdentity(t *testing.T) { require.NoError(t, err) }) + t.Run("RecycledInstanceID", func(t *testing.T) { + t.Parallel() + + instanceID := newTestInstanceID(t) + certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID) + setup := setupInstanceIDWorkspaceWithResources(t, &coderdtest.Options{ + AWSCertificates: certificates, + }, workspaceAgentsForInstanceID(instanceID, "dev")) + + successorVersion := coderdtest.CreateTemplateVersion(t, setup.client, setup.user.OrganizationID, &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionGraph: []*proto.Response{{ + Type: &proto.Response_Graph{ + Graph: &proto.GraphComplete{ + Resources: []*proto.Resource{{ + Name: "resource", + Type: "instance", + Agents: workspaceAgentsForInstanceID(newTestInstanceID(t), "dev"), + }}, + }, + }, + }}, + }, func(req *codersdk.CreateTemplateVersionRequest) { + req.TemplateID = setup.template.ID + }) + coderdtest.AwaitTemplateVersionJobCompleted(t, setup.client, successorVersion.ID) + build := coderdtest.CreateWorkspaceBuild(t, setup.client, setup.workspace, database.WorkspaceTransitionStart, func(req *codersdk.CreateWorkspaceBuildRequest) { + req.TemplateVersionID = successorVersion.ID + }) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, setup.client, build.ID) + + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancel() + + agentClient := agentsdk.New(setup.client.URL, agentsdk.WithAWSInstanceIdentity()) + agentClient.SDK.HTTPClient = metadataClient + + err := agentClient.RefreshToken(ctx) + var apiErr *codersdk.Error + require.ErrorAs(t, err, &apiErr) + require.Equal(t, http.StatusBadRequest, apiErr.StatusCode()) + require.Contains(t, apiErr.Message, "isn't registered on the latest history") + }) + t.Run("Ambiguous/MultipleAgentsNoSelector", func(t *testing.T) { t.Parallel() @@ -126,35 +170,11 @@ func TestPostWorkspaceAuthAWSInstanceIdentity(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - signatureReq, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://169.254.169.254/latest/dynamic/instance-identity/signature", nil) - require.NoError(t, err) - signatureRes, err := metadataClient.Do(signatureReq) - require.NoError(t, err) - defer signatureRes.Body.Close() - signature, err := io.ReadAll(signatureRes.Body) - require.NoError(t, err) - - documentReq, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://169.254.169.254/latest/dynamic/instance-identity/document", nil) - require.NoError(t, err) - documentRes, err := metadataClient.Do(documentReq) - require.NoError(t, err) - defer documentRes.Body.Close() - document, err := io.ReadAll(documentRes.Body) - require.NoError(t, err) - - reqBody, err := json.Marshal(map[string]string{ - "signature": string(signature), - "document": string(document), - "agent_name": "", - }) - require.NoError(t, err) - - res, err := client.RequestWithoutSessionToken(ctx, http.MethodPost, "/api/v2/workspaceagents/aws-instance-identity", reqBody) - require.NoError(t, err) + res := postAWSInstanceIdentity(ctx, t, client, metadataClient, "") defer res.Body.Close() require.Equal(t, http.StatusConflict, res.StatusCode) - err = codersdk.ReadBodyAsError(res) + err := codersdk.ReadBodyAsError(res) var apiErr *codersdk.Error require.ErrorAs(t, err, &apiErr) require.Equal(t, http.StatusConflict, apiErr.StatusCode()) @@ -174,35 +194,11 @@ func TestPostWorkspaceAuthAWSInstanceIdentity(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - signatureReq, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://169.254.169.254/latest/dynamic/instance-identity/signature", nil) - require.NoError(t, err) - signatureRes, err := metadataClient.Do(signatureReq) - require.NoError(t, err) - defer signatureRes.Body.Close() - signature, err := io.ReadAll(signatureRes.Body) - require.NoError(t, err) - - documentReq, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://169.254.169.254/latest/dynamic/instance-identity/document", nil) - require.NoError(t, err) - documentRes, err := metadataClient.Do(documentReq) - require.NoError(t, err) - defer documentRes.Body.Close() - document, err := io.ReadAll(documentRes.Body) - require.NoError(t, err) - - reqBody, err := json.Marshal(map[string]string{ - "signature": string(signature), - "document": string(document), - "agent_name": " ", - }) - require.NoError(t, err) - - res, err := client.RequestWithoutSessionToken(ctx, http.MethodPost, "/api/v2/workspaceagents/aws-instance-identity", reqBody) - require.NoError(t, err) + res := postAWSInstanceIdentity(ctx, t, client, metadataClient, " ") defer res.Body.Close() require.Equal(t, http.StatusConflict, res.StatusCode) - err = codersdk.ReadBodyAsError(res) + err := codersdk.ReadBodyAsError(res) var apiErr *codersdk.Error require.ErrorAs(t, err, &apiErr) require.Equal(t, http.StatusConflict, apiErr.StatusCode()) @@ -368,9 +364,28 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) { }) } +type instanceIDWorkspaceSetup struct { + client *codersdk.Client + store database.Store + user codersdk.CreateFirstUserResponse + template codersdk.Template + workspace codersdk.Workspace +} + func setupInstanceIDWorkspace(t *testing.T, opts *coderdtest.Options, agents []*proto.Agent) (*codersdk.Client, database.Store) { t.Helper() + setup := setupInstanceIDWorkspaceWithResources(t, opts, agents) + return setup.client, setup.store +} + +func setupInstanceIDWorkspaceWithResources( + t *testing.T, + opts *coderdtest.Options, + agents []*proto.Agent, +) instanceIDWorkspaceSetup { + t.Helper() + actualOpts := &coderdtest.Options{} if opts != nil { *actualOpts = *opts @@ -398,7 +413,13 @@ func setupInstanceIDWorkspace(t *testing.T, opts *coderdtest.Options, agents []* workspace := coderdtest.CreateWorkspace(t, client, template.ID) coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) - return client, store + return instanceIDWorkspaceSetup{ + client: client, + store: store, + user: user, + template: template, + workspace: workspace, + } } func workspaceAgentsForInstanceID(instanceID string, names ...string) []*proto.Agent { @@ -427,6 +448,59 @@ func requireWorkspaceAgentByInstanceIDAndName(t testing.TB, store database.Store return database.WorkspaceAgent{} } +const awsInstanceIdentityMetadataURL = "http://169.254.169.254/latest/dynamic/instance-identity" + +func postAWSInstanceIdentity( + ctx context.Context, + t testing.TB, + client *codersdk.Client, + metadataClient *http.Client, + agentName string, +) *http.Response { + t.Helper() + + signature := readAWSInstanceMetadata(ctx, t, metadataClient, "signature") + document := readAWSInstanceMetadata(ctx, t, metadataClient, "document") + reqBody, err := json.Marshal(map[string]string{ + "signature": signature, + "document": document, + "agent_name": agentName, + }) + require.NoError(t, err) + + res, err := client.RequestWithoutSessionToken( + ctx, + http.MethodPost, + "/api/v2/workspaceagents/aws-instance-identity", + reqBody, + ) + require.NoError(t, err) + return res +} + +func readAWSInstanceMetadata( + ctx context.Context, + t testing.TB, + metadataClient *http.Client, + path string, +) string { + t.Helper() + + req, err := http.NewRequestWithContext( + ctx, + http.MethodGet, + awsInstanceIdentityMetadataURL+"/"+path, + nil, + ) + require.NoError(t, err) + res, err := metadataClient.Do(req) + require.NoError(t, err) + defer res.Body.Close() + body, err := io.ReadAll(res.Body) + require.NoError(t, err) + return string(body) +} + func newTestInstanceID(t testing.TB) string { t.Helper() return fmt.Sprintf("instance-%d", time.Now().UnixNano())