From 840c0db240e548bab289b3336f3e22e6a8f19731 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Tue, 20 May 2025 13:04:41 +0200 Subject: [PATCH] fix: remove unnecessary user lookup in agent API calls Change-Id: I3b6e7581134e2374b364ee059e3b18ece3d98b41 Signed-off-by: Thomas Kosiewski --- coderd/agentapi/manifest.go | 11 +- coderd/agentapi/manifest_test.go | 10 +- coderd/workspaceagents_test.go | 64 ++++++--- coderd/workspaceagentsrpc.go | 13 +- coderd/workspaceagentsrpc_test.go | 212 +++++++++++++++++++----------- flake.nix | 1 + 6 files changed, 194 insertions(+), 117 deletions(-) diff --git a/coderd/agentapi/manifest.go b/coderd/agentapi/manifest.go index 66bfe4cb5f94f..855ff4b8acd37 100644 --- a/coderd/agentapi/manifest.go +++ b/coderd/agentapi/manifest.go @@ -47,7 +47,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest scripts []database.WorkspaceAgentScript metadata []database.WorkspaceAgentMetadatum workspace database.Workspace - owner database.User devcontainers []database.WorkspaceAgentDevcontainer ) @@ -76,10 +75,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest if err != nil { return xerrors.Errorf("getting workspace by id: %w", err) } - owner, err = a.Database.GetUserByID(ctx, workspace.OwnerID) - if err != nil { - return xerrors.Errorf("getting workspace owner by id: %w", err) - } return err }) eg.Go(func() (err error) { @@ -98,7 +93,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest AppSlugOrPort: "{{port}}", AgentName: workspaceAgent.Name, WorkspaceName: workspace.Name, - Username: owner.Username, + Username: workspace.OwnerUsername, } vscodeProxyURI := vscodeProxyURI(appSlug, a.AccessURL, a.AppHostname) @@ -115,7 +110,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest } } - apps, err := dbAppsToProto(dbApps, workspaceAgent, owner.Username, workspace) + apps, err := dbAppsToProto(dbApps, workspaceAgent, workspace.OwnerUsername, workspace) if err != nil { return nil, xerrors.Errorf("converting workspace apps: %w", err) } @@ -128,7 +123,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest return &agentproto.Manifest{ AgentId: workspaceAgent.ID[:], AgentName: workspaceAgent.Name, - OwnerUsername: owner.Username, + OwnerUsername: workspace.OwnerUsername, WorkspaceId: workspace.ID[:], WorkspaceName: workspace.Name, GitAuthConfigs: gitAuthConfigs, diff --git a/coderd/agentapi/manifest_test.go b/coderd/agentapi/manifest_test.go index 9273acb0c40ff..fc46f5fe480f8 100644 --- a/coderd/agentapi/manifest_test.go +++ b/coderd/agentapi/manifest_test.go @@ -46,9 +46,10 @@ func TestGetManifest(t *testing.T) { Username: "cool-user", } workspace = database.Workspace{ - ID: uuid.New(), - OwnerID: owner.ID, - Name: "cool-workspace", + ID: uuid.New(), + OwnerID: owner.ID, + OwnerUsername: owner.Username, + Name: "cool-workspace", } agent = database.WorkspaceAgent{ ID: uuid.New(), @@ -336,7 +337,6 @@ func TestGetManifest(t *testing.T) { }).Return(metadata, nil) mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil) mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil) - mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil) got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{}) require.NoError(t, err) @@ -404,7 +404,6 @@ func TestGetManifest(t *testing.T) { }).Return([]database.WorkspaceAgentMetadatum{}, nil) mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), childAgent.ID).Return([]database.WorkspaceAgentDevcontainer{}, nil) mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil) - mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil) got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{}) require.NoError(t, err) @@ -468,7 +467,6 @@ func TestGetManifest(t *testing.T) { }).Return(metadata, nil) mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil) mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil) - mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil) got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{}) require.NoError(t, err) diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 27da80b3c579b..f4f3dcdec9f89 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -439,25 +439,55 @@ func TestWorkspaceAgentConnectRPC(t *testing.T) { t.Run("Connect", func(t *testing.T) { t.Parallel() - client, db := coderdtest.NewWithDatabase(t, nil) - user := coderdtest.CreateFirstUser(t, client) - r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ - OrganizationID: user.OrganizationID, - OwnerID: user.UserID, - }).WithAgent().Do() - _ = agenttest.New(t, client.URL, r.AgentToken) - resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID) + for _, tc := range []struct { + name string + apiKeyScope rbac.ScopeName + }{ + { + name: "empty (backwards compat)", + apiKeyScope: "", + }, + { + name: "all", + apiKeyScope: rbac.ScopeAll, + }, + { + name: "no_user_data", + apiKeyScope: rbac.ScopeNoUserData, + }, + { + name: "application_connect", + apiKeyScope: rbac.ScopeApplicationConnect, + }, + } { + t.Run(tc.name, func(t *testing.T) { + client, db := coderdtest.NewWithDatabase(t, nil) + user := coderdtest.CreateFirstUser(t, client) + r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + }).WithAgent(func(agents []*proto.Agent) []*proto.Agent { + for _, agent := range agents { + agent.ApiKeyScope = string(tc.apiKeyScope) + } - ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) - defer cancel() + return agents + }).Do() + _ = agenttest.New(t, client.URL, r.AgentToken) + resources := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).AgentNames([]string{}).Wait() - conn, err := workspacesdk.New(client). - DialAgent(ctx, resources[0].Agents[0].ID, nil) - require.NoError(t, err) - defer func() { - _ = conn.Close() - }() - conn.AwaitReachable(ctx) + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancel() + + conn, err := workspacesdk.New(client). + DialAgent(ctx, resources[0].Agents[0].ID, nil) + require.NoError(t, err) + defer func() { + _ = conn.Close() + }() + conn.AwaitReachable(ctx) + }) + } }) t.Run("FailNonLatestBuild", func(t *testing.T) { diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go index 43da35410f632..2dcf65bd8c7d5 100644 --- a/coderd/workspaceagentsrpc.go +++ b/coderd/workspaceagentsrpc.go @@ -76,17 +76,8 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) { return } - owner, err := api.Database.GetUserByID(ctx, workspace.OwnerID) - if err != nil { - httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Internal error fetching user.", - Detail: err.Error(), - }) - return - } - logger = logger.With( - slog.F("owner", owner.Username), + slog.F("owner", workspace.OwnerUsername), slog.F("workspace_name", workspace.Name), slog.F("agent_name", workspaceAgent.Name), ) @@ -170,7 +161,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) { }) streamID := tailnet.StreamID{ - Name: fmt.Sprintf("%s-%s-%s", owner.Username, workspace.Name, workspaceAgent.Name), + Name: fmt.Sprintf("%s-%s-%s", workspace.OwnerUsername, workspace.Name, workspaceAgent.Name), ID: workspaceAgent.ID, Auth: tailnet.AgentCoordinateeAuth{ID: workspaceAgent.ID}, } diff --git a/coderd/workspaceagentsrpc_test.go b/coderd/workspaceagentsrpc_test.go index caea9b39c2f54..5175f80b0b723 100644 --- a/coderd/workspaceagentsrpc_test.go +++ b/coderd/workspaceagentsrpc_test.go @@ -13,6 +13,7 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/provisionersdk/proto" "github.com/coder/coder/v2/testutil" @@ -22,89 +23,150 @@ import ( func TestWorkspaceAgentReportStats(t *testing.T) { t.Parallel() - tickCh := make(chan time.Time) - flushCh := make(chan int, 1) - client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{ - WorkspaceUsageTrackerFlush: flushCh, - WorkspaceUsageTrackerTick: tickCh, - }) - user := coderdtest.CreateFirstUser(t, client) - r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ - OrganizationID: user.OrganizationID, - OwnerID: user.UserID, - LastUsedAt: dbtime.Now().Add(-time.Minute), - }).WithAgent().Do() + for _, tc := range []struct { + name string + apiKeyScope rbac.ScopeName + }{ + { + name: "empty (backwards compat)", + apiKeyScope: "", + }, + { + name: "all", + apiKeyScope: rbac.ScopeAll, + }, + { + name: "no_user_data", + apiKeyScope: rbac.ScopeNoUserData, + }, + { + name: "application_connect", + apiKeyScope: rbac.ScopeApplicationConnect, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() - ac := agentsdk.New(client.URL) - ac.SetSessionToken(r.AgentToken) - conn, err := ac.ConnectRPC(context.Background()) - require.NoError(t, err) - defer func() { - _ = conn.Close() - }() - agentAPI := agentproto.NewDRPCAgentClient(conn) + tickCh := make(chan time.Time) + flushCh := make(chan int, 1) + client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + WorkspaceUsageTrackerFlush: flushCh, + WorkspaceUsageTrackerTick: tickCh, + }) + user := coderdtest.CreateFirstUser(t, client) + r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + LastUsedAt: dbtime.Now().Add(-time.Minute), + }).WithAgent( + func(agent []*proto.Agent) []*proto.Agent { + for _, a := range agent { + a.ApiKeyScope = string(tc.apiKeyScope) + } - _, err = agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{ - Stats: &agentproto.Stats{ - ConnectionsByProto: map[string]int64{"TCP": 1}, - ConnectionCount: 1, - RxPackets: 1, - RxBytes: 1, - TxPackets: 1, - TxBytes: 1, - SessionCountVscode: 1, - SessionCountJetbrains: 0, - SessionCountReconnectingPty: 0, - SessionCountSsh: 0, - ConnectionMedianLatencyMs: 10, - }, - }) - require.NoError(t, err) + return agent + }, + ).Do() + + ac := agentsdk.New(client.URL) + ac.SetSessionToken(r.AgentToken) + conn, err := ac.ConnectRPC(context.Background()) + require.NoError(t, err) + defer func() { + _ = conn.Close() + }() + agentAPI := agentproto.NewDRPCAgentClient(conn) + + _, err = agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{ + Stats: &agentproto.Stats{ + ConnectionsByProto: map[string]int64{"TCP": 1}, + ConnectionCount: 1, + RxPackets: 1, + RxBytes: 1, + TxPackets: 1, + TxBytes: 1, + SessionCountVscode: 1, + SessionCountJetbrains: 0, + SessionCountReconnectingPty: 0, + SessionCountSsh: 0, + ConnectionMedianLatencyMs: 10, + }, + }) + require.NoError(t, err) - tickCh <- dbtime.Now() - count := <-flushCh - require.Equal(t, 1, count, "expected one flush with one id") + tickCh <- dbtime.Now() + count := <-flushCh + require.Equal(t, 1, count, "expected one flush with one id") - newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID) - require.NoError(t, err) + newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID) + require.NoError(t, err) - assert.True(t, - newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt), - "%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt, - ) + assert.True(t, + newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt), + "%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt, + ) + }) + } } func TestAgentAPI_LargeManifest(t *testing.T) { t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) - client, store := coderdtest.NewWithDatabase(t, nil) - adminUser := coderdtest.CreateFirstUser(t, client) - n := 512000 - longScript := make([]byte, n) - for i := range longScript { - longScript[i] = 'q' + + for _, tc := range []struct { + name string + apiKeyScope rbac.ScopeName + }{ + { + name: "empty (backwards compat)", + apiKeyScope: "", + }, + { + name: "all", + apiKeyScope: rbac.ScopeAll, + }, + { + name: "no_user_data", + apiKeyScope: rbac.ScopeNoUserData, + }, + { + name: "application_connect", + apiKeyScope: rbac.ScopeApplicationConnect, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + client, store := coderdtest.NewWithDatabase(t, nil) + adminUser := coderdtest.CreateFirstUser(t, client) + n := 512000 + longScript := make([]byte, n) + for i := range longScript { + longScript[i] = 'q' + } + r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{ + OrganizationID: adminUser.OrganizationID, + OwnerID: adminUser.UserID, + }).WithAgent(func(agents []*proto.Agent) []*proto.Agent { + agents[0].Scripts = []*proto.Script{ + { + Script: string(longScript), + }, + } + agents[0].ApiKeyScope = string(tc.apiKeyScope) + return agents + }).Do() + ac := agentsdk.New(client.URL) + ac.SetSessionToken(r.AgentToken) + conn, err := ac.ConnectRPC(ctx) + defer func() { + _ = conn.Close() + }() + require.NoError(t, err) + agentAPI := agentproto.NewDRPCAgentClient(conn) + manifest, err := agentAPI.GetManifest(ctx, &agentproto.GetManifestRequest{}) + require.NoError(t, err) + require.Len(t, manifest.Scripts, 1) + require.Len(t, manifest.Scripts[0].Script, n) + }) } - r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{ - OrganizationID: adminUser.OrganizationID, - OwnerID: adminUser.UserID, - }).WithAgent(func(agents []*proto.Agent) []*proto.Agent { - agents[0].Scripts = []*proto.Script{ - { - Script: string(longScript), - }, - } - return agents - }).Do() - ac := agentsdk.New(client.URL) - ac.SetSessionToken(r.AgentToken) - conn, err := ac.ConnectRPC(ctx) - defer func() { - _ = conn.Close() - }() - require.NoError(t, err) - agentAPI := agentproto.NewDRPCAgentClient(conn) - manifest, err := agentAPI.GetManifest(ctx, &agentproto.GetManifestRequest{}) - require.NoError(t, err) - require.Len(t, manifest.Scripts, 1) - require.Len(t, manifest.Scripts[0].Script, n) } diff --git a/flake.nix b/flake.nix index bff207662f913..c0f36c3be6e0f 100644 --- a/flake.nix +++ b/flake.nix @@ -141,6 +141,7 @@ kubectl kubectx kubernetes-helm + lazydocker lazygit less mockgen pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy