mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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>
This commit is contained in:
Thomas Kosiewski
@@ -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()
|
||||
|
||||
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,
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
apiKeyScope rbac.ScopeName
|
||||
}{
|
||||
{
|
||||
name: "empty (backwards compat)",
|
||||
apiKeyScope: "",
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
{
|
||||
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()
|
||||
|
||||
tickCh <- dbtime.Now()
|
||||
count := <-flushCh
|
||||
require.Equal(t, 1, count, "expected one flush with one id")
|
||||
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)
|
||||
}
|
||||
|
||||
newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
|
||||
require.NoError(t, err)
|
||||
return agent
|
||||
},
|
||||
).Do()
|
||||
|
||||
assert.True(t,
|
||||
newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
|
||||
"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
|
||||
)
|
||||
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")
|
||||
|
||||
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,
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user