mirror of https://github.com/coder/coder.git
fix: ensure agent token is from latest build in middleware (#12443)
This commit is contained in:
parent
63696d762f
commit
0723dd3abf
|
@ -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,
|
||||
}))
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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
|
||||
}
|
||||
|
||||
// Return the row related to the latest build
|
||||
return rows[latestBuildNumber], nil
|
||||
if rows[i].WorkspaceBuild.BuildNumber != latestBuildNumber[rows[i].Workspace.ID] {
|
||||
continue
|
||||
}
|
||||
|
||||
return rows[i], 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,
|
||||
|
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
@ -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.
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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 {
|
||||
type GetWorkspaceAgentAndLatestBuildByAuthTokenRow struct {
|
||||
Workspace Workspace `db:"workspace" json:"workspace"`
|
||||
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"`
|
||||
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
|
||||
}
|
||||
|
|
|
@ -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;
|
||||
;
|
||||
|
|
|
@ -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))
|
||||
|
|
|
@ -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
|
||||
}
|
||||
|
|
|
@ -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 {
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
}
|
||||
|
|
|
@ -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,
|
||||
}),
|
||||
|
|
Loading…
Reference in New Issue