Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit ae36bfa

Browse files
dannykoppingdannykopping
committed
chore: aibridge database & RBAC
1 parent f907c19 commit ae36bfa

25 files changed

+651
-0
lines changed

‎coderd/apidoc/docs.go

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,27 @@ func (q *querier) authorizePrebuiltWorkspace(ctx context.Context, action policy.
175175
return xerrors.Errorf("authorize context: %w", workspaceErr)
176176
}
177177

178+
// authorizeAIBridgeInterceptionUpdate validates that the context's actor matches the initiator of the AIBridgeInterception.
179+
// This is used by all of the sub-resources which fall under the [ResourceAibridgeInterception] umbrella.
180+
func (q *querier) authorizeAIBridgeInterceptionUpdate(ctx context.Context, sessID uuid.UUID) error {
181+
act, ok := ActorFromContext(ctx)
182+
if !ok {
183+
return ErrNoActor
184+
}
185+
186+
sess, err := q.db.GetAIBridgeInterceptionByID(ctx, sessID)
187+
if err != nil {
188+
return xerrors.Errorf("fetch aibridge session %q: %w", sessID, err)
189+
}
190+
191+
err = q.auth.Authorize(ctx, act, policy.ActionUpdate, sess.RBACObject())
192+
if err != nil {
193+
return logNotAuthorizedError(ctx, q.log, err)
194+
}
195+
196+
return nil
197+
}
198+
178199
type authContextKey struct{}
179200

180201
// ActorFromContext returns the authorization subject from the context.
@@ -542,6 +563,29 @@ var (
542563
}),
543564
Scope: rbac.ScopeAll,
544565
}.WithCachedASTValue()
566+
567+
// See aibridged package.
568+
subjectAibridged = rbac.Subject{
569+
Type: rbac.SubjectAibridged,
570+
FriendlyName: "AIBridge Daemon",
571+
ID: uuid.Nil.String(),
572+
Roles: rbac.Roles([]rbac.Role{
573+
{
574+
Identifier: rbac.RoleIdentifier{Name: "aibridged"},
575+
DisplayName: "AIBridge Daemon",
576+
Site: rbac.Permissions(map[string][]policy.Action{
577+
rbac.ResourceUser.Type: {
578+
policy.ActionReadPersonal, // Required to read users' external auth links. // TODO: this is too broad; reduce scope to just external_auth_links by creating separate resource.
579+
},
580+
rbac.ResourceApiKey.Type: {policy.ActionRead}, // Validate API keys.
581+
rbac.ResourceAibridgeInterception.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate},
582+
}),
583+
Org: map[string][]rbac.Permission{},
584+
User: []rbac.Permission{},
585+
},
586+
}),
587+
Scope: rbac.ScopeAll,
588+
}.WithCachedASTValue()
545589
)
546590

547591
// AsProvisionerd returns a context with an actor that has permissions required
@@ -624,6 +668,12 @@ func AsUsagePublisher(ctx context.Context) context.Context {
624668
return As(ctx, subjectUsagePublisher)
625669
}
626670

671+
// AsAIBridged returns a context with an actor that has permissions
672+
// required for creating, reading, and updating aibridge-related resources.
673+
func AsAIBridged(ctx context.Context) context.Context {
674+
return As(ctx, subjectAibridged)
675+
}
676+
627677
var AsRemoveActor = rbac.Subject{
628678
ID: "remove-actor",
629679
}
@@ -1866,6 +1916,10 @@ func (q *querier) FindMatchingPresetID(ctx context.Context, arg database.FindMat
18661916
return q.db.FindMatchingPresetID(ctx, arg)
18671917
}
18681918

1919+
func (q *querier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error) {
1920+
return fetch(q.log, q.auth, q.db.GetAIBridgeInterceptionByID)(ctx, id)
1921+
}
1922+
18691923
func (q *querier) GetAPIKeyByID(ctx context.Context, id string) (database.APIKey, error) {
18701924
return fetch(q.log, q.auth, q.db.GetAPIKeyByID)(ctx, id)
18711925
}
@@ -3745,6 +3799,34 @@ func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context, now ti
37453799
return q.db.GetWorkspacesEligibleForTransition(ctx, now)
37463800
}
37473801

3802+
func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) {
3803+
return insert(q.log, q.auth, rbac.ResourceAibridgeInterception.WithOwner(arg.InitiatorID.String()), q.db.InsertAIBridgeInterception)(ctx, arg)
3804+
}
3805+
3806+
func (q *querier) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error {
3807+
// All aibridge_token_usages records belong to the initiator of their associated session.
3808+
if err := q.authorizeAIBridgeInterceptionUpdate(ctx, arg.InterceptionID); err != nil {
3809+
return err
3810+
}
3811+
return q.db.InsertAIBridgeTokenUsage(ctx, arg)
3812+
}
3813+
3814+
func (q *querier) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error {
3815+
// All aibridge_tool_usages records belong to the initiator of their associated session.
3816+
if err := q.authorizeAIBridgeInterceptionUpdate(ctx, arg.InterceptionID); err != nil {
3817+
return err
3818+
}
3819+
return q.db.InsertAIBridgeToolUsage(ctx, arg)
3820+
}
3821+
3822+
func (q *querier) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error {
3823+
// All aibridge_user_prompts records belong to the initiator of their associated session.
3824+
if err := q.authorizeAIBridgeInterceptionUpdate(ctx, arg.InterceptionID); err != nil {
3825+
return err
3826+
}
3827+
return q.db.InsertAIBridgeUserPrompt(ctx, arg)
3828+
}
3829+
37483830
func (q *querier) InsertAPIKey(ctx context.Context, arg database.InsertAPIKeyParams) (database.APIKey, error) {
37493831
// TODO(Cian): ideally this would be encoded in the policy, but system users are just members and we
37503832
// don't currently have a capability to conditionally deny creating resources by owner ID in a role.

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4326,3 +4326,55 @@ func TestInsertAPIKey_AsPrebuildsUser(t *testing.T) {
43264326
_, err := dbz.InsertAPIKey(ctx, testutil.Fake(t, faker, database.InsertAPIKeyParams{}))
43274327
require.True(t, dbauthz.IsNotAuthorizedError(err))
43284328
}
4329+
4330+
func (s *MethodTestSuite) TestAIBridge() {
4331+
s.Run("GetAIBridgeInterceptionByID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
4332+
sessID := uuid.UUID{2}
4333+
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID})
4334+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), sessID).Return(sess, nil).AnyTimes()
4335+
check.Args(sessID).Asserts(sess, policy.ActionRead).Returns(sess)
4336+
}))
4337+
4338+
s.Run("InsertAIBridgeInterception", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
4339+
initID := uuid.UUID{3}
4340+
user := testutil.Fake(s.T(), faker, database.User{ID: initID})
4341+
// testutil.Fake cannot distinguish between a zero value and an explicitly requested value which is equivalent.
4342+
user.IsSystem = false
4343+
user.Deleted = false
4344+
4345+
sessID := uuid.UUID{2}
4346+
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID, InitiatorID: initID})
4347+
4348+
params := database.InsertAIBridgeInterceptionParams{ID: sess.ID, InitiatorID: sess.InitiatorID, Provider: sess.Provider, Model: sess.Model}
4349+
db.EXPECT().GetUserByID(gomock.Any(), initID).Return(user, nil).AnyTimes() // Validation.
4350+
db.EXPECT().InsertAIBridgeInterception(gomock.Any(), params).Return(sess, nil).AnyTimes()
4351+
check.Args(params).Asserts(sess, policy.ActionCreate)
4352+
}))
4353+
4354+
s.Run("InsertAIBridgeTokenUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
4355+
sessID := uuid.UUID{2}
4356+
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID})
4357+
params := database.InsertAIBridgeTokenUsageParams{InterceptionID: sess.ID}
4358+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), sessID).Return(sess, nil).AnyTimes() // Validation.
4359+
db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), params).Return(nil).AnyTimes()
4360+
check.Args(params).Asserts(sess, policy.ActionUpdate)
4361+
}))
4362+
4363+
s.Run("InsertAIBridgeToolUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
4364+
sessID := uuid.UUID{2}
4365+
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID})
4366+
params := database.InsertAIBridgeToolUsageParams{InterceptionID: sess.ID}
4367+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), sessID).Return(sess, nil).AnyTimes() // Validation.
4368+
db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), params).Return(nil).AnyTimes()
4369+
check.Args(params).Asserts(sess, policy.ActionUpdate)
4370+
}))
4371+
4372+
s.Run("InsertAIBridgeUserPrompt", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
4373+
sessID := uuid.UUID{2}
4374+
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID})
4375+
params := database.InsertAIBridgeUserPromptParams{InterceptionID: sess.ID}
4376+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), sessID).Return(sess, nil).AnyTimes() // Validation.
4377+
db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), params).Return(nil).AnyTimes()
4378+
check.Args(params).Asserts(sess, policy.ActionUpdate)
4379+
}))
4380+
}

‎coderd/database/dbmetrics/querymetrics.go

Lines changed: 35 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go

Lines changed: 72 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)