fix: ensure agent token is from latest build in middleware (#12443)

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,
 				}))

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) {

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) {

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,

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
 }
 

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.

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)

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
 }

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;
+;

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))

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
 }

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 {

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

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
+}

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,
 					}),