Skip to content

Commit 93f17bc

Browse files
ThomasK33ThomasK33
authored
fix: remove unnecessary user lookup in agent API calls (#17934)
# Use workspace.OwnerUsername instead of fetching the owner This PR optimizes the agent API by using the `workspace.OwnerUsername` field directly instead of making an additional database query to fetch the owner's username. The change removes the need to call `GetUserByID` in the manifest API and workspace agent RPC endpoints. An issue arose when the agent token was scoped without access to user data (`api_key_scope = "no_user_data"`), causing the agent to fail to fetch the manifest due to an RBAC issue. Change-Id: I3b6e7581134e2374b364ee059e3b18ece3d98b41 Signed-off-by: Thomas Kosiewski <tk@coder.com>
1 parent 1267c9c commit 93f17bc

File tree

6 files changed

+194
-117
lines changed

6 files changed

+194
-117
lines changed

‎coderd/agentapi/manifest.go

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
4747
scripts []database.WorkspaceAgentScript
4848
metadata []database.WorkspaceAgentMetadatum
4949
workspace database.Workspace
50-
owner database.User
5150
devcontainers []database.WorkspaceAgentDevcontainer
5251
)
5352

@@ -76,10 +75,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
7675
if err != nil {
7776
return xerrors.Errorf("getting workspace by id: %w", err)
7877
}
79-
owner, err = a.Database.GetUserByID(ctx, workspace.OwnerID)
80-
if err != nil {
81-
return xerrors.Errorf("getting workspace owner by id: %w", err)
82-
}
8378
return err
8479
})
8580
eg.Go(func() (err error) {
@@ -98,7 +93,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
9893
AppSlugOrPort: "{{port}}",
9994
AgentName: workspaceAgent.Name,
10095
WorkspaceName: workspace.Name,
101-
Username: owner.Username,
96+
Username: workspace.OwnerUsername,
10297
}
10398

10499
vscodeProxyURI := vscodeProxyURI(appSlug, a.AccessURL, a.AppHostname)
@@ -115,7 +110,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
115110
}
116111
}
117112

118-
apps, err := dbAppsToProto(dbApps, workspaceAgent, owner.Username, workspace)
113+
apps, err := dbAppsToProto(dbApps, workspaceAgent, workspace.OwnerUsername, workspace)
119114
if err != nil {
120115
return nil, xerrors.Errorf("converting workspace apps: %w", err)
121116
}
@@ -128,7 +123,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
128123
return &agentproto.Manifest{
129124
AgentId: workspaceAgent.ID[:],
130125
AgentName: workspaceAgent.Name,
131-
OwnerUsername: owner.Username,
126+
OwnerUsername: workspace.OwnerUsername,
132127
WorkspaceId: workspace.ID[:],
133128
WorkspaceName: workspace.Name,
134129
GitAuthConfigs: gitAuthConfigs,

‎coderd/agentapi/manifest_test.go

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,10 @@ func TestGetManifest(t *testing.T) {
4646
Username: "cool-user",
4747
}
4848
workspace = database.Workspace{
49-
ID: uuid.New(),
50-
OwnerID: owner.ID,
51-
Name: "cool-workspace",
49+
ID: uuid.New(),
50+
OwnerID: owner.ID,
51+
OwnerUsername: owner.Username,
52+
Name: "cool-workspace",
5253
}
5354
agent = database.WorkspaceAgent{
5455
ID: uuid.New(),
@@ -336,7 +337,6 @@ func TestGetManifest(t *testing.T) {
336337
}).Return(metadata, nil)
337338
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil)
338339
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
339-
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
340340

341341
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
342342
require.NoError(t, err)
@@ -404,7 +404,6 @@ func TestGetManifest(t *testing.T) {
404404
}).Return([]database.WorkspaceAgentMetadatum{}, nil)
405405
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), childAgent.ID).Return([]database.WorkspaceAgentDevcontainer{}, nil)
406406
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
407-
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
408407

409408
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
410409
require.NoError(t, err)
@@ -468,7 +467,6 @@ func TestGetManifest(t *testing.T) {
468467
}).Return(metadata, nil)
469468
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil)
470469
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
471-
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
472470

473471
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
474472
require.NoError(t, err)

‎coderd/workspaceagents_test.go

Lines changed: 47 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -439,25 +439,55 @@ func TestWorkspaceAgentConnectRPC(t *testing.T) {
439439
t.Run("Connect", func(t *testing.T) {
440440
t.Parallel()
441441

442-
client, db := coderdtest.NewWithDatabase(t, nil)
443-
user := coderdtest.CreateFirstUser(t, client)
444-
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
445-
OrganizationID: user.OrganizationID,
446-
OwnerID: user.UserID,
447-
}).WithAgent().Do()
448-
_ = agenttest.New(t, client.URL, r.AgentToken)
449-
resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
442+
for _, tc := range []struct {
443+
name string
444+
apiKeyScope rbac.ScopeName
445+
}{
446+
{
447+
name: "empty (backwards compat)",
448+
apiKeyScope: "",
449+
},
450+
{
451+
name: "all",
452+
apiKeyScope: rbac.ScopeAll,
453+
},
454+
{
455+
name: "no_user_data",
456+
apiKeyScope: rbac.ScopeNoUserData,
457+
},
458+
{
459+
name: "application_connect",
460+
apiKeyScope: rbac.ScopeApplicationConnect,
461+
},
462+
} {
463+
t.Run(tc.name, func(t *testing.T) {
464+
client, db := coderdtest.NewWithDatabase(t, nil)
465+
user := coderdtest.CreateFirstUser(t, client)
466+
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
467+
OrganizationID: user.OrganizationID,
468+
OwnerID: user.UserID,
469+
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
470+
for _, agent := range agents {
471+
agent.ApiKeyScope = string(tc.apiKeyScope)
472+
}
450473

451-
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
452-
defer cancel()
474+
return agents
475+
}).Do()
476+
_ = agenttest.New(t, client.URL, r.AgentToken)
477+
resources := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).AgentNames([]string{}).Wait()
453478

454-
conn, err := workspacesdk.New(client).
455-
DialAgent(ctx, resources[0].Agents[0].ID, nil)
456-
require.NoError(t, err)
457-
defer func() {
458-
_ = conn.Close()
459-
}()
460-
conn.AwaitReachable(ctx)
479+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
480+
defer cancel()
481+
482+
conn, err := workspacesdk.New(client).
483+
DialAgent(ctx, resources[0].Agents[0].ID, nil)
484+
require.NoError(t, err)
485+
defer func() {
486+
_ = conn.Close()
487+
}()
488+
conn.AwaitReachable(ctx)
489+
})
490+
}
461491
})
462492

463493
t.Run("FailNonLatestBuild", func(t *testing.T) {

‎coderd/workspaceagentsrpc.go

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -76,17 +76,8 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
7676
return
7777
}
7878

79-
owner, err := api.Database.GetUserByID(ctx, workspace.OwnerID)
80-
if err != nil {
81-
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
82-
Message: "Internal error fetching user.",
83-
Detail: err.Error(),
84-
})
85-
return
86-
}
87-
8879
logger = logger.With(
89-
slog.F("owner", owner.Username),
80+
slog.F("owner", workspace.OwnerUsername),
9081
slog.F("workspace_name", workspace.Name),
9182
slog.F("agent_name", workspaceAgent.Name),
9283
)
@@ -170,7 +161,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
170161
})
171162

172163
streamID := tailnet.StreamID{
173-
Name: fmt.Sprintf("%s-%s-%s", owner.Username, workspace.Name, workspaceAgent.Name),
164+
Name: fmt.Sprintf("%s-%s-%s", workspace.OwnerUsername, workspace.Name, workspaceAgent.Name),
174165
ID: workspaceAgent.ID,
175166
Auth: tailnet.AgentCoordinateeAuth{ID: workspaceAgent.ID},
176167
}

0 commit comments

Comments
 (0)