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

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions coderd/database/dbauthz/dbauthz.go
Original file line number Diff line number Diff line change
Expand Up @@ -473,6 +473,27 @@ var (
}.WithCachedASTValue()
}

subjectChatdKeyMinter = func(userID uuid.UUID) rbac.Subject {
Comment thread
ibetitsmike marked this conversation as resolved.
return rbac.Subject{
Type: rbac.SubjectTypeChatdKeyMinter,
FriendlyName: "Chatd Key Minter",
ID: userID.String(),
Roles: rbac.Roles([]rbac.Role{
{
Identifier: rbac.RoleIdentifier{Name: "chatdkeyminter"},
DisplayName: "Chatd Key Minter",
Site: []rbac.Permission{},
User: rbac.Permissions(map[string][]policy.Action{
rbac.ResourceApiKey.Type: {policy.ActionRead, policy.ActionCreate, policy.ActionUpdate, policy.ActionDelete},
rbac.ResourceUser.Type: {policy.ActionReadPersonal},
}),
ByOrgID: map[string]rbac.OrgPermissions{},
},
}),
Scope: rbac.ScopeAll,
}.WithCachedASTValue()
}

subjectSystemRestricted = rbac.Subject{
Type: rbac.SubjectTypeSystemRestricted,
FriendlyName: "System",
Expand Down Expand Up @@ -874,6 +895,12 @@ func AsAPIKeyRevoker(ctx context.Context, userID uuid.UUID) context.Context {
return As(ctx, subjectAPIKeyRevoker(userID))
}

// AsChatdKeyMinter returns a context with an actor that manages the synthetic
// gateway API key owned by the specified user.
func AsChatdKeyMinter(ctx context.Context, userID uuid.UUID) context.Context {
return As(ctx, subjectChatdKeyMinter(userID))
}

// AsSystemRestricted returns a context with an actor that has permissions
// required for various system operations (login, logout, metrics cache).
// DO NOT USE THIS UNLESS YOU HAVE ABSOLUTELY NO OTHER CHOICE. Prefer using a
Expand Down Expand Up @@ -3315,6 +3342,10 @@ func (q *querier) GetChatFilesByIDs(ctx context.Context, ids []uuid.UUID) ([]dat
return files, nil
}

func (q *querier) GetChatGatewayAPIKey(ctx context.Context, arg database.GetChatGatewayAPIKeyParams) (database.APIKey, error) {
return fetch(q.log, q.auth, q.db.GetChatGatewayAPIKey)(ctx, arg)
}

func (q *querier) GetChatGeneralModelOverride(ctx context.Context) (string, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return "", err
Expand Down Expand Up @@ -5072,6 +5103,10 @@ func (q *querier) GetUserCount(ctx context.Context, includeSystem bool) (int64,
return q.db.GetUserCount(ctx, includeSystem)
}

func (q *querier) GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid.UUID) (database.User, error) {
return fetchWithAction(q.log, q.auth, policy.ActionReadPersonal, q.db.GetUserForChatSyntheticAPIKeyByID)(ctx, id)
}

func (q *querier) GetUserGroupSpendLimit(ctx context.Context, arg database.GetUserGroupSpendLimitParams) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil {
return 0, err
Expand Down
32 changes: 32 additions & 0 deletions coderd/database/dbauthz/dbauthz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -339,6 +339,20 @@ func defaultIPAddress() pqtype.Inet {
}
}

func (s *MethodTestSuite) TestChatGatewayAPIKey() {
s.Run("GetUserForChatSyntheticAPIKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
user := testutil.Fake(s.T(), faker, database.User{})
dbm.EXPECT().GetUserForChatSyntheticAPIKeyByID(gomock.Any(), user.ID).Return(user, nil).AnyTimes()
check.Args(user.ID).Asserts(user, policy.ActionReadPersonal).Returns(user)
}))
s.Run("GetChatGatewayAPIKey", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
key := testutil.Fake(s.T(), faker, database.APIKey{})
arg := database.GetChatGatewayAPIKeyParams{UserID: key.UserID, TokenName: key.TokenName}
dbm.EXPECT().GetChatGatewayAPIKey(gomock.Any(), arg).Return(key, nil).AnyTimes()
check.Args(arg).Asserts(key, policy.ActionRead).Returns(key)
}))
}

func (s *MethodTestSuite) TestAPIKey() {
s.Run("DeleteAPIKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
key := testutil.Fake(s.T(), faker, database.APIKey{})
Expand Down Expand Up @@ -7481,6 +7495,24 @@ func TestAsAPIKeyRevoker(t *testing.T) {
})
}

func TestAsChatdKeyMinter(t *testing.T) {
t.Parallel()

userID := uuid.New()
ctx := dbauthz.AsChatdKeyMinter(context.Background(), userID)
actor, ok := dbauthz.ActorFromContext(ctx)
require.True(t, ok)
require.Equal(t, rbac.SubjectTypeChatdKeyMinter, actor.Type)
require.Equal(t, userID.String(), actor.ID)

auth := rbac.NewStrictCachingAuthorizer(prometheus.NewRegistry())
for _, action := range []policy.Action{policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete} {
require.NoError(t, auth.Authorize(ctx, actor, action, rbac.ResourceApiKey.WithOwner(userID.String())))
require.Error(t, auth.Authorize(ctx, actor, action, rbac.ResourceApiKey.WithOwner(uuid.NewString())))
}
require.NoError(t, auth.Authorize(ctx, actor, policy.ActionReadPersonal, rbac.ResourceUserObject(userID)))
}

func TestAsChatd(t *testing.T) {
t.Parallel()

Expand Down
16 changes: 16 additions & 0 deletions coderd/database/dbmetrics/querymetrics.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

30 changes: 30 additions & 0 deletions coderd/database/dbmock/dbmock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 0 additions & 6 deletions coderd/database/dump.sql

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 0 additions & 2 deletions coderd/database/foreign_key_constraint.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
UPDATE chat_messages
SET api_key_id = NULL
WHERE api_key_id IS NOT NULL
AND NOT EXISTS (
SELECT 1
FROM api_keys
WHERE api_keys.id = chat_messages.api_key_id
);

UPDATE chat_queued_messages
SET api_key_id = NULL
WHERE api_key_id IS NOT NULL
AND NOT EXISTS (
SELECT 1
FROM api_keys
WHERE api_keys.id = chat_queued_messages.api_key_id
);

ALTER TABLE chat_messages
ADD CONSTRAINT chat_messages_api_key_id_fkey
FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL;

ALTER TABLE chat_queued_messages
ADD CONSTRAINT chat_queued_messages_api_key_id_fkey
FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL;
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
ALTER TABLE chat_messages
DROP CONSTRAINT chat_messages_api_key_id_fkey;

ALTER TABLE chat_queued_messages
DROP CONSTRAINT chat_queued_messages_api_key_id_fkey;
59 changes: 59 additions & 0 deletions coderd/database/migrations/migrate_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1658,6 +1658,65 @@ func TestMigration000542ChatReasoningEffortBackfill(t *testing.T) {
require.Equal(t, sql.NullString{}, got["bedrock:anthropic.invalid-effort"])
}

func TestMigration000546ChatHistoryAPIKeyConstraints(t *testing.T) {
t.Parallel()

const priorMigrationVersion = 545

sqlDB := testSQLDB(t)
next, err := migrations.Stepper(sqlDB)
require.NoError(t, err)
for {
version, more, err := next()
require.NoError(t, err)
if !more || version == priorMigrationVersion {
break
}
}

ctx := testutil.Context(t, testutil.WaitSuperLong)
constraintNames := []string{
"chat_messages_api_key_id_fkey",
"chat_queued_messages_api_key_id_fkey",
}
assertConstraintCount := func(t *testing.T, want int) {
t.Helper()
for _, name := range constraintNames {
var got int
err := sqlDB.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM pg_constraint
WHERE conname = $1
`, name).Scan(&got)
require.NoError(t, err)
require.Equal(t, want, got, name)
}
}

upSQL, err := os.ReadFile("000546_drop_chat_history_api_key_fks.up.sql")
require.NoError(t, err)
_, err = sqlDB.ExecContext(ctx, string(upSQL))
require.NoError(t, err)
assertConstraintCount(t, 0)

downSQL, err := os.ReadFile("000546_drop_chat_history_api_key_fks.down.sql")
require.NoError(t, err)
_, err = sqlDB.ExecContext(ctx, string(downSQL))
require.NoError(t, err)
assertConstraintCount(t, 1)

for _, name := range constraintNames {
var count int
err := sqlDB.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM pg_constraint
WHERE conname = $1 AND confdeltype = 'n'
`, name).Scan(&count)
require.NoError(t, err)
require.Equal(t, 1, count, name)
}
}

func TestMigration000498SoftDeleteStaleWorkspaceAgents(t *testing.T) {
t.Parallel()

Expand Down
2 changes: 2 additions & 0 deletions coderd/database/querier.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading