refactor: update code according to CR

logto-io · Jul 22, 2024 · 9bc2bc0 · 9bc2bc0
1 parent d8a40d1
commit 9bc2bc0
Show file tree

Hide file tree

Showing 14 changed files with 70 additions and 32 deletions.
diff --git a/packages/core/src/libraries/quota.ts b/packages/core/src/libraries/quota.ts
@@ -143,8 +143,7 @@ export const createQuotaLibrary = (
     }
   };
 
-  // `SubscriptionQuota` and `SubscriptionUsage` are sharing keys.
-  const newGuardKey = async (key: keyof SubscriptionQuota) => {
+  const guardTenantUsageByKey = async (key: keyof SubscriptionQuota) => {
     const { isCloud, isIntegrationTest } = EnvSet.values;
 
     // Cloud only feature, skip in non-cloud environments
@@ -160,6 +159,8 @@ export const createQuotaLibrary = (
     const { quota: fullQuota, usage: fullUsage } = await getTenantSubscriptionQuotaAndUsage(
       cloudConnection
     );
+
+    // Type `SubscriptionQuota` and type `SubscriptionUsage` are sharing keys, this design helps us to compare the usage with the quota limit in a easier way.
     const { [key]: limit } = fullQuota;
     const { [key]: usage } = fullUsage;
 
@@ -178,7 +179,10 @@ export const createQuotaLibrary = (
           },
         })
       );
-    } else if (typeof limit === 'number') {
+      return;
+    }
+
+    if (typeof limit === 'number') {
       // See the definition of `SubscriptionQuota` and `SubscriptionUsage` in `types.ts`, this should never happen.
       assertThat(
         typeof usage === 'number',
@@ -197,12 +201,14 @@ export const createQuotaLibrary = (
           },
         })
       );
-    } else {
-      throw new TypeError('Unsupported subscription quota type');
+
+      return;
     }
+
+    throw new TypeError('Unsupported subscription quota type');
   };
 
-  const scopesGuardKey = async (entityName: 'resources' | 'roles', entityId: string) => {
+  const guardEntityScopesUsage = async (entityName: 'resources' | 'roles', entityId: string) => {
     const { isCloud, isIntegrationTest } = EnvSet.values;
 
     // Cloud only feature, skip in non-cloud environments
@@ -215,10 +221,15 @@ export const createQuotaLibrary = (
       return;
     }
 
-    const {
-      quota: { scopesPerResourceLimit, scopesPerRoleLimit },
-    } = await getTenantSubscriptionQuotaAndUsage(cloudConnection);
-    const scopeUsages = await getTenantSubscriptionScopeUsage(cloudConnection, entityName);
+    const [
+      {
+        quota: { scopesPerResourceLimit, scopesPerRoleLimit },
+      },
+      scopeUsages,
+    ] = await Promise.all([
+      getTenantSubscriptionQuotaAndUsage(cloudConnection),
+      getTenantSubscriptionScopeUsage(cloudConnection, entityName),
+    ]);
     const usage = scopeUsages[entityId] ?? 0;
 
     if (entityName === 'resources') {
@@ -251,5 +262,5 @@ export const createQuotaLibrary = (
     );
   };
 
-  return { guardKey, newGuardKey, scopesGuardKey };
+  return { guardKey, guardTenantUsageByKey, guardEntityScopesUsage };
 };
diff --git a/packages/core/src/libraries/sign-in-experience/index.ts b/packages/core/src/libraries/sign-in-experience/index.ts
@@ -7,7 +7,7 @@ import type {
   SignInExperience,
   SsoConnectorMetadata,
 } from '@logto/schemas';
-import { ConnectorType } from '@logto/schemas';
+import { ConnectorType, ReservedPlanId } from '@logto/schemas';
 import { deduplicate, pick, trySafe } from '@silverhand/essentials';
 import deepmerge from 'deepmerge';
 
@@ -19,7 +19,7 @@ import type { SsoConnectorLibrary } from '#src/libraries/sso-connector.js';
 import { ssoConnectorFactories } from '#src/sso/index.js';
 import type Queries from '#src/tenants/Queries.js';
 import assertThat from '#src/utils/assert-that.js';
-import { getTenantSubscriptionPlan } from '#src/utils/subscription/index.js';
+import { getTenantSubscription, getTenantSubscriptionPlan } from '#src/utils/subscription/index.js';
 import { isKeyOfI18nPhrases } from '#src/utils/translation.js';
 
 import { type CloudConnectionLibrary } from '../cloud-connection.js';
@@ -113,8 +113,12 @@ export const createSignInExperienceLibrary = (
       return false;
     }
 
-    const plan = await getTenantSubscriptionPlan(cloudConnection);
+    if (EnvSet.values.isDevFeaturesEnabled) {
+      const subscription = await getTenantSubscription(cloudConnection);
+      return subscription.planId === ReservedPlanId.Development;
+    }
 
+    const plan = await getTenantSubscriptionPlan(cloudConnection);
     return plan.id === developmentTenantPlanId;
   }, ['is-development-tenant']);
 

diff --git a/packages/core/src/middleware/koa-quota-guard.ts b/packages/core/src/middleware/koa-quota-guard.ts
@@ -40,7 +40,7 @@ export function newKoaQuotaGuard<StateT, ContextT, ResponseBodyT>({
   return async (ctx, next) => {
     // eslint-disable-next-line no-restricted-syntax
     if (!methods || methods.includes(ctx.method.toUpperCase() as Method)) {
-      await quota.newGuardKey(key);
+      await quota.guardTenantUsageByKey(key);
     }
     return next();
   };

diff --git a/packages/core/src/routes/applications/application.ts b/packages/core/src/routes/applications/application.ts
@@ -171,19 +171,19 @@ export default function applicationRoutes<T extends ManagementApiRouter>(
       // When creating a m2m app, should check both m2m limit and application limit.
       if (rest.type === ApplicationType.MachineToMachine) {
         await (EnvSet.values.isDevFeaturesEnabled
-          ? quota.newGuardKey('machineToMachineLimit')
+          ? quota.guardTenantUsageByKey('machineToMachineLimit')
           : quota.guardKey('machineToMachineLimit'));
       }
 
       // Guard third party application limit
       if (rest.isThirdParty) {
         await (EnvSet.values.isDevFeaturesEnabled
-          ? quota.newGuardKey('thirdPartyApplicationsLimit')
+          ? quota.guardTenantUsageByKey('thirdPartyApplicationsLimit')
           : quota.guardKey('thirdPartyApplicationsLimit'));
       }
 
       await (EnvSet.values.isDevFeaturesEnabled
-        ? quota.newGuardKey('applicationsLimit')
+        ? quota.guardTenantUsageByKey('applicationsLimit')
         : quota.guardKey('applicationsLimit'));
 
       assertThat(

diff --git a/packages/core/src/routes/connector/index.ts b/packages/core/src/routes/connector/index.ts
@@ -28,7 +28,7 @@ const guardConnectorsQuota = async (
 ) => {
   if (factory.type === ConnectorType.Social) {
     await (EnvSet.values.isDevFeaturesEnabled
-      ? quota.newGuardKey('socialConnectorsLimit')
+      ? quota.guardTenantUsageByKey('socialConnectorsLimit')
       : quota.guardKey('socialConnectorsLimit'));
   }
 };

diff --git a/packages/core/src/routes/resource.scope.ts b/packages/core/src/routes/resource.scope.ts
@@ -91,7 +91,7 @@ export default function resourceScopeRoutes<T extends ManagementApiRouter>(
       } = ctx.guard;
 
       await (EnvSet.values.isDevFeaturesEnabled
-        ? quota.scopesGuardKey('resources', resourceId)
+        ? quota.guardEntityScopesUsage('resources', resourceId)
         : quota.guardKey('scopesPerResourceLimit', 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
@@ -95,7 +95,7 @@ export default function roleScopeRoutes<T extends ManagementApiRouter>(
       } = ctx.guard;
 
       await (EnvSet.values.isDevFeaturesEnabled
-        ? quota.scopesGuardKey('roles', id)
+        ? quota.guardEntityScopesUsage('roles', id)
         : quota.guardKey('scopesPerRoleLimit', id));
 
       await validateRoleScopeAssignment(scopeIds, id);

diff --git a/packages/core/src/routes/role.ts b/packages/core/src/routes/role.ts
@@ -152,7 +152,7 @@ export default function roleRoutes<T extends ManagementApiRouter>(
       // We have optional `type` when creating a new role, if `type` is not provided, use `User` as default.
       // `machineToMachineRolesLimit` is the limit of machine to machine roles, and is independent to `rolesLimit`.
       await (EnvSet.values.isDevFeaturesEnabled
-        ? quota.newGuardKey(
+        ? quota.guardTenantUsageByKey(
             roleBody.type === RoleType.MachineToMachine
               ? 'machineToMachineRolesLimit'
               : // In new pricing model, we rename `rolesLimit` to `userRolesLimit`, which is easier to be distinguished from `machineToMachineRolesLimit`.

diff --git a/packages/core/src/routes/sign-in-experience/custom-ui-assets/index.ts b/packages/core/src/routes/sign-in-experience/custom-ui-assets/index.ts
@@ -8,7 +8,7 @@ import { object, z } from 'zod';
 import { EnvSet } from '#src/env-set/index.js';
 import RequestError from '#src/errors/RequestError/index.js';
 import koaGuard from '#src/middleware/koa-guard.js';
-import koaQuotaGuard from '#src/middleware/koa-quota-guard.js';
+import koaQuotaGuard, { newKoaQuotaGuard } from '#src/middleware/koa-quota-guard.js';
 import SystemContext from '#src/tenants/SystemContext.js';
 import assertThat from '#src/utils/assert-that.js';
 import { getConsoleLogFromContext } from '#src/utils/console.js';
@@ -35,7 +35,11 @@ export default function customUiAssetsRoutes<T extends ManagementApiRouter>(
 
   router.post(
     '/sign-in-exp/default/custom-ui-assets',
-    koaQuotaGuard({ key: 'bringYourUiEnabled', quota }),
+    // Manually add this to avoid the case that the dev feature guard is removed but the quota guard is not being updated accordingly.
+    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+    EnvSet.values.isDevFeaturesEnabled
+      ? newKoaQuotaGuard({ key: 'bringYourUiEnabled', quota })
+      : koaQuotaGuard({ key: 'bringYourUiEnabled', quota }),
     koaGuard({
       files: object({
         file: uploadFileGuard.array().min(1).max(1),

diff --git a/packages/core/src/routes/sign-in-experience/index.ts b/packages/core/src/routes/sign-in-experience/index.ts
@@ -19,7 +19,7 @@ export default function signInExperiencesRoutes<T extends ManagementApiRouter>(
   const { deleteConnectorById } = queries.connectors;
   const {
     signInExperiences: { validateLanguageInfo },
-    quota: { guardKey, newGuardKey },
+    quota: { guardKey, guardTenantUsageByKey },
   } = libraries;
   const { getLogtoConnectors } = connectors;
 
@@ -92,7 +92,7 @@ export default function signInExperiencesRoutes<T extends ManagementApiRouter>(
       if (mfa) {
         if (mfa.factors.length > 0) {
           await (EnvSet.values.isDevFeaturesEnabled
-            ? newGuardKey('mfaEnabled')
+            ? guardTenantUsageByKey('mfaEnabled')
             : guardKey('mfaEnabled'));
         }
         validateMfa(mfa);

diff --git a/packages/core/src/routes/subject-token.ts b/packages/core/src/routes/subject-token.ts
@@ -4,8 +4,9 @@ import { addSeconds } from 'date-fns';
 import { object, string } from 'zod';
 
 import { subjectTokenExpiresIn, subjectTokenPrefix } from '#src/constants/index.js';
+import { EnvSet } from '#src/env-set/index.js';
 import koaGuard from '#src/middleware/koa-guard.js';
-import { newKoaQuotaGuard } from '#src/middleware/koa-quota-guard.js';
+import koaQuotaGuard, { newKoaQuotaGuard } from '#src/middleware/koa-quota-guard.js';
 
 import { type RouterInitArgs, type ManagementApiRouter } from './types.js';
 
@@ -25,7 +26,9 @@ export default function subjectTokenRoutes<T extends ManagementApiRouter>(
 
   router.post(
     '/subject-tokens',
-    newKoaQuotaGuard({ key: 'subjectTokenEnabled', quota }),
+    EnvSet.values.isDevFeaturesEnabled
+      ? newKoaQuotaGuard({ key: 'subjectTokenEnabled', quota })
+      : koaQuotaGuard({ key: 'subjectTokenEnabled', quota }),
     koaGuard({
       body: object({
         userId: string(),

diff --git a/packages/core/src/test-utils/quota.ts b/packages/core/src/test-utils/quota.ts
@@ -5,7 +5,7 @@ const { jest } = import.meta;
 export const createMockQuotaLibrary = (): QuotaLibrary => {
   return {
     guardKey: jest.fn(),
-    newGuardKey: jest.fn(),
-    scopesGuardKey: jest.fn(),
+    guardTenantUsageByKey: jest.fn(),
+    guardEntityScopesUsage: jest.fn(),
   };
 };
diff --git a/packages/core/src/utils/subscription/index.ts b/packages/core/src/utils/subscription/index.ts
@@ -2,14 +2,28 @@ import { type CloudConnectionLibrary } from '#src/libraries/cloud-connection.js'
 
 import assertThat from '../assert-that.js';
 
-import { type SubscriptionQuota, type SubscriptionUsage, type SubscriptionPlan } from './types.js';
+import {
+  type SubscriptionQuota,
+  type SubscriptionUsage,
+  type SubscriptionPlan,
+  type Subscription,
+} from './types.js';
+
+export const getTenantSubscription = async (
+  cloudConnection: CloudConnectionLibrary
+): Promise<Subscription> => {
+  const client = await cloudConnection.getClient();
+  const subscription = await client.get('/api/tenants/my/subscription');
+
+  return subscription;
+};
 
 export const getTenantSubscriptionPlan = async (
   cloudConnection: CloudConnectionLibrary
 ): Promise<SubscriptionPlan> => {
   const client = await cloudConnection.getClient();
   const [subscription, plans] = await Promise.all([
-    client.get('/api/tenants/my/subscription'),
+    getTenantSubscription(cloudConnection),
     client.get('/api/subscription-plans'),
   ]);
   const plan = plans.find(({ id }) => id === subscription.planId);

diff --git a/packages/core/src/utils/subscription/types.ts b/packages/core/src/utils/subscription/types.ts
@@ -9,6 +9,8 @@ type RouteResponseType<T extends { search?: unknown; body?: unknown; response?:
 
 export type SubscriptionPlan = RouteResponseType<GetRoutes['/api/subscription-plans']>[number];
 
+export type Subscription = RouteResponseType<GetRoutes['/api/tenants/:tenantId/subscription']>;
+
 // Since `standardConnectorsLimit` will be removed in the upcoming pricing V2, no need to guard it.
 // `tokenLimit` is not guarded in backend.
 export type FeatureQuota = Omit<