Skip to content

Commit 242cf70

Browse files
johnstcnpull[bot]
johnstcn
authored and
pull[bot]
committed
refactor(coderd): fetch owner information when authorizing workspace agent (#9123)
* Refactors the existing httpmw tests to use dbtestutil so that we can test them against a real database if desired, * Modifies the GetWorkspaceAgentByAuthToken to return the owner and associated roles, removing the need for additional queries
1 parent 304f9af commit 242cf70

File tree

11 files changed

+313
-170
lines changed

11 files changed

+313
-170
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1493,13 +1493,12 @@ func (q *querier) GetUsersByIDs(ctx context.Context, ids []uuid.UUID) ([]databas
14931493
return q.db.GetUsersByIDs(ctx, ids)
14941494
}
14951495

1496-
// GetWorkspaceAgentByAuthToken is used in http middleware to get the workspace agent.
1497-
// This should only be used by a system user in that middleware.
1498-
func (q *querier) GetWorkspaceAgentByAuthToken(ctx context.Context, authToken uuid.UUID) (database.WorkspaceAgent, error) {
1496+
func (q *querier) GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) {
1497+
// This is a system function
14991498
if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceSystem); err != nil {
1500-
return database.WorkspaceAgent{}, err
1499+
return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, err
15011500
}
1502-
return q.db.GetWorkspaceAgentByAuthToken(ctx, authToken)
1501+
return q.db.GetWorkspaceAgentAndOwnerByAuthToken(ctx, authToken)
15031502
}
15041503

15051504
func (q *querier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) {

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1319,10 +1319,6 @@ func (s *MethodTestSuite) TestSystemFunctions() {
13191319
dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{})
13201320
check.Args().Asserts(rbac.ResourceSystem, rbac.ActionRead)
13211321
}))
1322-
s.Run("GetWorkspaceAgentByAuthToken", s.Subtest(func(db database.Store, check *expects) {
1323-
agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{})
1324-
check.Args(agt.AuthToken).Asserts(rbac.ResourceSystem, rbac.ActionRead).Returns(agt)
1325-
}))
13261322
s.Run("GetActiveUserCount", s.Subtest(func(db database.Store, check *expects) {
13271323
check.Args().Asserts(rbac.ResourceSystem, rbac.ActionRead).Returns(int64(0))
13281324
}))

‎coderd/database/dbfake/dbfake.go

Lines changed: 61 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2962,18 +2962,72 @@ func (q *FakeQuerier) GetUsersByIDs(_ context.Context, ids []uuid.UUID) ([]datab
29622962
return users, nil
29632963
}
29642964

2965-
func (q *FakeQuerier) GetWorkspaceAgentByAuthToken(_ context.Context, authToken uuid.UUID) (database.WorkspaceAgent, error) {
2965+
func (q *FakeQuerier) GetWorkspaceAgentAndOwnerByAuthToken(_ context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) {
29662966
q.mutex.RLock()
29672967
defer q.mutex.RUnlock()
29682968

2969-
// The schema sorts this by created at, so we iterate the array backwards.
2970-
for i := len(q.workspaceAgents) - 1; i >= 0; i-- {
2971-
agent := q.workspaceAgents[i]
2972-
if agent.AuthToken == authToken {
2973-
return agent, nil
2969+
// map of build number -> row
2970+
rows := make(map[int32]database.GetWorkspaceAgentAndOwnerByAuthTokenRow)
2971+
2972+
// We want to return the latest build number
2973+
var latestBuildNumber int32
2974+
2975+
for _, agt := range q.workspaceAgents {
2976+
if agt.AuthToken != authToken {
2977+
continue
2978+
}
2979+
// get the related workspace and user
2980+
for _, res := range q.workspaceResources {
2981+
if agt.ResourceID != res.ID {
2982+
continue
2983+
}
2984+
for _, build := range q.workspaceBuilds {
2985+
if build.JobID != res.JobID {
2986+
continue
2987+
}
2988+
for _, ws := range q.workspaces {
2989+
if build.WorkspaceID != ws.ID {
2990+
continue
2991+
}
2992+
var row database.GetWorkspaceAgentAndOwnerByAuthTokenRow
2993+
row.WorkspaceID = ws.ID
2994+
usr, err := q.getUserByIDNoLock(ws.OwnerID)
2995+
if err != nil {
2996+
return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, sql.ErrNoRows
2997+
}
2998+
row.OwnerID = usr.ID
2999+
row.OwnerRoles = append(usr.RBACRoles, "member")
3000+
// We also need to get org roles for the user
3001+
row.OwnerName = usr.Username
3002+
row.WorkspaceAgent = agt
3003+
for _, mem := range q.organizationMembers {
3004+
if mem.UserID == usr.ID {
3005+
row.OwnerRoles = append(row.OwnerRoles, fmt.Sprintf("organization-member:%s", mem.OrganizationID.String()))
3006+
}
3007+
}
3008+
// And group memberships
3009+
for _, groupMem := range q.groupMembers {
3010+
if groupMem.UserID == usr.ID {
3011+
row.OwnerGroups = append(row.OwnerGroups, groupMem.GroupID.String())
3012+
}
3013+
}
3014+
3015+
// Keep track of the latest build number
3016+
rows[build.BuildNumber] = row
3017+
if build.BuildNumber > latestBuildNumber {
3018+
latestBuildNumber = build.BuildNumber
3019+
}
3020+
}
3021+
}
29743022
}
29753023
}
2976-
return database.WorkspaceAgent{}, sql.ErrNoRows
3024+
3025+
if len(rows) == 0 {
3026+
return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, sql.ErrNoRows
3027+
}
3028+
3029+
// Return the row related to the latest build
3030+
return rows[latestBuildNumber], nil
29773031
}
29783032

29793033
func (q *FakeQuerier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) {

‎coderd/database/dbmetrics/dbmetrics.go

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go

Lines changed: 7 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

Lines changed: 101 additions & 42 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaceagents.sql

Lines changed: 53 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,3 @@
1-
-- name: GetWorkspaceAgentByAuthToken :one
2-
SELECT
3-
*
4-
FROM
5-
workspace_agents
6-
WHERE
7-
auth_token = $1
8-
ORDER BY
9-
created_at DESC;
10-
111
-- name: GetWorkspaceAgentByID :one
122
SELECT
133
*
@@ -200,3 +190,56 @@ WHERE
200190
WHERE
201191
wb.workspace_id = @workspace_id :: uuid
202192
);
193+
194+
-- name: GetWorkspaceAgentAndOwnerByAuthToken :one
195+
SELECT
196+
sqlc.embed(workspace_agents),
197+
workspaces.id AS workspace_id,
198+
users.id AS owner_id,
199+
users.username AS owner_name,
200+
users.status AS owner_status,
201+
array_cat(
202+
array_append(users.rbac_roles, 'member'),
203+
array_append(ARRAY[]::text[], 'organization-member:' || organization_members.organization_id::text)
204+
)::text[] as owner_roles,
205+
array_agg(COALESCE(group_members.group_id::text, ''))::text[] AS owner_groups
206+
FROM users
207+
INNER JOIN
208+
workspaces
209+
ON
210+
workspaces.owner_id = users.id
211+
INNER JOIN
212+
workspace_builds
213+
ON
214+
workspace_builds.workspace_id = workspaces.id
215+
INNER JOIN
216+
workspace_resources
217+
ON
218+
workspace_resources.job_id = workspace_builds.job_id
219+
INNER JOIN
220+
workspace_agents
221+
ON
222+
workspace_agents.resource_id = workspace_resources.id
223+
INNER JOIN -- every user is a member of some org
224+
organization_members
225+
ON
226+
organization_members.user_id = users.id
227+
LEFT JOIN -- as they may not be a member of any groups
228+
group_members
229+
ON
230+
group_members.user_id = users.id
231+
WHERE
232+
-- TODO: we can add more conditions here, such as:
233+
-- 1) The user must be active
234+
-- 2) The user must not be deleted
235+
-- 3) The workspace must be running
236+
workspace_agents.auth_token = @auth_token
237+
GROUP BY
238+
workspace_agents.id,
239+
workspaces.id,
240+
users.id,
241+
organization_members.organization_id,
242+
workspace_builds.build_number
243+
ORDER BY
244+
workspace_builds.build_number DESC
245+
LIMIT 1;

0 commit comments

Comments
 (0)