diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go index 252e6b5c08..59ce378844 100644 --- a/coderd/agentapi/api.go +++ b/coderd/agentapi/api.go @@ -197,6 +197,7 @@ func New(opts Options, workspace database.Workspace) *API { AgentFn: api.agent, ConnectionLogger: opts.ConnectionLogger, Database: opts.Database, + Workspace: api.cachedWorkspaceFields, Log: opts.Log, } diff --git a/coderd/agentapi/connectionlog.go b/coderd/agentapi/connectionlog.go index bd11f9e726..e38a312c4e 100644 --- a/coderd/agentapi/connectionlog.go +++ b/coderd/agentapi/connectionlog.go @@ -14,11 +14,13 @@ import ( "github.com/coder/coder/v2/coderd/connectionlog" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" + "github.com/coder/coder/v2/coderd/database/dbauthz" ) type ConnLogAPI struct { AgentFn func(context.Context) (database.WorkspaceAgent, error) ConnectionLogger *atomic.Pointer[connectionlog.ConnectionLogger] + Workspace *CachedWorkspaceFields Database database.Store Log slog.Logger } @@ -51,14 +53,31 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor } } + // Inject RBAC object into context for dbauthz fast path, avoid having to + // call GetWorkspaceByAgentID on every metadata update. + rbacCtx := ctx + var ws database.WorkspaceIdentity + if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok { + ws = dbws + rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject()) + if err != nil { + // Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID. + //nolint:gocritic + a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err)) + } + } + // Fetch contextual data for this connection log event. - workspaceAgent, err := a.AgentFn(ctx) + workspaceAgent, err := a.AgentFn(rbacCtx) if err != nil { return nil, xerrors.Errorf("get agent: %w", err) } - workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID) - if err != nil { - return nil, xerrors.Errorf("get workspace by agent id: %w", err) + if ws.Equal(database.WorkspaceIdentity{}) { + workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID) + if err != nil { + return nil, xerrors.Errorf("get workspace by agent id: %w", err) + } + ws = database.WorkspaceIdentityFromWorkspace(workspace) } // Some older clients may incorrectly report "localhost" as the IP address. @@ -74,10 +93,10 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor err = connLogger.Upsert(ctx, database.UpsertConnectionLogParams{ ID: uuid.New(), Time: req.GetConnection().GetTimestamp().AsTime(), - OrganizationID: workspace.OrganizationID, - WorkspaceOwnerID: workspace.OwnerID, - WorkspaceID: workspace.ID, - WorkspaceName: workspace.Name, + OrganizationID: ws.OrganizationID, + WorkspaceOwnerID: ws.OwnerID, + WorkspaceID: ws.ID, + WorkspaceName: ws.Name, AgentName: workspaceAgent.Name, Type: connectionType, Code: code, diff --git a/coderd/agentapi/connectionlog_test.go b/coderd/agentapi/connectionlog_test.go index 81d969e5ba..306220dce2 100644 --- a/coderd/agentapi/connectionlog_test.go +++ b/coderd/agentapi/connectionlog_test.go @@ -117,6 +117,7 @@ func TestConnectionLog(t *testing.T) { AgentFn: func(context.Context) (database.WorkspaceAgent, error) { return agent, nil }, + Workspace: &agentapi.CachedWorkspaceFields{}, } api.ReportConnection(context.Background(), &agentproto.ReportConnectionRequest{ Connection: &agentproto.Connection{ diff --git a/coderd/agentapi/metadata.go b/coderd/agentapi/metadata.go index 756422f856..1d4e23ab88 100644 --- a/coderd/agentapi/metadata.go +++ b/coderd/agentapi/metadata.go @@ -47,7 +47,20 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B maxErrorLen = maxValueLen ) - workspaceAgent, err := a.AgentFn(ctx) + // Inject RBAC object into context for dbauthz fast path, avoid having to + // call GetWorkspaceByAgentID on every metadata update. + var err error + rbacCtx := ctx + if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok { + rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject()) + if err != nil { + // Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID. + //nolint:gocritic + a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err)) + } + } + + workspaceAgent, err := a.AgentFn(rbacCtx) if err != nil { return nil, err } @@ -109,18 +122,6 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B ) } - // Inject RBAC object into context for dbauthz fast path, avoid having to - // call GetWorkspaceByAgentID on every metadata update. - rbacCtx := ctx - if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok { - rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject()) - if err != nil { - // Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID. - //nolint:gocritic - a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err)) - } - } - err = a.Database.UpdateWorkspaceAgentMetadata(rbacCtx, dbUpdate) if err != nil { return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err) diff --git a/coderd/agentapi/metadata_test.go b/coderd/agentapi/metadata_test.go index 1ba02d037f..866b2a8bf2 100644 --- a/coderd/agentapi/metadata_test.go +++ b/coderd/agentapi/metadata_test.go @@ -295,6 +295,7 @@ func TestBatchUpdateMetadata(t *testing.T) { now = dbtime.Now() // Set up consistent IDs that represent a valid workspace->agent relationship workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012") + templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000") ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321") orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111") agentID = uuid.MustParse("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") @@ -358,8 +359,48 @@ func TestBatchUpdateMetadata(t *testing.T) { OrganizationID: orgID, }) - // Create context with system actor so authorization passes - ctx := dbauthz.AsSystemRestricted(context.Background()) + // Create roles with workspace permissions + userRoles := rbac.Roles([]rbac.Role{ + { + Identifier: rbac.RoleMember(), + User: []rbac.Permission{ + { + Negate: false, + ResourceType: rbac.ResourceWorkspace.Type, + Action: policy.WildcardSymbol, + }, + }, + ByOrgID: map[string]rbac.OrgPermissions{ + orgID.String(): { + Member: []rbac.Permission{ + { + Negate: false, + ResourceType: rbac.ResourceWorkspace.Type, + Action: policy.WildcardSymbol, + }, + }, + }, + }, + }, + }) + + agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{ + WorkspaceID: workspaceID, + OwnerID: ownerID, + TemplateID: templateID, + VersionID: uuid.New(), + }) + + ctx := dbauthz.As(context.Background(), rbac.Subject{ + Type: rbac.SubjectTypeUser, + FriendlyName: "testuser", + Email: "testuser@example.com", + ID: ownerID.String(), + Roles: userRoles, + Groups: []string{orgID.String()}, + Scope: agentScope, + }.WithCachedASTValue()) + resp, err := api.BatchUpdateMetadata(ctx, req) require.NoError(t, err) require.NotNil(t, resp) @@ -376,6 +417,7 @@ func TestBatchUpdateMetadata(t *testing.T) { pub = &fakePublisher{} now = dbtime.Now() workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012") + templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000") ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321") orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111") agentID = uuid.MustParse("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb") @@ -445,12 +487,53 @@ func TestBatchUpdateMetadata(t *testing.T) { OrganizationID: uuid.Nil, // Invalid: fails dbauthz fast path validation }) - // Create context with system actor so authorization passes - ctx := dbauthz.AsSystemRestricted(context.Background()) + // Create roles with workspace permissions + userRoles := rbac.Roles([]rbac.Role{ + { + Identifier: rbac.RoleMember(), + User: []rbac.Permission{ + { + Negate: false, + ResourceType: rbac.ResourceWorkspace.Type, + Action: policy.WildcardSymbol, + }, + }, + ByOrgID: map[string]rbac.OrgPermissions{ + orgID.String(): { + Member: []rbac.Permission{ + { + Negate: false, + ResourceType: rbac.ResourceWorkspace.Type, + Action: policy.WildcardSymbol, + }, + }, + }, + }, + }, + }) + + agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{ + WorkspaceID: workspaceID, + OwnerID: ownerID, + TemplateID: templateID, + VersionID: uuid.New(), + }) + + ctx := dbauthz.As(context.Background(), rbac.Subject{ + Type: rbac.SubjectTypeUser, + FriendlyName: "testuser", + Email: "testuser@example.com", + ID: ownerID.String(), + Roles: userRoles, + Groups: []string{orgID.String()}, + Scope: agentScope, + }.WithCachedASTValue()) + resp, err := api.BatchUpdateMetadata(ctx, req) require.NoError(t, err) require.NotNil(t, resp) }) + // Test RBAC slow path - no RBAC object in context // This test verifies that when no RBAC object is present in context, the dbauthz layer // falls back to the slow path and calls GetWorkspaceByAgentID. @@ -463,6 +546,7 @@ func TestBatchUpdateMetadata(t *testing.T) { pub = &fakePublisher{} now = dbtime.Now() workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012") + templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000") ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321") orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111") agentID = uuid.MustParse("dddddddd-dddd-dddd-dddd-dddddddddddd") @@ -523,8 +607,48 @@ func TestBatchUpdateMetadata(t *testing.T) { }, } - // Create context with system actor so authorization passes - ctx := dbauthz.AsSystemRestricted(context.Background()) + // Create roles with workspace permissions + userRoles := rbac.Roles([]rbac.Role{ + { + Identifier: rbac.RoleMember(), + User: []rbac.Permission{ + { + Negate: false, + ResourceType: rbac.ResourceWorkspace.Type, + Action: policy.WildcardSymbol, + }, + }, + ByOrgID: map[string]rbac.OrgPermissions{ + orgID.String(): { + Member: []rbac.Permission{ + { + Negate: false, + ResourceType: rbac.ResourceWorkspace.Type, + Action: policy.WildcardSymbol, + }, + }, + }, + }, + }, + }) + + agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{ + WorkspaceID: workspaceID, + OwnerID: ownerID, + TemplateID: templateID, + VersionID: uuid.New(), + }) + + ctx := dbauthz.As(context.Background(), rbac.Subject{ + Type: rbac.SubjectTypeUser, + FriendlyName: "testuser", + Email: "testuser@example.com", + ID: ownerID.String(), + Roles: userRoles, + Groups: []string{orgID.String()}, + Scope: agentScope, + }.WithCachedASTValue()) + resp, err := api.BatchUpdateMetadata(ctx, req) require.NoError(t, err) require.NotNil(t, resp) diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go index 40533ea3fe..8da0d33930 100644 --- a/coderd/agentapi/stats.go +++ b/coderd/agentapi/stats.go @@ -10,6 +10,7 @@ import ( "cdr.dev/slog" agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" @@ -43,7 +44,21 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR return res, nil } - workspaceAgent, err := a.AgentFn(ctx) + // Inject RBAC object into context for dbauthz fast path, avoid having to + // call GetWorkspaceAgentByID on every stats update. + + rbacCtx := ctx + if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok { + var err error + rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject()) + if err != nil { + // Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID. + //nolint:gocritic + a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err)) + } + } + + workspaceAgent, err := a.AgentFn(rbacCtx) if err != nil { return nil, err } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 3f8138fb10..4962949f7f 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -3560,6 +3560,21 @@ func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context } func (q *querier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) { + // Fast path: Check if we have a workspace RBAC object in context. + // In the agent API this is set at agent connection time to avoid the expensive + // GetWorkspaceByAgentID query for every agent operation. + // NOTE: The cached RBAC object is refreshed every 5 minutes in agentapi/api.go. + if rbacObj, ok := WorkspaceRBACFromContext(ctx); ok { + // Errors here will result in falling back to GetWorkspaceByAgentID, + // in case the cached data is stale. + if err := q.authorizeContext(ctx, policy.ActionRead, rbacObj); err == nil { + return q.db.GetWorkspaceAgentByID(ctx, id) + } + q.log.Debug(ctx, "fast path authorization failed for GetWorkspaceAgentByID, using slow path", + slog.F("agent_id", id)) + } + + // Slow path: Fallback to fetching the workspace for authorization if _, err := q.GetWorkspaceByAgentID(ctx, id); err != nil { return database.WorkspaceAgent{}, err } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index b3a3acb890..11909b1a65 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -4805,3 +4805,81 @@ func TestGetLatestWorkspaceBuildByWorkspaceID_FastPath(t *testing.T) { require.Equal(t, build, result) }) } + +func TestGetWorkspaceAgentByID_FastPath(t *testing.T) { + t.Parallel() + + agentID := uuid.New() + ownerID := uuid.New() + wsID := uuid.New() + orgID := uuid.New() + + agent := database.WorkspaceAgent{ + ID: agentID, + Name: "test-agent", + } + + workspace := database.Workspace{ + ID: wsID, + OwnerID: ownerID, + OrganizationID: orgID, + } + + wsIdentity := database.WorkspaceIdentity{ + ID: wsID, + OwnerID: ownerID, + OrganizationID: orgID, + } + + actor := rbac.Subject{ + ID: ownerID.String(), + Roles: rbac.RoleIdentifiers{rbac.RoleOwner()}, + Groups: []string{orgID.String()}, + Scope: rbac.ScopeAll, + } + + authorizer := &coderdtest.RecordingAuthorizer{ + Wrapped: (&coderdtest.FakeAuthorizer{}).AlwaysReturn(nil), + } + + t.Run("WithWorkspaceRBAC", func(t *testing.T) { + t.Parallel() + + ctx := dbauthz.As(context.Background(), actor) + ctrl := gomock.NewController(t) + mockDB := dbmock.NewMockStore(ctrl) + + rbacObj := wsIdentity.RBACObject() + ctx, err := dbauthz.WithWorkspaceRBAC(ctx, rbacObj) + require.NoError(t, err) + + mockDB.EXPECT().Wrappers().Return([]string{}) + // GetWorkspaceByAgentID should NOT be called + mockDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(agent, nil) + + q := dbauthz.New(mockDB, authorizer, slogtest.Make(t, nil), coderdtest.AccessControlStorePointer()) + + result, err := q.GetWorkspaceAgentByID(ctx, agentID) + require.NoError(t, err) + require.Equal(t, agent, result) + }) + + t.Run("WithoutWorkspaceRBAC", func(t *testing.T) { + t.Parallel() + + ctx := dbauthz.As(context.Background(), actor) + ctrl := gomock.NewController(t) + mockDB := dbmock.NewMockStore(ctrl) + + mockDB.EXPECT().Wrappers().Return([]string{}) + // GetWorkspaceByAgentID SHOULD be called + mockDB.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agentID).Return(workspace, nil) + mockDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(agent, nil) + + q := dbauthz.New(mockDB, authorizer, slogtest.Make(t, nil), coderdtest.AccessControlStorePointer()) + + result, err := q.GetWorkspaceAgentByID(ctx, agentID) + require.NoError(t, err) + require.Equal(t, agent, result) + }) +}