mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat(coderd/database): add AI Gateway key auth lookup and last-used queries (#26505)
Adds DB methods`GetAIGatewayKeyIDByHashedSecret` and `UpdateAIGatewayKeyLastUsedAt`. `GetAIGatewayKeyIDByHashedSecret` - returns AI Gateway key ID by hashed secret value. `UpdateAIGatewayKeyLastUsedAt` - updates last used timestamp for given AI Gateway key. Used by standalone AI Gateway for authentication and keeping track of currently used keys.
This commit is contained in:
@@ -486,6 +486,7 @@ var (
|
||||
rbac.ResourceOauth2AppSecret.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
|
||||
rbac.ResourceChat.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
|
||||
rbac.ResourceAIProvider.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
|
||||
rbac.ResourceAIGatewayKey.Type: {policy.ActionRead, policy.ActionUpdate},
|
||||
}),
|
||||
User: []rbac.Permission{},
|
||||
ByOrgID: map[string]rbac.OrgPermissions{},
|
||||
@@ -2753,6 +2754,16 @@ func (q *querier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, in
|
||||
return q.db.GetAIBridgeUserPromptsByInterceptionID(ctx, interceptionID)
|
||||
}
|
||||
|
||||
// Authenticates a standalone AI Gateway replica by its hashed key secret, returning the matched key.
|
||||
func (q *querier) GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (database.AIGatewayKey, error) {
|
||||
// Standalone AI Gateway has no Coder identity, so this runs under the
|
||||
// system actor reading the AI Gateway key it authenticates against.
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIGatewayKey); err != nil {
|
||||
return database.AIGatewayKey{}, err
|
||||
}
|
||||
return q.db.GetAIGatewayKeyByHashedSecret(ctx, hashedSecret)
|
||||
}
|
||||
|
||||
func (q *querier) GetAIModelPriceByProviderModel(ctx context.Context, arg database.GetAIModelPriceByProviderModelParams) (database.AIModelPrice, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAiModelPrice); err != nil {
|
||||
return database.AIModelPrice{}, err
|
||||
@@ -7022,6 +7033,16 @@ func (q *querier) UpdateAIBridgeInterceptionEnded(ctx context.Context, params da
|
||||
return q.db.UpdateAIBridgeInterceptionEnded(ctx, params)
|
||||
}
|
||||
|
||||
// Records liveness for a key used in active DRPC session between coderd and standalone AI Gateway.
|
||||
func (q *querier) UpdateAIGatewayKeyLastUsedAt(ctx context.Context, id uuid.UUID) (int64, error) {
|
||||
// Standalone AI Gateway has no Coder identity, so this runs under the
|
||||
// system actor recording connection liveness on the AI Gateway key.
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIGatewayKey); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return q.db.UpdateAIGatewayKeyLastUsedAt(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) UpdateAIProvider(ctx context.Context, arg database.UpdateAIProviderParams) (database.AIProvider, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIProvider); err != nil {
|
||||
return database.AIProvider{}, err
|
||||
|
||||
@@ -6968,6 +6968,17 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
dbm.EXPECT().DeleteAIGatewayKey(gomock.Any(), id).Return(database.DeleteAIGatewayKeyRow{}, nil).AnyTimes()
|
||||
check.Args(id).Asserts(rbac.ResourceAIGatewayKey, policy.ActionDelete).Returns(database.DeleteAIGatewayKeyRow{})
|
||||
}))
|
||||
s.Run("GetAIGatewayKeyByHashedSecret", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
hashedSecret := []byte("hashed-secret")
|
||||
key := database.AIGatewayKey{ID: uuid.New(), HashedSecret: hashedSecret}
|
||||
dbm.EXPECT().GetAIGatewayKeyByHashedSecret(gomock.Any(), hashedSecret).Return(key, nil).AnyTimes()
|
||||
check.Args(hashedSecret).Asserts(rbac.ResourceAIGatewayKey, policy.ActionRead).Returns(key)
|
||||
}))
|
||||
s.Run("UpdateAIGatewayKeyLastUsedAt", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
id := uuid.New()
|
||||
dbm.EXPECT().UpdateAIGatewayKeyLastUsedAt(gomock.Any(), id).Return(int64(1), nil).AnyTimes()
|
||||
check.Args(id).Asserts(rbac.ResourceAIGatewayKey, policy.ActionUpdate).Returns(int64(1))
|
||||
}))
|
||||
}
|
||||
|
||||
func (s *MethodTestSuite) TestTelemetry() {
|
||||
|
||||
Reference in New Issue
Block a user