Skip to content

Commit da03c4b

Browse files
johnstcnjohnstcn
committed
add new query
1 parent 11c20d6 commit da03c4b

File tree

8 files changed

+250
-12
lines changed

8 files changed

+250
-12
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1474,6 +1474,14 @@ func (q *querier) GetUsersByIDs(ctx context.Context, ids []uuid.UUID) ([]databas
14741474
return q.db.GetUsersByIDs(ctx, ids)
14751475
}
14761476

1477+
func (q *querier) GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, agentID uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) {
1478+
// This is a system function
1479+
if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceSystem); err != nil {
1480+
return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, err
1481+
}
1482+
return q.db.GetWorkspaceAgentAndOwnerByAuthToken(ctx, agentID)
1483+
}
1484+
14771485
// GetWorkspaceAgentByAuthToken is used in http middleware to get the workspace agent.
14781486
// This should only be used by a system user in that middleware.
14791487
func (q *querier) GetWorkspaceAgentByAuthToken(ctx context.Context, authToken uuid.UUID) (database.WorkspaceAgent, error) {

‎coderd/database/dbfake/dbfake.go

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -651,6 +651,48 @@ func (q *FakeQuerier) isEveryoneGroup(id uuid.UUID) bool {
651651
return false
652652
}
653653

654+
func (q *FakeQuerier) GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) {
655+
q.mutex.RLock()
656+
defer q.mutex.RUnlock()
657+
var resp database.GetWorkspaceAgentAndOwnerByAuthTokenRow
658+
var found bool
659+
for _, agt := range q.workspaceAgents {
660+
if agt.AuthToken == authToken {
661+
resp.WorkspaceAgent = agt
662+
found = true
663+
break
664+
}
665+
}
666+
if !found {
667+
return resp, sql.ErrNoRows
668+
}
669+
670+
// get the related workspace and user
671+
for _, res := range q.workspaceResources {
672+
if resp.WorkspaceAgent.ResourceID != res.ID {
673+
continue
674+
}
675+
for _, build := range q.workspaceBuilds {
676+
if build.JobID != res.JobID {
677+
continue
678+
}
679+
for _, ws := range q.workspaces {
680+
if build.WorkspaceID != ws.ID {
681+
continue
682+
}
683+
resp.WorkspaceID = ws.ID
684+
if usr, err := q.getUserByIDNoLock(ws.OwnerID); err == nil {
685+
resp.OwnerID = usr.ID
686+
resp.OwnerRoles = usr.RBACRoles
687+
resp.OwnerName = usr.Username
688+
return resp, nil
689+
}
690+
}
691+
}
692+
}
693+
return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, sql.ErrNoRows
694+
}
695+
654696
func (*FakeQuerier) AcquireLock(_ context.Context, _ int64) error {
655697
return xerrors.New("AcquireLock must only be called within a transaction")
656698
}

‎coderd/database/dbmetrics/dbmetrics.go

Lines changed: 7 additions & 0 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: 15 additions & 0 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 & 0 deletions
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: 112 additions & 0 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: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,3 +200,57 @@ WHERE
200200
WHERE
201201
wb.workspace_id = @workspace_id :: uuid
202202
);
203+
204+
-- name: GetWorkspaceAgentAndOwnerByAuthToken :one
205+
SELECT
206+
sqlc.embed(workspace_agents),
207+
workspaces.id AS workspace_id,
208+
users.id AS owner_id,
209+
users.username AS owner_name,
210+
users.status AS owner_status,
211+
array_cat(
212+
-- All users are members
213+
array_append(users.rbac_roles, 'member'),
214+
(
215+
SELECT
216+
array_agg(org_roles)
217+
FROM
218+
organization_members,
219+
-- All org_members get the org-member role for their orgs
220+
unnest(
221+
array_append(roles, 'organization-member:' || organization_members.organization_id::text)
222+
) AS org_roles
223+
WHERE
224+
user_id = users.id
225+
)
226+
) :: text[] AS owner_roles,
227+
(
228+
SELECT
229+
array_agg(
230+
group_members.group_id :: text
231+
)
232+
FROM
233+
group_members
234+
WHERE
235+
user_id = users.id
236+
) :: text[] AS owner_groups
237+
FROM users
238+
INNER JOIN
239+
workspaces
240+
ON
241+
workspaces.owner_id = users.id
242+
INNER JOIN
243+
workspace_builds
244+
ON
245+
workspace_builds.workspace_id = workspaces.id
246+
INNER JOIN
247+
workspace_resources
248+
ON
249+
workspace_resources.job_id = workspace_builds.job_id
250+
INNER JOIN
251+
workspace_agents
252+
ON
253+
workspace_agents.resource_id = workspace_resources.id
254+
WHERE
255+
workspace_agents.auth_token = @auth_token
256+
LIMIT 1;

‎coderd/httpmw/workspaceagent.go

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,9 @@ func ExtractWorkspaceAgent(opts ExtractWorkspaceAgentConfig) func(http.Handler)
7474
})
7575
return
7676
}
77+
7778
//nolint:gocritic // System needs to be able to get workspace agents.
78-
agent, err := opts.DB.GetWorkspaceAgentByAuthToken(dbauthz.AsSystemRestricted(ctx), token)
79+
row, err := opts.DB.GetWorkspaceAgentAndOwnerByAuthToken(dbauthz.AsSystemRestricted(ctx), token)
7980
if err != nil {
8081
if errors.Is(err, sql.ErrNoRows) {
8182
optionalWrite(http.StatusUnauthorized, codersdk.Response{
@@ -86,23 +87,21 @@ func ExtractWorkspaceAgent(opts ExtractWorkspaceAgentConfig) func(http.Handler)
8687
}
8788

8889
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
89-
Message: "Internal error fetching workspace agent.",
90+
Message: "Internal error authorizing workspace agent.",
9091
Detail: err.Error(),
9192
})
9293
return
9394
}
9495

95-
//nolint:gocritic // System needs to be able to get workspace agents.
96-
subject, err := getAgentSubject(dbauthz.AsSystemRestricted(ctx), opts.DB, agent)
97-
if err != nil {
98-
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
99-
Message: "Internal error fetching workspace agent.",
100-
Detail: err.Error(),
101-
})
102-
return
103-
}
96+
subject := rbac.Subject{
97+
ID: row.OwnerID.String(),
98+
Roles: rbac.RoleNames(row.OwnerRoles),
99+
Groups: row.OwnerGroups,
100+
// Note: this is generated as a NullUUID even though it shouldn't be nullable based on the query.
101+
Scope: rbac.WorkspaceAgentScope(row.WorkspaceID, row.OwnerID),
102+
}.WithCachedASTValue()
104103

105-
ctx = context.WithValue(ctx, workspaceAgentContextKey{}, agent)
104+
ctx = context.WithValue(ctx, workspaceAgentContextKey{}, row.WorkspaceAgent)
106105
// Also set the dbauthz actor for the request.
107106
ctx = dbauthz.As(ctx, subject)
108107
next.ServeHTTP(rw, r.WithContext(ctx))

0 commit comments

Comments
 (0)