Skip to content

Commit 0723dd3

Browse files
f0sself0ssel
authored
fix: ensure agent token is from latest build in middleware (coder#12443)
1 parent 63696d7 commit 0723dd3

15 files changed

+243
-261
lines changed

‎coderd/coderd.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -902,7 +902,7 @@ func New(options *Options) *API {
902902
httpmw.RequireAPIKeyOrWorkspaceProxyAuth(),
903903
).Get("/connection", api.workspaceAgentConnectionGeneric)
904904
r.Route("/me", func(r chi.Router) {
905-
r.Use(httpmw.ExtractWorkspaceAgent(httpmw.ExtractWorkspaceAgentConfig{
905+
r.Use(httpmw.ExtractWorkspaceAgentAndLatestBuild(httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
906906
DB: options.Database,
907907
Optional: false,
908908
}))

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1880,12 +1880,12 @@ func (q *querier) GetUsersByIDs(ctx context.Context, ids []uuid.UUID) ([]databas
18801880
return q.db.GetUsersByIDs(ctx, ids)
18811881
}
18821882

1883-
func (q *querier) GetWorkspaceAgentAndOwnerByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) {
1883+
func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) {
18841884
// This is a system function
18851885
if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceSystem); err != nil {
1886-
return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, err
1886+
return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, err
18871887
}
1888-
return q.db.GetWorkspaceAgentAndOwnerByAuthToken(ctx, authToken)
1888+
return q.db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken)
18891889
}
18901890

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

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2274,7 +2274,7 @@ func (s *MethodTestSuite) TestSystemFunctions() {
22742274
s.Run("GetReplicaByID", s.Subtest(func(db database.Store, check *expects) {
22752275
check.Args(uuid.New()).Asserts(rbac.ResourceSystem, rbac.ActionRead).Errors(sql.ErrNoRows)
22762276
}))
2277-
s.Run("GetWorkspaceAgentAndOwnerByAuthToken", s.Subtest(func(db database.Store, check *expects) {
2277+
s.Run("GetWorkspaceAgentAndLatestBuildByAuthToken", s.Subtest(func(db database.Store, check *expects) {
22782278
check.Args(uuid.New()).Asserts(rbac.ResourceSystem, rbac.ActionRead).Errors(sql.ErrNoRows)
22792279
}))
22802280
s.Run("GetUserLinksByUserID", s.Subtest(func(db database.Store, check *expects) {

‎coderd/database/dbmem/dbmem.go

Lines changed: 34 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ func New() database.Store {
6969
templates: make([]database.TemplateTable, 0),
7070
workspaceAgentStats: make([]database.WorkspaceAgentStat, 0),
7171
workspaceAgentLogs: make([]database.WorkspaceAgentLog, 0),
72-
workspaceBuilds: make([]database.WorkspaceBuildTable, 0),
72+
workspaceBuilds: make([]database.WorkspaceBuild, 0),
7373
workspaceApps: make([]database.WorkspaceApp, 0),
7474
workspaces: make([]database.Workspace, 0),
7575
licenses: make([]database.License, 0),
@@ -171,7 +171,7 @@ type data struct {
171171
workspaceApps []database.WorkspaceApp
172172
workspaceAppStatsLastInsertID int64
173173
workspaceAppStats []database.WorkspaceAppStat
174-
workspaceBuilds []database.WorkspaceBuildTable
174+
workspaceBuilds []database.WorkspaceBuild
175175
workspaceBuildParameters []database.WorkspaceBuildParameter
176176
workspaceResourceMetadata []database.WorkspaceResourceMetadatum
177177
workspaceResources []database.WorkspaceResource
@@ -542,7 +542,7 @@ func (q *FakeQuerier) templateVersionWithUserNoLock(tpl database.TemplateVersion
542542
return withUser
543543
}
544544

545-
func (q *FakeQuerier) workspaceBuildWithUserNoLock(tpl database.WorkspaceBuildTable) database.WorkspaceBuild {
545+
func (q *FakeQuerier) workspaceBuildWithUserNoLock(tpl database.WorkspaceBuild) database.WorkspaceBuild {
546546
var user database.User
547547
for _, _user := range q.users {
548548
if _user.ID == tpl.InitiatorID {
@@ -2801,7 +2801,7 @@ func (q *FakeQuerier) GetQuotaConsumedForUser(_ context.Context, userID uuid.UUI
28012801
continue
28022802
}
28032803

2804-
var lastBuild database.WorkspaceBuildTable
2804+
var lastBuild database.WorkspaceBuild
28052805
for _, build := range q.workspaceBuilds {
28062806
if build.WorkspaceID != workspace.ID {
28072807
continue
@@ -3488,7 +3488,7 @@ func (q *FakeQuerier) GetTemplateParameterInsights(ctx context.Context, arg data
34883488
defer q.mutex.RUnlock()
34893489

34903490
// WITH latest_workspace_builds ...
3491-
latestWorkspaceBuilds := make(map[uuid.UUID]database.WorkspaceBuildTable)
3491+
latestWorkspaceBuilds := make(map[uuid.UUID]database.WorkspaceBuild)
34923492
for _, wb := range q.workspaceBuilds {
34933493
if wb.CreatedAt.Before(arg.StartTime) || wb.CreatedAt.Equal(arg.EndTime) || wb.CreatedAt.After(arg.EndTime) {
34943494
continue
@@ -4270,20 +4270,14 @@ func (q *FakeQuerier) GetUsersByIDs(_ context.Context, ids []uuid.UUID) ([]datab
42704270
return users, nil
42714271
}
42724272

4273-
func (q *FakeQuerier) GetWorkspaceAgentAndOwnerByAuthToken(_ context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndOwnerByAuthTokenRow, error) {
4273+
func (q *FakeQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(_ context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) {
42744274
q.mutex.RLock()
42754275
defer q.mutex.RUnlock()
4276-
4277-
// map of build number -> row
4278-
rows := make(map[int32]database.GetWorkspaceAgentAndOwnerByAuthTokenRow)
4279-
4280-
// We want to return the latest build number
4281-
var latestBuildNumber int32
4276+
rows := []database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}
4277+
// We want to return the latest build number for each workspace
4278+
latestBuildNumber := make(map[uuid.UUID]int32)
42824279

42834280
for _, agt := range q.workspaceAgents {
4284-
if agt.AuthToken != authToken {
4285-
continue
4286-
}
42874281
// get the related workspace and user
42884282
for _, res := range q.workspaceResources {
42894283
if agt.ResourceID != res.ID {
@@ -4300,47 +4294,43 @@ func (q *FakeQuerier) GetWorkspaceAgentAndOwnerByAuthToken(_ context.Context, au
43004294
if ws.Deleted {
43014295
continue
43024296
}
4303-
var row database.GetWorkspaceAgentAndOwnerByAuthTokenRow
4304-
row.WorkspaceID = ws.ID
4305-
row.TemplateID = ws.TemplateID
4297+
row := database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{
4298+
Workspace: database.Workspace{
4299+
ID: ws.ID,
4300+
TemplateID: ws.TemplateID,
4301+
},
4302+
WorkspaceAgent: agt,
4303+
WorkspaceBuild: build,
4304+
}
43064305
usr, err := q.getUserByIDNoLock(ws.OwnerID)
43074306
if err != nil {
4308-
return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, sql.ErrNoRows
4309-
}
4310-
row.OwnerID = usr.ID
4311-
row.OwnerRoles = append(usr.RBACRoles, "member")
4312-
// We also need to get org roles for the user
4313-
row.OwnerName = usr.Username
4314-
row.WorkspaceAgent = agt
4315-
row.TemplateVersionID = build.TemplateVersionID
4316-
for _, mem := range q.organizationMembers {
4317-
if mem.UserID == usr.ID {
4318-
row.OwnerRoles = append(row.OwnerRoles, fmt.Sprintf("organization-member:%s", mem.OrganizationID.String()))
4319-
}
4320-
}
4321-
// And group memberships
4322-
for _, groupMem := range q.groupMembers {
4323-
if groupMem.UserID == usr.ID {
4324-
row.OwnerGroups = append(row.OwnerGroups, groupMem.GroupID.String())
4325-
}
4307+
return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, sql.ErrNoRows
43264308
}
4309+
row.Workspace.OwnerID = usr.ID
43274310

43284311
// Keep track of the latest build number
4329-
rows[build.BuildNumber] = row
4330-
if build.BuildNumber > latestBuildNumber {
4331-
latestBuildNumber = build.BuildNumber
4312+
rows = append(rows, row)
4313+
if build.BuildNumber > latestBuildNumber[ws.ID] {
4314+
latestBuildNumber[ws.ID] = build.BuildNumber
43324315
}
43334316
}
43344317
}
43354318
}
43364319
}
43374320

4338-
if len(rows) == 0 {
4339-
return database.GetWorkspaceAgentAndOwnerByAuthTokenRow{}, sql.ErrNoRows
4321+
for i := range rows {
4322+
if rows[i].WorkspaceAgent.AuthToken != authToken {
4323+
continue
4324+
}
4325+
4326+
if rows[i].WorkspaceBuild.BuildNumber != latestBuildNumber[rows[i].Workspace.ID] {
4327+
continue
4328+
}
4329+
4330+
return rows[i], nil
43404331
}
43414332

4342-
// Return the row related to the latest build
4343-
return rows[latestBuildNumber], nil
4333+
return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, sql.ErrNoRows
43444334
}
43454335

43464336
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
62436233
q.mutex.Lock()
62446234
defer q.mutex.Unlock()
62456235

6246-
workspaceBuild := database.WorkspaceBuildTable{
6236+
workspaceBuild := database.WorkspaceBuild{
62476237
ID: arg.ID,
62486238
CreatedAt: arg.CreatedAt,
62496239
UpdatedAt: arg.UpdatedAt,

‎coderd/database/dbmetrics/dbmetrics.go

Lines changed: 3 additions & 3 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: 66 additions & 72 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)