From 0723dd3abf702279aabc3fbfc6cb026bd4baaa70 Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Thu, 14 Mar 2024 12:27:32 -0400 Subject: [PATCH] fix: ensure agent token is from latest build in middleware (#12443) --- coderd/coderd.go | 2 +- coderd/database/dbauthz/dbauthz.go | 6 +- coderd/database/dbauthz/dbauthz_test.go | 2 +- coderd/database/dbmem/dbmem.go | 78 +++++------ coderd/database/dbmetrics/dbmetrics.go | 6 +- coderd/database/dbmock/dbmock.go | 14 +- coderd/database/querier.go | 2 +- coderd/database/queries.sql.go | 136 ++++++++++---------- coderd/database/queries/workspaceagents.sql | 78 ++++------- coderd/httpmw/workspaceagent.go | 49 +++++-- coderd/httpmw/workspaceagent_test.go | 50 ++++++- coderd/workspaceagents.go | 8 +- coderd/workspaceagents_test.go | 4 +- coderd/workspaceagentsrpc.go | 65 ++-------- enterprise/coderd/coderd.go | 2 +- 15 files changed, 242 insertions(+), 260 deletions(-) diff --git a/coderd/coderd.go b/coderd/coderd.go index 759c9450d9..a06116c8b0 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -902,7 +902,7 @@ func New(options *Options) *API { httpmw.RequireAPIKeyOrWorkspaceProxyAuth(), ).Get("/connection", api.workspaceAgentConnectionGeneric) r.Route("/me", func(r chi.Router) { - r.Use(httpmw.ExtractWorkspaceAgent(httpmw.ExtractWorkspaceAgentConfig{ + r.Use(httpmw.ExtractWorkspaceAgentAndLatestBuild(httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{ DB: options.Database, Optional: false, })) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 76f497e1e3..2dd15e073b 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1880,12 +1880,12 @@ func (q *querier) GetUsersByIDs(ctx context.Context, ids []uuid.UUID) ([]databas return q.db.GetUsersByIDs(ctx, ids) } -func (q *querier) GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) { +func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { // This is a system function if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceSystem); err != nil { - return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, err + return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, err } - return q.db.GetWorkspaceAgentAndOwnerByAuthToken(ctx, authToken) + return q.db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken) } func (q *querier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 65b18b01fb..a19512797f 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -2274,7 +2274,7 @@ func (s *MethodTestSuite) TestSystemFunctions() { s.Run("GetReplicaByID", s.Subtest(func(db database.Store, check *expects) { check.Args(uuid.New()).Asserts(rbac.ResourceSystem, rbac.ActionRead).Errors(sql.ErrNoRows) })) - s.Run("GetWorkspaceAgentAndOwnerByAuthToken", s.Subtest(func(db database.Store, check *expects) { + s.Run("GetWorkspaceAgentAndLatestBuildByAuthToken", s.Subtest(func(db database.Store, check *expects) { check.Args(uuid.New()).Asserts(rbac.ResourceSystem, rbac.ActionRead).Errors(sql.ErrNoRows) })) s.Run("GetUserLinksByUserID", s.Subtest(func(db database.Store, check *expects) { diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go index 898d257790..9c78858e86 100644 --- a/coderd/database/dbmem/dbmem.go +++ b/coderd/database/dbmem/dbmem.go @@ -69,7 +69,7 @@ func New() database.Store { templates: make([]database.TemplateTable, 0), workspaceAgentStats: make([]database.WorkspaceAgentStat, 0), workspaceAgentLogs: make([]database.WorkspaceAgentLog, 0), - workspaceBuilds: make([]database.WorkspaceBuildTable, 0), + workspaceBuilds: make([]database.WorkspaceBuild, 0), workspaceApps: make([]database.WorkspaceApp, 0), workspaces: make([]database.Workspace, 0), licenses: make([]database.License, 0), @@ -171,7 +171,7 @@ type data struct { workspaceApps []database.WorkspaceApp workspaceAppStatsLastInsertID int64 workspaceAppStats []database.WorkspaceAppStat - workspaceBuilds []database.WorkspaceBuildTable + workspaceBuilds []database.WorkspaceBuild workspaceBuildParameters []database.WorkspaceBuildParameter workspaceResourceMetadata []database.WorkspaceResourceMetadatum workspaceResources []database.WorkspaceResource @@ -542,7 +542,7 @@ func (q *FakeQuerier) templateVersionWithUserNoLock(tpl database.TemplateVersion return withUser } -func (q *FakeQuerier) workspaceBuildWithUserNoLock(tpl database.WorkspaceBuildTable) database.WorkspaceBuild { +func (q *FakeQuerier) workspaceBuildWithUserNoLock(tpl database.WorkspaceBuild) database.WorkspaceBuild { var user database.User for _, _user := range q.users { if _user.ID == tpl.InitiatorID { @@ -2801,7 +2801,7 @@ func (q *FakeQuerier) GetQuotaConsumedForUser(_ context.Context, userID uuid.UUI continue } - var lastBuild database.WorkspaceBuildTable + var lastBuild database.WorkspaceBuild for _, build := range q.workspaceBuilds { if build.WorkspaceID != workspace.ID { continue @@ -3488,7 +3488,7 @@ func (q *FakeQuerier) GetTemplateParameterInsights(ctx context.Context, arg data defer q.mutex.RUnlock() // WITH latest_workspace_builds ... - latestWorkspaceBuilds := make(map[uuid.UUID]database.WorkspaceBuildTable) + latestWorkspaceBuilds := make(map[uuid.UUID]database.WorkspaceBuild) for _, wb := range q.workspaceBuilds { if wb.CreatedAt.Before(arg.StartTime) || wb.CreatedAt.Equal(arg.EndTime) || wb.CreatedAt.After(arg.EndTime) { continue @@ -4270,20 +4270,14 @@ func (q *FakeQuerier) GetUsersByIDs(_ context.Context, ids []uuid.UUID) ([]datab return users, nil } -func (q *FakeQuerier) GetWorkspaceAgentAndOwnerByAuthToken(_ context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) { +func (q *FakeQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(_ context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { q.mutex.RLock() defer q.mutex.RUnlock() - - // map of build number -> row - rows := make(map[int32]database.GetWorkspaceAgentAndOwnerByAuthTokenRow) - - // We want to return the latest build number - var latestBuildNumber int32 + rows := []database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{} + // We want to return the latest build number for each workspace + latestBuildNumber := make(map[uuid.UUID]int32) for _, agt := range q.workspaceAgents { - if agt.AuthToken != authToken { - continue - } // get the related workspace and user for _, res := range q.workspaceResources { if agt.ResourceID != res.ID { @@ -4300,47 +4294,43 @@ func (q *FakeQuerier) GetWorkspaceAgentAndOwnerByAuthToken(_ context.Context, au if ws.Deleted { continue } - var row database.GetWorkspaceAgentAndOwnerByAuthTokenRow - row.WorkspaceID = ws.ID - row.TemplateID = ws.TemplateID + row := database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{ + Workspace: database.Workspace{ + ID: ws.ID, + TemplateID: ws.TemplateID, + }, + WorkspaceAgent: agt, + WorkspaceBuild: build, + } usr, err := q.getUserByIDNoLock(ws.OwnerID) if err != nil { - return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, sql.ErrNoRows - } - row.OwnerID = usr.ID - row.OwnerRoles = append(usr.RBACRoles, "member") - // We also need to get org roles for the user - row.OwnerName = usr.Username - row.WorkspaceAgent = agt - row.TemplateVersionID = build.TemplateVersionID - for _, mem := range q.organizationMembers { - if mem.UserID == usr.ID { - row.OwnerRoles = append(row.OwnerRoles, fmt.Sprintf("organization-member:%s", mem.OrganizationID.String())) - } - } - // And group memberships - for _, groupMem := range q.groupMembers { - if groupMem.UserID == usr.ID { - row.OwnerGroups = append(row.OwnerGroups, groupMem.GroupID.String()) - } + return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, sql.ErrNoRows } + row.Workspace.OwnerID = usr.ID // Keep track of the latest build number - rows[build.BuildNumber] = row - if build.BuildNumber > latestBuildNumber { - latestBuildNumber = build.BuildNumber + rows = append(rows, row) + if build.BuildNumber > latestBuildNumber[ws.ID] { + latestBuildNumber[ws.ID] = build.BuildNumber } } } } } - if len(rows) == 0 { - return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, sql.ErrNoRows + for i := range rows { + if rows[i].WorkspaceAgent.AuthToken != authToken { + continue + } + + if rows[i].WorkspaceBuild.BuildNumber != latestBuildNumber[rows[i].Workspace.ID] { + continue + } + + return rows[i], nil } - // Return the row related to the latest build - return rows[latestBuildNumber], nil + return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, sql.ErrNoRows } func (q *FakeQuerier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) { @@ -6243,7 +6233,7 @@ func (q *FakeQuerier) InsertWorkspaceBuild(_ context.Context, arg database.Inser q.mutex.Lock() defer q.mutex.Unlock() - workspaceBuild := database.WorkspaceBuildTable{ + workspaceBuild := database.WorkspaceBuild{ ID: arg.ID, CreatedAt: arg.CreatedAt, UpdatedAt: arg.UpdatedAt, diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go index bbd67cf064..2350588511 100644 --- a/coderd/database/dbmetrics/dbmetrics.go +++ b/coderd/database/dbmetrics/dbmetrics.go @@ -1103,10 +1103,10 @@ func (m metricsStore) GetUsersByIDs(ctx context.Context, ids []uuid.UUID) ([]dat return users, err } -func (m metricsStore) GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) { +func (m metricsStore) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { start := time.Now() - r0, r1 := m.s.GetWorkspaceAgentAndOwnerByAuthToken(ctx, authToken) - m.queryLatencies.WithLabelValues("GetWorkspaceAgentAndOwnerByAuthToken").Observe(time.Since(start).Seconds()) + r0, r1 := m.s.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken) + m.queryLatencies.WithLabelValues("GetWorkspaceAgentAndLatestBuildByAuthToken").Observe(time.Since(start).Seconds()) return r0, r1 } diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 8ef54b213e..ba5b3c965c 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -2295,19 +2295,19 @@ func (mr *MockStoreMockRecorder) GetUsersByIDs(arg0, arg1 any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUsersByIDs", reflect.TypeOf((*MockStore)(nil).GetUsersByIDs), arg0, arg1) } -// GetWorkspaceAgentAndOwnerByAuthToken mocks base method. -func (m *MockStore) GetWorkspaceAgentAndOwnerByAuthToken(arg0 context.Context, arg1 uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) { +// GetWorkspaceAgentAndLatestBuildByAuthToken mocks base method. +func (m *MockStore) GetWorkspaceAgentAndLatestBuildByAuthToken(arg0 context.Context, arg1 uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetWorkspaceAgentAndOwnerByAuthToken", arg0, arg1) - ret0, _ := ret[0].(database.GetWorkspaceAgentAndOwnerByAuthTokenRow) + ret := m.ctrl.Call(m, "GetWorkspaceAgentAndLatestBuildByAuthToken", arg0, arg1) + ret0, _ := ret[0].(database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow) ret1, _ := ret[1].(error) return ret0, ret1 } -// GetWorkspaceAgentAndOwnerByAuthToken indicates an expected call of GetWorkspaceAgentAndOwnerByAuthToken. -func (mr *MockStoreMockRecorder) GetWorkspaceAgentAndOwnerByAuthToken(arg0, arg1 any) *gomock.Call { +// GetWorkspaceAgentAndLatestBuildByAuthToken indicates an expected call of GetWorkspaceAgentAndLatestBuildByAuthToken. +func (mr *MockStoreMockRecorder) GetWorkspaceAgentAndLatestBuildByAuthToken(arg0, arg1 any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentAndOwnerByAuthToken", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentAndOwnerByAuthToken), arg0, arg1) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentAndLatestBuildByAuthToken", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentAndLatestBuildByAuthToken), arg0, arg1) } // GetWorkspaceAgentByID mocks base method. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index b0c48f5275..fc47f58204 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -230,7 +230,7 @@ type sqlcQuerier interface { // to look up references to actions. eg. a user could build a workspace // for another user, then be deleted... we still want them to appear! GetUsersByIDs(ctx context.Context, ids []uuid.UUID) ([]User, error) - GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndOwnerByAuthTokenRow, error) + GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (WorkspaceAgent, error) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (WorkspaceAgent, error) GetWorkspaceAgentLifecycleStateByID(ctx context.Context, id uuid.UUID) (GetWorkspaceAgentLifecycleStateByIDRow, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 888515f49f..d5bcfdd40f 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -8671,80 +8671,66 @@ func (q *sqlQuerier) DeleteOldWorkspaceAgentLogs(ctx context.Context) error { return err } -const getWorkspaceAgentAndOwnerByAuthToken = `-- name: GetWorkspaceAgentAndOwnerByAuthToken :one +const getWorkspaceAgentAndLatestBuildByAuthToken = `-- name: GetWorkspaceAgentAndLatestBuildByAuthToken :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, 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, - workspaces.id AS workspace_id, - users.id AS owner_id, - users.username AS owner_name, - users.status AS owner_status, - workspaces.template_id AS template_id, - workspace_builds.template_version_id AS template_version_id, - array_cat( - array_append(users.rbac_roles, 'member'), - array_append(ARRAY[]::text[], 'organization-member:' || organization_members.organization_id::text) - )::text[] as owner_roles, - array_agg(COALESCE(group_members.group_id::text, ''))::text[] AS owner_groups -FROM users - INNER JOIN - workspaces - ON - workspaces.owner_id = users.id - INNER JOIN - workspace_builds - ON - workspace_builds.workspace_id = workspaces.id - INNER JOIN - workspace_resources - ON - workspace_resources.job_id = workspace_builds.job_id - INNER JOIN - workspace_agents - ON - workspace_agents.resource_id = workspace_resources.id - INNER JOIN -- every user is a member of some org - organization_members - ON - organization_members.user_id = users.id - LEFT JOIN -- as they may not be a member of any groups - group_members - ON - group_members.user_id = users.id + workspace_build_with_user.id, workspace_build_with_user.created_at, workspace_build_with_user.updated_at, workspace_build_with_user.workspace_id, workspace_build_with_user.template_version_id, workspace_build_with_user.build_number, workspace_build_with_user.transition, workspace_build_with_user.initiator_id, workspace_build_with_user.provisioner_state, workspace_build_with_user.job_id, workspace_build_with_user.deadline, workspace_build_with_user.reason, workspace_build_with_user.daily_cost, workspace_build_with_user.max_deadline, workspace_build_with_user.initiator_by_avatar_url, workspace_build_with_user.initiator_by_username +FROM + -- Only get the latest build for each workspace + ( + SELECT + workspace_id, MAX(build_number) as max_build_number + FROM + workspace_build_with_user + GROUP BY + workspace_id + ) as latest_builds + -- Pull the workspace_build rows for returning +INNER JOIN workspace_build_with_user + ON workspace_build_with_user.workspace_id = latest_builds.workspace_id + AND workspace_build_with_user.build_number = latest_builds.max_build_number + -- For each latest build, grab the resources to relate to an agent +INNER JOIN workspace_resources + ON workspace_resources.job_id = workspace_build_with_user.job_id + -- Agent <-> Resource is 1:1 +INNER JOIN workspace_agents + ON workspace_agents.resource_id = workspace_resources.id + -- We need the owner ID +INNER JOIN workspaces + ON workspace_build_with_user.workspace_id = workspaces.id WHERE - -- TODO: we can add more conditions here, such as: - -- 1) The user must be active - -- 2) The workspace must be running + -- This should only match 1 agent, so 1 returned row or 0 workspace_agents.auth_token = $1 AND workspaces.deleted = FALSE -GROUP BY - workspace_agents.id, - workspaces.id, - users.id, - organization_members.organization_id, - workspace_builds.build_number, - workspace_builds.template_version_id -ORDER BY - workspace_builds.build_number DESC -LIMIT 1 ` -type GetWorkspaceAgentAndOwnerByAuthTokenRow struct { - WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"` - WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"` - OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` - OwnerName string `db:"owner_name" json:"owner_name"` - OwnerStatus UserStatus `db:"owner_status" json:"owner_status"` - TemplateID uuid.UUID `db:"template_id" json:"template_id"` - TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"` - OwnerRoles []string `db:"owner_roles" json:"owner_roles"` - OwnerGroups []string `db:"owner_groups" json:"owner_groups"` +type GetWorkspaceAgentAndLatestBuildByAuthTokenRow struct { + Workspace Workspace `db:"workspace" json:"workspace"` + WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"` + WorkspaceBuild WorkspaceBuild `db:"workspace_build" json:"workspace_build"` } -func (q *sqlQuerier) GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndOwnerByAuthTokenRow, error) { - row := q.db.QueryRowContext(ctx, getWorkspaceAgentAndOwnerByAuthToken, authToken) - var i GetWorkspaceAgentAndOwnerByAuthTokenRow +func (q *sqlQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) { + row := q.db.QueryRowContext(ctx, getWorkspaceAgentAndLatestBuildByAuthToken, authToken) + var i GetWorkspaceAgentAndLatestBuildByAuthTokenRow err := row.Scan( + &i.Workspace.ID, + &i.Workspace.CreatedAt, + &i.Workspace.UpdatedAt, + &i.Workspace.OwnerID, + &i.Workspace.OrganizationID, + &i.Workspace.TemplateID, + &i.Workspace.Deleted, + &i.Workspace.Name, + &i.Workspace.AutostartSchedule, + &i.Workspace.Ttl, + &i.Workspace.LastUsedAt, + &i.Workspace.DormantAt, + &i.Workspace.DeletingAt, + &i.Workspace.AutomaticUpdates, + &i.Workspace.Favorite, &i.WorkspaceAgent.ID, &i.WorkspaceAgent.CreatedAt, &i.WorkspaceAgent.UpdatedAt, @@ -8776,14 +8762,22 @@ func (q *sqlQuerier) GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, a pq.Array(&i.WorkspaceAgent.DisplayApps), &i.WorkspaceAgent.APIVersion, &i.WorkspaceAgent.DisplayOrder, - &i.WorkspaceID, - &i.OwnerID, - &i.OwnerName, - &i.OwnerStatus, - &i.TemplateID, - &i.TemplateVersionID, - pq.Array(&i.OwnerRoles), - pq.Array(&i.OwnerGroups), + &i.WorkspaceBuild.ID, + &i.WorkspaceBuild.CreatedAt, + &i.WorkspaceBuild.UpdatedAt, + &i.WorkspaceBuild.WorkspaceID, + &i.WorkspaceBuild.TemplateVersionID, + &i.WorkspaceBuild.BuildNumber, + &i.WorkspaceBuild.Transition, + &i.WorkspaceBuild.InitiatorID, + &i.WorkspaceBuild.ProvisionerState, + &i.WorkspaceBuild.JobID, + &i.WorkspaceBuild.Deadline, + &i.WorkspaceBuild.Reason, + &i.WorkspaceBuild.DailyCost, + &i.WorkspaceBuild.MaxDeadline, + &i.WorkspaceBuild.InitiatorByAvatarUrl, + &i.WorkspaceBuild.InitiatorByUsername, ) return i, err } diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql index d1021f866b..e1ed91f221 100644 --- a/coderd/database/queries/workspaceagents.sql +++ b/coderd/database/queries/workspaceagents.sql @@ -214,59 +214,37 @@ WHERE wb.workspace_id = @workspace_id :: uuid ); --- name: GetWorkspaceAgentAndOwnerByAuthToken :one +-- name: GetWorkspaceAgentAndLatestBuildByAuthToken :one SELECT + sqlc.embed(workspaces), sqlc.embed(workspace_agents), - workspaces.id AS workspace_id, - users.id AS owner_id, - users.username AS owner_name, - users.status AS owner_status, - workspaces.template_id AS template_id, - workspace_builds.template_version_id AS template_version_id, - array_cat( - array_append(users.rbac_roles, 'member'), - array_append(ARRAY[]::text[], 'organization-member:' || organization_members.organization_id::text) - )::text[] as owner_roles, - array_agg(COALESCE(group_members.group_id::text, ''))::text[] AS owner_groups -FROM users - INNER JOIN - workspaces - ON - workspaces.owner_id = users.id - INNER JOIN - workspace_builds - ON - workspace_builds.workspace_id = workspaces.id - INNER JOIN - workspace_resources - ON - workspace_resources.job_id = workspace_builds.job_id - INNER JOIN - workspace_agents - ON - workspace_agents.resource_id = workspace_resources.id - INNER JOIN -- every user is a member of some org - organization_members - ON - organization_members.user_id = users.id - LEFT JOIN -- as they may not be a member of any groups - group_members - ON - group_members.user_id = users.id + sqlc.embed(workspace_build_with_user) +FROM + -- Only get the latest build for each workspace + ( + SELECT + workspace_id, MAX(build_number) as max_build_number + FROM + workspace_build_with_user + GROUP BY + workspace_id + ) as latest_builds + -- Pull the workspace_build rows for returning +INNER JOIN workspace_build_with_user + ON workspace_build_with_user.workspace_id = latest_builds.workspace_id + AND workspace_build_with_user.build_number = latest_builds.max_build_number + -- For each latest build, grab the resources to relate to an agent +INNER JOIN workspace_resources + ON workspace_resources.job_id = workspace_build_with_user.job_id + -- Agent <-> Resource is 1:1 +INNER JOIN workspace_agents + ON workspace_agents.resource_id = workspace_resources.id + -- We need the owner ID +INNER JOIN workspaces + ON workspace_build_with_user.workspace_id = workspaces.id WHERE - -- TODO: we can add more conditions here, such as: - -- 1) The user must be active - -- 2) The workspace must be running + -- This should only match 1 agent, so 1 returned row or 0 workspace_agents.auth_token = @auth_token AND workspaces.deleted = FALSE -GROUP BY - workspace_agents.id, - workspaces.id, - users.id, - organization_members.organization_id, - workspace_builds.build_number, - workspace_builds.template_version_id -ORDER BY - workspace_builds.build_number DESC -LIMIT 1; +; diff --git a/coderd/httpmw/workspaceagent.go b/coderd/httpmw/workspaceagent.go index 31caad0b7c..a72d05caec 100644 --- a/coderd/httpmw/workspaceagent.go +++ b/coderd/httpmw/workspaceagent.go @@ -32,7 +32,23 @@ func WorkspaceAgent(r *http.Request) database.WorkspaceAgent { return user } -type ExtractWorkspaceAgentConfig struct { +type latestBuildContextKey struct{} + +func latestBuildOptional(r *http.Request) (database.WorkspaceBuild, bool) { + wb, ok := r.Context().Value(latestBuildContextKey{}).(database.WorkspaceBuild) + return wb, ok +} + +// LatestBuild returns the Latest Build from the ExtractLatestBuild handler. +func LatestBuild(r *http.Request) database.WorkspaceBuild { + wb, ok := latestBuildOptional(r) + if !ok { + panic("developer error: agent middleware not provided or was made optional") + } + return wb +} + +type ExtractWorkspaceAgentAndLatestBuildConfig struct { DB database.Store // Optional indicates whether the middleware should be optional. If true, any // requests without the a token or with an invalid token will be allowed to @@ -40,8 +56,8 @@ type ExtractWorkspaceAgentConfig struct { Optional bool } -// ExtractWorkspaceAgent requires authentication using a valid agent token. -func ExtractWorkspaceAgent(opts ExtractWorkspaceAgentConfig) func(http.Handler) http.Handler { +// ExtractWorkspaceAgentAndLatestBuild requires authentication using a valid agent token. +func ExtractWorkspaceAgentAndLatestBuild(opts ExtractWorkspaceAgentAndLatestBuildConfig) func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() @@ -76,7 +92,7 @@ func ExtractWorkspaceAgent(opts ExtractWorkspaceAgentConfig) func(http.Handler) } //nolint:gocritic // System needs to be able to get workspace agents. - row, err := opts.DB.GetWorkspaceAgentAndOwnerByAuthToken(dbauthz.AsSystemRestricted(ctx), token) + row, err := opts.DB.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), token) if err != nil { if errors.Is(err, sql.ErrNoRows) { optionalWrite(http.StatusUnauthorized, codersdk.Response{ @@ -93,19 +109,30 @@ func ExtractWorkspaceAgent(opts ExtractWorkspaceAgentConfig) func(http.Handler) return } + //nolint:gocritic // System needs to be able to get owner roles. + roles, err := opts.DB.GetAuthorizationUserRoles(dbauthz.AsSystemRestricted(ctx), row.Workspace.OwnerID) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Internal error checking workspace agent authorization.", + Detail: err.Error(), + }) + return + } + subject := rbac.Subject{ - ID: row.OwnerID.String(), - Roles: rbac.RoleNames(row.OwnerRoles), - Groups: row.OwnerGroups, + ID: row.Workspace.OwnerID.String(), + Roles: rbac.RoleNames(roles.Roles), + Groups: roles.Groups, Scope: rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{ - WorkspaceID: row.WorkspaceID, - OwnerID: row.OwnerID, - TemplateID: row.TemplateID, - VersionID: row.TemplateVersionID, + WorkspaceID: row.Workspace.ID, + OwnerID: row.Workspace.OwnerID, + TemplateID: row.Workspace.TemplateID, + VersionID: row.WorkspaceBuild.TemplateVersionID, }), }.WithCachedASTValue() ctx = context.WithValue(ctx, workspaceAgentContextKey{}, row.WorkspaceAgent) + ctx = context.WithValue(ctx, latestBuildContextKey{}, row.WorkspaceBuild) // Also set the dbauthz actor for the request. ctx = dbauthz.As(ctx, subject) next.ServeHTTP(rw, r.WithContext(ctx)) diff --git a/coderd/httpmw/workspaceagent_test.go b/coderd/httpmw/workspaceagent_test.go index 5788540628..0bc4b04a35 100644 --- a/coderd/httpmw/workspaceagent_test.go +++ b/coderd/httpmw/workspaceagent_test.go @@ -23,8 +23,8 @@ func TestWorkspaceAgent(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) - req, rtr := setup(t, db, uuid.New(), httpmw.ExtractWorkspaceAgent( - httpmw.ExtractWorkspaceAgentConfig{ + req, rtr, _, _ := setup(t, db, uuid.New(), httpmw.ExtractWorkspaceAgentAndLatestBuild( + httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{ DB: db, Optional: false, })) @@ -42,8 +42,8 @@ func TestWorkspaceAgent(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) authToken := uuid.New() - req, rtr := setup(t, db, authToken, httpmw.ExtractWorkspaceAgent( - httpmw.ExtractWorkspaceAgentConfig{ + req, rtr, _, _ := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild( + httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{ DB: db, Optional: false, })) @@ -57,9 +57,47 @@ func TestWorkspaceAgent(t *testing.T) { t.Cleanup(func() { _ = res.Body.Close() }) require.Equal(t, http.StatusOK, res.StatusCode) }) + + t.Run("Latest", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + authToken := uuid.New() + req, rtr, ws, tpv := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild( + httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{ + DB: db, + Optional: false, + }), + ) + + // Create a newer build + job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: ws.OrganizationID, + }) + resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: job.ID, + }) + _ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: ws.ID, + JobID: job.ID, + TemplateVersionID: tpv.ID, + BuildNumber: 2, + }) + _ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: resource.ID, + }) + + 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) + }) } -func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Handler) http.Handler) (*http.Request, http.Handler) { +func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Handler) http.Handler) (*http.Request, http.Handler, database.Workspace, database.TemplateVersion) { t.Helper() org := dbgen.Organization(t, db, database.Organization{}) user := dbgen.User(t, db, database.User{ @@ -107,5 +145,5 @@ func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Ha rw.WriteHeader(http.StatusOK) }) - return req, rtr + return req, rtr, workspace, templateVersion } diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index c254eb9e40..cf64d92d28 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -954,13 +954,9 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request api.WebsocketWaitGroup.Add(1) api.WebsocketWaitMutex.Unlock() defer api.WebsocketWaitGroup.Done() + // The middleware only accept agents for resources on the latest build. workspaceAgent := httpmw.WorkspaceAgent(r) - // Ensure the resource is still valid! - // We only accept agents for resources on the latest build. - build, ok := ensureLatestBuild(ctx, api.Database, api.Logger, rw, workspaceAgent) - if !ok { - return - } + build := httpmw.LatestBuild(r) workspace, err := api.Database.GetWorkspaceByID(ctx, build.WorkspaceID) if err != nil { diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 27a222b82b..68a1492e79 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -404,7 +404,7 @@ func TestWorkspaceAgentConnectRPC(t *testing.T) { require.Error(t, err) var sdkErr *codersdk.Error require.ErrorAs(t, err, &sdkErr) - require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + require.Equal(t, http.StatusUnauthorized, sdkErr.StatusCode()) }) t.Run("FailDeleted", func(t *testing.T) { @@ -488,7 +488,7 @@ func TestWorkspaceAgentClientCoordinate_BadVersion(t *testing.T) { agentToken, err := uuid.Parse(r.AgentToken) require.NoError(t, err) //nolint: gocritic // testing - ao, err := db.GetWorkspaceAgentAndOwnerByAuthToken(dbauthz.AsSystemRestricted(ctx), agentToken) + ao, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentToken) require.NoError(t, err) //nolint: bodyclose // closed by ReadBodyAsError diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go index 2656adf374..7130d0b88e 100644 --- a/coderd/workspaceagentsrpc.go +++ b/coderd/workspaceagentsrpc.go @@ -61,11 +61,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) { api.WebsocketWaitMutex.Unlock() defer api.WebsocketWaitGroup.Done() workspaceAgent := httpmw.WorkspaceAgent(r) - - build, ok := ensureLatestBuild(ctx, api.Database, logger, rw, workspaceAgent) - if !ok { - return - } + build := httpmw.LatestBuild(r) workspace, err := api.Database.GetWorkspaceByID(ctx, build.WorkspaceID) if err != nil { @@ -167,54 +163,6 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) { } } -func ensureLatestBuild(ctx context.Context, db database.Store, logger slog.Logger, rw http.ResponseWriter, workspaceAgent database.WorkspaceAgent) (database.WorkspaceBuild, bool) { - resource, err := db.GetWorkspaceResourceByID(ctx, workspaceAgent.ResourceID) - if err != nil { - httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Internal error fetching workspace agent resource.", - Detail: err.Error(), - }) - return database.WorkspaceBuild{}, false - } - - build, err := db.GetWorkspaceBuildByJobID(ctx, resource.JobID) - if err != nil { - httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Internal error fetching workspace build job.", - Detail: err.Error(), - }) - return database.WorkspaceBuild{}, false - } - - // Ensure the resource is still valid! - // We only accept agents for resources on the latest build. - err = checkBuildIsLatest(ctx, db, build) - if err != nil { - logger.Debug(ctx, "agent tried to connect from non-latest build", - slog.F("resource", resource), - slog.F("agent", workspaceAgent), - ) - httpapi.Write(ctx, rw, http.StatusForbidden, codersdk.Response{ - Message: "Agent trying to connect from non-latest build.", - Detail: err.Error(), - }) - return database.WorkspaceBuild{}, false - } - - return build, true -} - -func checkBuildIsLatest(ctx context.Context, db database.Store, build database.WorkspaceBuild) error { - latestBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, build.WorkspaceID) - if err != nil { - return err - } - if build.ID != latestBuild.ID { - return xerrors.New("build is outdated") - } - return nil -} - func (api *API) startAgentWebsocketMonitor(ctx context.Context, workspaceAgent database.WorkspaceAgent, workspaceBuild database.WorkspaceBuild, conn *websocket.Conn, @@ -494,3 +442,14 @@ func (m *agentConnectionMonitor) close() { m.cancel() m.wg.Wait() } + +func checkBuildIsLatest(ctx context.Context, db database.Store, build database.WorkspaceBuild) error { + latestBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, build.WorkspaceID) + if err != nil { + return err + } + if build.ID != latestBuild.ID { + return xerrors.New("build is outdated") + } + return nil +} diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index e68c7d8d17..b6c41badaa 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -336,7 +336,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { r.Group(func(r chi.Router) { r.Use( apiKeyMiddlewareOptional, - httpmw.ExtractWorkspaceAgent(httpmw.ExtractWorkspaceAgentConfig{ + httpmw.ExtractWorkspaceAgentAndLatestBuild(httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{ DB: options.Database, Optional: true, }),