fix: stop group spend limits from leaking across org boundaries (#24294)

Three SQL queries (`GetUserGroupSpendLimit`,
`ResolveUserChatSpendLimit`, `GetUserChatSpendInPeriod`) aggregated chat
spend limits and usage globally across all organizations. A restrictive
group limit in org A would bleed into org B.

## Changes

- Add `organization_id` parameter to all three SQL queries in
`coderd/database/queries/chats.sql`
- When nil UUID is passed, queries fall back to global behavior
(backward compat for HTTP dashboard endpoints)
- When real org ID is passed, limits and spend are scoped to that
organization
- Thread `organizationID` through `ResolveUsageLimitStatus` →
`checkUsageLimit` → all chatd call sites
- Update dbauthz wrappers for new param structs
- HTTP endpoints (`chatCostSummary`, `getMyChatUsageLimitStatus`) pass
`uuid.Nil` with TODO for future org-scoped UI
- Add `TestResolveUsageLimitStatus_OrgScoped` with 5 test cases covering
org isolation, nil-UUID fallback, spend scoping, and user override
priority

Closes coder/internal#1466

> 🤖
This commit is contained in:
Cian Johnston
2026-04-14 16:56:17 +01:00
committed by GitHub
parent 4d4266a4ad
commit c552f9f281
12 changed files with 542 additions and 79 deletions
+7 -7
View File
@@ -4222,11 +4222,11 @@ func (q *querier) GetUserCount(ctx context.Context, includeSystem bool) (int64,
return q.db.GetUserCount(ctx, includeSystem)
}
func (q *querier) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(userID.String())); err != nil {
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
}
return q.db.GetUserGroupSpendLimit(ctx, userID)
return q.db.GetUserGroupSpendLimit(ctx, arg)
}
func (q *querier) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
@@ -5816,11 +5816,11 @@ func (q *querier) RemoveUserFromGroups(ctx context.Context, arg database.RemoveU
return q.db.RemoveUserFromGroups(ctx, arg)
}
func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(userID.String())); err != nil {
return 0, err
func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil {
return database.ResolveUserChatSpendLimitRow{}, err
}
return q.db.ResolveUserChatSpendLimit(ctx, userID)
return q.db.ResolveUserChatSpendLimit(ctx, arg)
}
func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
+18 -8
View File
@@ -1125,7 +1125,9 @@ func (s *MethodTestSuite) TestChats() {
}))
s.Run("GetUserChatSpendInPeriod", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
arg := database.GetUserChatSpendInPeriodParams{
UserID: uuid.New(),
UserID: uuid.New(),
OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
StartTime: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
EndTime: time.Date(2025, 2, 1, 0, 0, 0, 0, time.UTC),
}
@@ -1134,17 +1136,25 @@ func (s *MethodTestSuite) TestChats() {
check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(spend)
}))
s.Run("GetUserGroupSpendLimit", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
userID := uuid.New()
arg := database.GetUserGroupSpendLimitParams{
UserID: uuid.New(),
OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
}
limit := int64(456)
dbm.EXPECT().GetUserGroupSpendLimit(gomock.Any(), userID).Return(limit, nil).AnyTimes()
check.Args(userID).Asserts(rbac.ResourceChat.WithOwner(userID.String()), policy.ActionRead).Returns(limit)
dbm.EXPECT().GetUserGroupSpendLimit(gomock.Any(), arg).Return(limit, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(limit)
}))
s.Run("ResolveUserChatSpendLimit", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
userID := uuid.New()
limit := int64(789)
dbm.EXPECT().ResolveUserChatSpendLimit(gomock.Any(), userID).Return(limit, nil).AnyTimes()
check.Args(userID).Asserts(rbac.ResourceChat.WithOwner(userID.String()), policy.ActionRead).Returns(limit)
arg := database.ResolveUserChatSpendLimitParams{
UserID: uuid.New(),
OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
}
row := database.ResolveUserChatSpendLimitRow{EffectiveLimitMicros: 789, LimitSource: "group"}
dbm.EXPECT().ResolveUserChatSpendLimit(gomock.Any(), arg).Return(row, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(row)
}))
s.Run("GetChatUsageLimitConfig", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
now := dbtime.Now()
config := database.ChatUsageLimitConfig{
+2 -2
View File
@@ -2704,7 +2704,7 @@ func (m queryMetricsStore) GetUserCount(ctx context.Context, includeSystem bool)
return r0, r1
}
func (m queryMetricsStore) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
func (m queryMetricsStore) GetUserGroupSpendLimit(ctx context.Context, userID database.GetUserGroupSpendLimitParams) (int64, error) {
start := time.Now()
r0, r1 := m.s.GetUserGroupSpendLimit(ctx, userID)
m.queryLatencies.WithLabelValues("GetUserGroupSpendLimit").Observe(time.Since(start).Seconds())
@@ -4160,7 +4160,7 @@ func (m queryMetricsStore) RemoveUserFromGroups(ctx context.Context, arg databas
return r0, r1
}
func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
start := time.Now()
r0, r1 := m.s.ResolveUserChatSpendLimit(ctx, userID)
m.queryLatencies.WithLabelValues("ResolveUserChatSpendLimit").Observe(time.Since(start).Seconds())
+9 -9
View File
@@ -5058,18 +5058,18 @@ func (mr *MockStoreMockRecorder) GetUserCount(ctx, includeSystem any) *gomock.Ca
}
// GetUserGroupSpendLimit mocks base method.
func (m *MockStore) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
func (m *MockStore) GetUserGroupSpendLimit(ctx context.Context, arg database.GetUserGroupSpendLimitParams) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetUserGroupSpendLimit", ctx, userID)
ret := m.ctrl.Call(m, "GetUserGroupSpendLimit", ctx, arg)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetUserGroupSpendLimit indicates an expected call of GetUserGroupSpendLimit.
func (mr *MockStoreMockRecorder) GetUserGroupSpendLimit(ctx, userID any) *gomock.Call {
func (mr *MockStoreMockRecorder) GetUserGroupSpendLimit(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserGroupSpendLimit", reflect.TypeOf((*MockStore)(nil).GetUserGroupSpendLimit), ctx, userID)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserGroupSpendLimit", reflect.TypeOf((*MockStore)(nil).GetUserGroupSpendLimit), ctx, arg)
}
// GetUserLatencyInsights mocks base method.
@@ -7889,18 +7889,18 @@ func (mr *MockStoreMockRecorder) RemoveUserFromGroups(ctx, arg any) *gomock.Call
}
// ResolveUserChatSpendLimit mocks base method.
func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ResolveUserChatSpendLimit", ctx, userID)
ret0, _ := ret[0].(int64)
ret := m.ctrl.Call(m, "ResolveUserChatSpendLimit", ctx, arg)
ret0, _ := ret[0].(database.ResolveUserChatSpendLimitRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ResolveUserChatSpendLimit indicates an expected call of ResolveUserChatSpendLimit.
func (mr *MockStoreMockRecorder) ResolveUserChatSpendLimit(ctx, userID any) *gomock.Call {
func (mr *MockStoreMockRecorder) ResolveUserChatSpendLimit(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResolveUserChatSpendLimit", reflect.TypeOf((*MockStore)(nil).ResolveUserChatSpendLimit), ctx, userID)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResolveUserChatSpendLimit", reflect.TypeOf((*MockStore)(nil).ResolveUserChatSpendLimit), ctx, arg)
}
// RevokeDBCryptKey mocks base method.
+18 -5
View File
@@ -642,11 +642,18 @@ type sqlcQuerier interface {
GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID) (string, error)
GetUserChatDebugLoggingEnabled(ctx context.Context, userID uuid.UUID) (bool, error)
GetUserChatProviderKeys(ctx context.Context, userID uuid.UUID) ([]UserChatProviderKey, error)
// Returns the total spend for a user in the given period.
// When organization_id is NULL, spend across all organizations is
// returned (global behavior). Otherwise only spend within the
// specified organization is included.
GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error)
GetUserCount(ctx context.Context, includeSystem bool) (int64, error)
// Returns the minimum (most restrictive) group limit for a user.
// Returns -1 if the user has no group limits applied.
GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error)
// Returns -1 if no group limits match the specified scope.
// When organization_id is NULL, groups across all organizations are
// considered (global behavior). Otherwise only groups within the
// specified organization are considered.
GetUserGroupSpendLimit(ctx context.Context, arg GetUserGroupSpendLimitParams) (int64, error)
// GetUserLatencyInsights returns the median and 95th percentile connection
// latency that users have experienced. The result can be filtered on
// template_ids, meaning only user data from workspaces based on those templates
@@ -907,11 +914,17 @@ type sqlcQuerier interface {
RegisterWorkspaceProxy(ctx context.Context, arg RegisterWorkspaceProxyParams) (WorkspaceProxy, error)
RemoveUserFromGroups(ctx context.Context, arg RemoveUserFromGroupsParams) ([]uuid.UUID, error)
// Resolves the effective spend limit for a user using the hierarchy:
// 1. Individual user override (highest priority)
// 2. Minimum group limit across all user's groups
// 1. Individual user override (highest priority, applies globally across
// all organizations since it lives on the users table)
// 2. Minimum group limit across the user's groups
// 3. Global default from config
// Returns -1 if limits are not enabled.
ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error)
// When organization_id is NULL, groups across all organizations are
// considered (global behavior). Otherwise only groups within the
// specified organization are considered.
// limit_source indicates which tier won: 'user', 'group', 'default',
// or 'disabled'.
ResolveUserChatSpendLimit(ctx context.Context, arg ResolveUserChatSpendLimitParams) (ResolveUserChatSpendLimitRow, error)
RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error
// Note that this selects from the CTE, not the original table. The CTE is named
// the same as the original table to trick sqlc into reusing the existing struct
+65 -22
View File
@@ -6747,19 +6747,31 @@ SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros
FROM chat_messages cm
JOIN chats c ON c.id = cm.chat_id
WHERE c.owner_id = $1::uuid
AND cm.created_at >= $2::timestamptz
AND cm.created_at < $3::timestamptz
AND ($2::uuid IS NULL
OR c.organization_id = $2::uuid)
AND cm.created_at >= $3::timestamptz
AND cm.created_at < $4::timestamptz
AND cm.total_cost_micros IS NOT NULL
`
type GetUserChatSpendInPeriodParams struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
StartTime time.Time `db:"start_time" json:"start_time"`
EndTime time.Time `db:"end_time" json:"end_time"`
UserID uuid.UUID `db:"user_id" json:"user_id"`
OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"`
StartTime time.Time `db:"start_time" json:"start_time"`
EndTime time.Time `db:"end_time" json:"end_time"`
}
// Returns the total spend for a user in the given period.
// When organization_id is NULL, spend across all organizations is
// returned (global behavior). Otherwise only spend within the
// specified organization is included.
func (q *sqlQuerier) GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error) {
row := q.db.QueryRowContext(ctx, getUserChatSpendInPeriod, arg.UserID, arg.StartTime, arg.EndTime)
row := q.db.QueryRowContext(ctx, getUserChatSpendInPeriod,
arg.UserID,
arg.OrganizationID,
arg.StartTime,
arg.EndTime,
)
var total_spend_micros int64
err := row.Scan(&total_spend_micros)
return total_spend_micros, err
@@ -6770,13 +6782,23 @@ SELECT COALESCE(MIN(g.chat_spend_limit_micros), -1)::bigint AS limit_micros
FROM groups g
JOIN group_members_expanded gme ON gme.group_id = g.id
WHERE gme.user_id = $1::uuid
AND ($2::uuid IS NULL
OR g.organization_id = $2::uuid)
AND g.chat_spend_limit_micros IS NOT NULL
`
type GetUserGroupSpendLimitParams struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"`
}
// Returns the minimum (most restrictive) group limit for a user.
// Returns -1 if the user has no group limits applied.
func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
row := q.db.QueryRowContext(ctx, getUserGroupSpendLimit, userID)
// Returns -1 if no group limits match the specified scope.
// When organization_id is NULL, groups across all organizations are
// considered (global behavior). Otherwise only groups within the
// specified organization are considered.
func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, arg GetUserGroupSpendLimitParams) (int64, error) {
row := q.db.QueryRowContext(ctx, getUserGroupSpendLimit, arg.UserID, arg.OrganizationID)
var limit_micros int64
err := row.Scan(&limit_micros)
return limit_micros, err
@@ -7300,15 +7322,17 @@ func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID)
const resolveUserChatSpendLimit = `-- name: ResolveUserChatSpendLimit :one
SELECT CASE
-- If limits are disabled, return -1.
WHEN NOT cfg.enabled THEN -1
-- Individual override takes priority.
WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros
-- Group limit (minimum across all user's groups) is next.
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
-- Fall back to global default.
ELSE cfg.default_limit_micros
END::bigint AS effective_limit_micros
END::bigint AS effective_limit_micros,
CASE
WHEN NOT cfg.enabled THEN 'disabled'
WHEN u.chat_spend_limit_micros IS NOT NULL THEN 'user'
WHEN gl.limit_micros IS NOT NULL THEN 'group'
ELSE 'default'
END AS limit_source
FROM chat_usage_limit_config cfg
CROSS JOIN users u
LEFT JOIN LATERAL (
@@ -7316,22 +7340,41 @@ LEFT JOIN LATERAL (
FROM groups g
JOIN group_members_expanded gme ON gme.group_id = g.id
WHERE gme.user_id = $1::uuid
AND ($2::uuid IS NULL
OR g.organization_id = $2::uuid)
AND g.chat_spend_limit_micros IS NOT NULL
) gl ON TRUE
WHERE u.id = $1::uuid
LIMIT 1
`
type ResolveUserChatSpendLimitParams struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"`
}
type ResolveUserChatSpendLimitRow struct {
EffectiveLimitMicros int64 `db:"effective_limit_micros" json:"effective_limit_micros"`
LimitSource string `db:"limit_source" json:"limit_source"`
}
// Resolves the effective spend limit for a user using the hierarchy:
// 1. Individual user override (highest priority)
// 2. Minimum group limit across all user's groups
// 3. Global default from config
// 1. Individual user override (highest priority, applies globally across
// all organizations since it lives on the users table)
// 2. Minimum group limit across the user's groups
// 3. Global default from config
//
// Returns -1 if limits are not enabled.
func (q *sqlQuerier) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
row := q.db.QueryRowContext(ctx, resolveUserChatSpendLimit, userID)
var effective_limit_micros int64
err := row.Scan(&effective_limit_micros)
return effective_limit_micros, err
// When organization_id is NULL, groups across all organizations are
// considered (global behavior). Otherwise only groups within the
// specified organization are considered.
// limit_source indicates which tier won: 'user', 'group', 'default',
// or 'disabled'.
func (q *sqlQuerier) ResolveUserChatSpendLimit(ctx context.Context, arg ResolveUserChatSpendLimitParams) (ResolveUserChatSpendLimitRow, error) {
row := q.db.QueryRowContext(ctx, resolveUserChatSpendLimit, arg.UserID, arg.OrganizationID)
var i ResolveUserChatSpendLimitRow
err := row.Scan(&i.EffectiveLimitMicros, &i.LimitSource)
return i, err
}
const softDeleteChatMessageByID = `-- name: SoftDeleteChatMessageByID :exec
+29 -8
View File
@@ -1134,10 +1134,16 @@ FROM users
WHERE id = @user_id::uuid AND chat_spend_limit_micros IS NOT NULL;
-- name: GetUserChatSpendInPeriod :one
-- Returns the total spend for a user in the given period.
-- When organization_id is NULL, spend across all organizations is
-- returned (global behavior). Otherwise only spend within the
-- specified organization is included.
SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros
FROM chat_messages cm
JOIN chats c ON c.id = cm.chat_id
WHERE c.owner_id = @user_id::uuid
AND (sqlc.narg('organization_id')::uuid IS NULL
OR c.organization_id = sqlc.narg('organization_id')::uuid)
AND cm.created_at >= @start_time::timestamptz
AND cm.created_at < @end_time::timestamptz
AND cm.total_cost_micros IS NOT NULL;
@@ -1189,11 +1195,16 @@ WHERE id = @group_id::uuid AND chat_spend_limit_micros IS NOT NULL;
-- name: GetUserGroupSpendLimit :one
-- Returns the minimum (most restrictive) group limit for a user.
-- Returns -1 if the user has no group limits applied.
-- Returns -1 if no group limits match the specified scope.
-- When organization_id is NULL, groups across all organizations are
-- considered (global behavior). Otherwise only groups within the
-- specified organization are considered.
SELECT COALESCE(MIN(g.chat_spend_limit_micros), -1)::bigint AS limit_micros
FROM groups g
JOIN group_members_expanded gme ON gme.group_id = g.id
WHERE gme.user_id = @user_id::uuid
AND (sqlc.narg('organization_id')::uuid IS NULL
OR g.organization_id = sqlc.narg('organization_id')::uuid)
AND g.chat_spend_limit_micros IS NOT NULL;
-- name: GetChatsByWorkspaceIDs :many
@@ -1205,20 +1216,28 @@ ORDER BY workspace_id, updated_at DESC;
-- name: ResolveUserChatSpendLimit :one
-- Resolves the effective spend limit for a user using the hierarchy:
-- 1. Individual user override (highest priority)
-- 2. Minimum group limit across all user's groups
-- 1. Individual user override (highest priority, applies globally across
-- all organizations since it lives on the users table)
-- 2. Minimum group limit across the user's groups
-- 3. Global default from config
-- Returns -1 if limits are not enabled.
-- When organization_id is NULL, groups across all organizations are
-- considered (global behavior). Otherwise only groups within the
-- specified organization are considered.
-- limit_source indicates which tier won: 'user', 'group', 'default',
-- or 'disabled'.
SELECT CASE
-- If limits are disabled, return -1.
WHEN NOT cfg.enabled THEN -1
-- Individual override takes priority.
WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros
-- Group limit (minimum across all user's groups) is next.
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
-- Fall back to global default.
ELSE cfg.default_limit_micros
END::bigint AS effective_limit_micros
END::bigint AS effective_limit_micros,
CASE
WHEN NOT cfg.enabled THEN 'disabled'
WHEN u.chat_spend_limit_micros IS NOT NULL THEN 'user'
WHEN gl.limit_micros IS NOT NULL THEN 'group'
ELSE 'default'
END AS limit_source
FROM chat_usage_limit_config cfg
CROSS JOIN users u
LEFT JOIN LATERAL (
@@ -1226,6 +1245,8 @@ LEFT JOIN LATERAL (
FROM groups g
JOIN group_members_expanded gme ON gme.group_id = g.id
WHERE gme.user_id = @user_id::uuid
AND (sqlc.narg('organization_id')::uuid IS NULL
OR g.organization_id = sqlc.narg('organization_id')::uuid)
AND g.chat_spend_limit_micros IS NOT NULL
) gl ON TRUE
WHERE u.id = @user_id::uuid