fix(core): validate batch usage consumption in quota guard (#7939)

xiaoyijun · web-flow · commit 7c14743b0932 · 2025-11-04T10:25:50.000+08:00
diff --git a/packages/core/src/libraries/quota.test.ts b/packages/core/src/libraries/quota.test.ts
@@ -1,3 +1,4 @@
+/* eslint-disable max-lines */
 import { ConnectorType, ReservedPlanId } from '@logto/schemas';
 import { createMockUtils } from '@logto/shared/esm';
 
@@ -313,7 +314,7 @@ describe('guardTenantUsageByKey', () => {
       },
     });
 
-    await quotaLibrary.guardTenantUsageByKey('scopesPerResourceLimit', 'resource_1');
+    await quotaLibrary.guardTenantUsageByKey('scopesPerResourceLimit', { entityId: 'resource_1' });
 
     expect(getSelfComputedUsageByKey).toHaveBeenCalledWith(
       tenant.id,
@@ -340,7 +341,7 @@ describe('guardTenantUsageByKey', () => {
     });
 
     await expect(
-      quotaLibrary.guardTenantUsageByKey('scopesPerRoleLimit', 'role_1')
+      quotaLibrary.guardTenantUsageByKey('scopesPerRoleLimit', { entityId: 'role_1' })
     ).rejects.toMatchObject({
       code: 'subscription.limit_exceeded',
       status: 403,
@@ -399,6 +400,141 @@ describe('guardTenantUsageByKey', () => {
     // The key point: usage should be fetched only once, not twice
     expect(getSelfComputedUsageByKey).toHaveBeenCalledTimes(1);
   });
+
+  it('allows usage when usage + consumeUsageCount equals limit', async () => {
+    mockGetTenantSubscription.mockResolvedValueOnce({
+      ...mockSubscriptionData,
+      quota: {
+        ...mockSubscriptionData.quota,
+        applicationsLimit: 10,
+      },
+    });
+
+    const getSelfComputedUsageByKey = jest.fn().mockResolvedValue(7);
+
+    const { quotaLibrary } = createQuotaLibrary({
+      queriesOverride: {
+        tenantUsage: { getSelfComputedUsageByKey },
+      },
+    });
+
+    // 7 + 3 = 10, should be allowed (<=)
+    await expect(
+      quotaLibrary.guardTenantUsageByKey('applicationsLimit', { consumeUsageCount: 3 })
+    ).resolves.not.toThrow();
+  });
+
+  it('throws when usage + consumeUsageCount exceeds limit', async () => {
+    mockGetTenantSubscription.mockResolvedValueOnce({
+      ...mockSubscriptionData,
+      quota: {
+        ...mockSubscriptionData.quota,
+        applicationsLimit: 10,
+      },
+    });
+
+    const getSelfComputedUsageByKey = jest.fn().mockResolvedValue(7);
+
+    const { quotaLibrary } = createQuotaLibrary({
+      queriesOverride: {
+        tenantUsage: { getSelfComputedUsageByKey },
+      },
+    });
+
+    // 7 + 4 = 11 > 10, should throw
+    await expect(
+      quotaLibrary.guardTenantUsageByKey('applicationsLimit', { consumeUsageCount: 4 })
+    ).rejects.toMatchObject({
+      code: 'subscription.limit_exceeded',
+      status: 403,
+    });
+  });
+
+  it('allows batch operation when total usage equals limit for entity-based quota', async () => {
+    mockGetTenantSubscription.mockResolvedValueOnce({
+      ...mockSubscriptionData,
+      quota: {
+        ...mockSubscriptionData.quota,
+        scopesPerRoleLimit: 10,
+      },
+    });
+
+    const getSelfComputedUsageByKey = jest.fn().mockResolvedValue(5);
+
+    const { quotaLibrary } = createQuotaLibrary({
+      queriesOverride: {
+        tenantUsage: { getSelfComputedUsageByKey },
+      },
+    });
+
+    // Simulating adding 5 scopes at once: 5 + 5 = 10, should be allowed
+    await expect(
+      quotaLibrary.guardTenantUsageByKey('scopesPerRoleLimit', {
+        entityId: 'role_1',
+        consumeUsageCount: 5,
+      })
+    ).resolves.not.toThrow();
+  });
+
+  it('blocks batch operation that would exceed entity-based quota limit', async () => {
+    mockGetTenantSubscription.mockResolvedValueOnce({
+      ...mockSubscriptionData,
+      quota: {
+        ...mockSubscriptionData.quota,
+        scopesPerRoleLimit: 10,
+      },
+    });
+
+    const getSelfComputedUsageByKey = jest.fn().mockResolvedValue(5);
+
+    const { quotaLibrary } = createQuotaLibrary({
+      queriesOverride: {
+        tenantUsage: { getSelfComputedUsageByKey },
+      },
+    });
+
+    // Simulating adding 100 scopes at once: 5 + 100 = 105 > 10, should throw
+    await expect(
+      quotaLibrary.guardTenantUsageByKey('scopesPerRoleLimit', {
+        entityId: 'role_1',
+        consumeUsageCount: 100,
+      })
+    ).rejects.toMatchObject({
+      code: 'subscription.limit_exceeded',
+      status: 403,
+    });
+  });
+
+  it('respects system limit with batch consumption', async () => {
+    setEnvFlag('isDevFeaturesEnabled', true);
+    mockGetTenantSubscription.mockResolvedValueOnce({
+      ...mockSubscriptionData,
+      quota: {
+        ...mockSubscriptionData.quota,
+        applicationsLimit: null,
+      },
+      systemLimit: {
+        ...mockSubscriptionData.systemLimit,
+        applicationsLimit: 10,
+      },
+    });
+
+    const getSelfComputedUsageByKey = jest.fn().mockResolvedValue(8);
+
+    const { quotaLibrary } = createQuotaLibrary({
+      queriesOverride: {
+        tenantUsage: { getSelfComputedUsageByKey },
+      },
+    });
+
+    // 8 + 3 = 11 > 10, should throw system limit error
+    await expect(
+      quotaLibrary.guardTenantUsageByKey('applicationsLimit', { consumeUsageCount: 3 })
+    ).rejects.toMatchObject({
+      code: 'system_limit.limit_exceeded',
+      status: 403,
+    });
+  });
 });
 
 describe('reportSubscriptionUpdatesUsage', () => {
@@ -485,3 +621,4 @@ describe('reportSubscriptionUpdatesUsage', () => {
     expect(mockReportSubscriptionUpdates).not.toHaveBeenCalled();
   });
 });
+/* eslint-enable max-lines */
diff --git a/packages/core/src/libraries/quota.ts b/packages/core/src/libraries/quota.ts
@@ -36,10 +36,77 @@ const paidReservedPlans = new Set<string>([
   ReservedPlanId.Pro202509,
 ]);
 
+/**
+ * Options for quota guard operations.
+ */
+type GuardTenantUsageOptions = {
+  /**
+   * Entity ID for entity-based quota limits.
+   *
+   * @remarks
+   * Required for entity-based quotas like `scopesPerRoleLimit`, `scopesPerResourceLimit`, etc.
+   * Optional for tenant-level quotas like `applicationsLimit`, `resourcesLimit`, etc.
+   *
+   * @example
+   * ```ts
+   * // Entity-based quota: entityId is required
+   * await quota.guardTenantUsageByKey('scopesPerRoleLimit', { entityId: roleId });
+   *
+   * // Tenant-level quota: entityId not needed
+   * await quota.guardTenantUsageByKey('applicationsLimit');
+   * ```
+   */
+  entityId?: string;
+
+  /**
+   * The number of resources to consume in this operation.
+   *
+   * @remarks
+   * This is used to validate batch operations where multiple resources are added at once.
+   * For example, when assigning 10 scopes to a role in a single request, set this to 10.
+   *
+   * The guard will check: `currentUsage + consumeUsageCount <= limit`
+   *
+   * @default 1 - Assumes single resource consumption if not specified
+   *
+   * @example
+   * ```ts
+   * // Adding a single application (default behavior)
+   * await quota.guardTenantUsageByKey('applicationsLimit');
+   *
+   * // Adding 5 scopes to a role at once
+   * await quota.guardTenantUsageByKey('scopesPerRoleLimit', {
+   *   entityId: roleId,
+   *   consumeUsageCount: 5,
+   * });
+   * ```
+   */
+  consumeUsageCount?: number;
+};
+
+/**
+ * Function type for guarding tenant usage against quota limits.
+ * Provides type-safe overloads for tenant-level and entity-based quotas.
+ */
 type GuardTenantUsageByKeyFunction = {
-  (key: Exclude<UsageKey, EntityBasedUsageKey>, entityId?: string): Promise<void>;
-  (key: EntityBasedUsageKey, entityId: string): Promise<void>;
+  /**
+   * Guard tenant-level quota (e.g., applicationsLimit, resourcesLimit)
+   * @param key - Tenant-level quota key
+   * @param options - Optional guard options
+   */
+  (key: Exclude<UsageKey, EntityBasedUsageKey>, options?: GuardTenantUsageOptions): Promise<void>;
+
+  /**
+   * Guard entity-based quota (e.g., scopesPerRoleLimit, scopesPerResourceLimit)
+   * @param key - Entity-based quota key
+   * @param options - Guard options with required entityId
+   */
+  (
+    key: EntityBasedUsageKey,
+    options: GuardTenantUsageOptions & { entityId: string }
+  ): Promise<void>;
 };
+
 export class QuotaLibrary {
   constructor(
     public readonly tenantId: string,
@@ -49,7 +116,44 @@ export class QuotaLibrary {
     private readonly subscription: SubscriptionLibrary
   ) {}
 
-  guardTenantUsageByKey: GuardTenantUsageByKeyFunction = async (key, entityId) => {
+  /**
+   * Guards tenant usage against quota and system limits before performing an operation.
+   *
+   * @remarks
+   * This method checks if the current usage plus the resources to be consumed would exceed
+   * the configured limits. It validates against both system limits (hard caps) and
+   * subscription quota limits.
+   *
+   * The check formula is: `currentUsage + consumeUsageCount <= limit`
+   *
+   * If the limit would be exceeded, this method throws a RequestError. Otherwise, it
+   * returns silently, allowing the operation to proceed.
+   *
+   * @param key - The quota key to check (e.g., 'applicationsLimit', 'scopesPerRoleLimit')
+   * @param options - Optional configuration for the guard operation
+   * @param options.entityId - Entity ID for entity-based quotas (required for entity-based keys)
+   * @param options.consumeUsageCount - Number of resources to consume (default: 1)
+   *
+   * @throws {RequestError} With code 'system_limit.limit_exceeded' if system limit is exceeded
+   * @throws {RequestError} With code 'subscription.limit_exceeded' if quota limit is exceeded
+   *
+   * @example
+   * ```ts
+   * // Guard before adding a single application
+   * await quota.guardTenantUsageByKey('applicationsLimit');
+   *
+   * // Guard before adding multiple scopes to a role
+   * const scopeIds = ['scope1', 'scope2', 'scope3'];
+   * await quota.guardTenantUsageByKey('scopesPerRoleLimit', {
+   *   entityId: 'role_id',
+   *   consumeUsageCount: scopeIds.length,
+   * });
+   * ```
+   */
+  guardTenantUsageByKey: GuardTenantUsageByKeyFunction = async (
+    key,
+    { entityId, consumeUsageCount = 1 } = {}
+  ) => {
     const { isCloud } = EnvSet.values;
 
     // Cloud only feature, skip in non-cloud environments
@@ -67,7 +171,13 @@ export class QuotaLibrary {
 
     // Todo: @xiaoyijun: remove dev feature flag
     if (EnvSet.values.isDevFeaturesEnabled && isSystemUsageKey(key)) {
-      await this.assertSystemLimit({ key, entityId, subscriptionData, tenantUsageQuery });
+      await this.assertSystemLimit({
+        key,
+        entityId,
+        subscriptionData,
+        tenantUsageQuery,
+        consumeUsageCount,
+      });
     }
 
     if (isQuotaUsageKey(key)) {
@@ -76,6 +186,7 @@ export class QuotaLibrary {
         entityId,
         subscriptionData,
         tenantUsageQuery,
+        consumeUsageCount,
       });
     }
   };
@@ -118,11 +229,13 @@ export class QuotaLibrary {
     entityId,
     subscriptionData,
     tenantUsageQuery,
+    consumeUsageCount,
   }: {
     key: SystemUsageKey;
     entityId?: string;
     subscriptionData: Subscription;
     tenantUsageQuery: TenantUsageQuery;
+    consumeUsageCount: number;
   }) => {
     const { systemLimit } = subscriptionData;
 
@@ -139,7 +252,7 @@ export class QuotaLibrary {
     const usage = await tenantUsageQuery.get(key, entityId);
 
     assertThat(
-      usage < limit,
+      usage + consumeUsageCount <= limit,
       new RequestError(
         {
           code: 'system_limit.limit_exceeded',
@@ -155,11 +268,13 @@ export class QuotaLibrary {
     entityId,
     subscriptionData,
     tenantUsageQuery,
+    consumeUsageCount,
   }: {
     key: QuotaUsageKey;
     entityId?: string;
     subscriptionData: Subscription;
     tenantUsageQuery: TenantUsageQuery;
+    consumeUsageCount: number;
   }) => {
     const { planId, isEnterprisePlan, quota } = subscriptionData;
 
@@ -205,7 +320,7 @@ export class QuotaLibrary {
       const usage = await tenantUsageQuery.get(key, entityId);
 
       assertThat(
-        usage < limit,
+        usage + consumeUsageCount <= limit,
         new RequestError({
           code: 'subscription.limit_exceeded',
           status: 403,
diff --git a/packages/core/src/routes/resource.scope.ts b/packages/core/src/routes/resource.scope.ts
@@ -89,7 +89,7 @@ export default function resourceScopeRoutes<T extends ManagementApiRouter>(
         body,
       } = ctx.guard;
 
-      await quota.guardTenantUsageByKey('scopesPerResourceLimit', resourceId);
+      await quota.guardTenantUsageByKey('scopesPerResourceLimit', { entityId: resourceId });
 
       assertThat(!/\s/.test(body.name), 'scope.name_with_space');
 
diff --git a/packages/core/src/routes/role.scope.ts b/packages/core/src/routes/role.scope.ts
@@ -93,7 +93,10 @@ export default function roleScopeRoutes<T extends ManagementApiRouter>(
         body: { scopeIds },
       } = ctx.guard;
 
-      await quota.guardTenantUsageByKey('scopesPerRoleLimit', id);
+      await quota.guardTenantUsageByKey('scopesPerRoleLimit', {
+        entityId: id,
+        consumeUsageCount: scopeIds.length,
+      });
 
       await validateRoleScopeAssignment(scopeIds, id);
       await insertRolesScopes(
