fix: allow schema router to config middlewares status codes (#7926)

Author: xiaoyijun

packages/core/src/middleware/koa-quota-guard.ts

@@ -10,42 +10,30 @@ type Method = 'GET' | 'POST' | 'PUT' | 'PATCH' | 'DELETE' | 'COPY' | 'HEAD' | 'O
 type UsageGuardConfig = {
   key: Exclude<AllLimitKey, EntityBasedUsageKey>;
   quota: QuotaLibrary;
-  /** Guard usage only for the specified method types. Guard all if not provided. */
-  methods?: Method[];
 };
 
 export function koaQuotaGuard<StateT, ContextT, ResponseBodyT>({
   key,
   quota,
-  methods,
 }: UsageGuardConfig): MiddlewareType<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.guardTenantUsageByKey(key);
-    }
+  return async (_, next) => {
+    await quota.guardTenantUsageByKey(key);
     return next();
   };
 }
 
 type UsageReportConfig = {
   key: keyof SubscriptionUsage;
   quota: QuotaLibrary;
-  /** Report usage only for the specified method types. Report for all if not provided. */
-  methods?: Method[];
 };
 
 export function koaReportSubscriptionUpdates<StateT, ContextT, ResponseBodyT>({
   key,
   quota,
-  methods = ['POST', 'PUT', 'DELETE'],
 }: UsageReportConfig): MiddlewareType<StateT, ContextT, Nullable<ResponseBodyT>> {
-  return async (ctx, next) => {
+  return async (_, next) => {
     await next();
 
-    // eslint-disable-next-line no-restricted-syntax
-    if (methods.includes(ctx.method.toUpperCase() as Method)) {
-      void quota.reportSubscriptionUpdatesUsage(key);
-    }
+    void quota.reportSubscriptionUpdatesUsage(key);
   };
 }

packages/core/src/routes/hook.ts

@@ -172,7 +172,6 @@ export default function hookRoutes<T extends ManagementApiRouter>(
     koaReportSubscriptionUpdates({
       key: 'hooksLimit',
       quota,
-      methods: ['POST'],
     }),
     async (ctx, next) => {
       const { event, events, enabled, ...rest } = ctx.guard.body;
@@ -266,7 +265,6 @@ export default function hookRoutes<T extends ManagementApiRouter>(
     koaReportSubscriptionUpdates({
       key: 'hooksLimit',
       quota,
-      methods: ['DELETE'],
     }),
     async (ctx, next) => {
       const { id } = ctx.guard.params;

packages/core/src/routes/organization/index.ts

@@ -39,24 +39,19 @@ export default function organizationRoutes<T extends ManagementApiRouter>(
   const router = new SchemaRouter(Organizations, organizations, {
     middlewares: [
       {
-        middlewares: [
-          koaQuotaGuard({ key: 'organizationsLimit', quota, methods: ['POST', 'PUT'] }),
-        ],
-        scope: {
-          native: ['post', 'put'],
-        },
+        middleware: koaQuotaGuard({ key: 'organizationsLimit', quota }),
+        scope: 'native',
+        method: ['post', 'put'],
+        // Throw 403 when quota exceeded
+        status: [403],
       },
       {
-        middlewares: [
-          koaReportSubscriptionUpdates({
-            key: 'organizationsLimit',
-            quota,
-            methods: ['POST', 'PUT', 'DELETE'],
-          }),
-        ],
-        scope: {
-          native: ['post', 'put', 'delete'],
-        },
+        middleware: koaReportSubscriptionUpdates({
+          key: 'organizationsLimit',
+          quota,
+        }),
+        scope: 'native',
+        method: ['post', 'put', 'delete'],
       },
     ],
     errorHandler,

packages/core/src/routes/resource.ts

@@ -90,7 +90,6 @@ export default function resourceRoutes<T extends ManagementApiRouter>(
     koaReportSubscriptionUpdates({
       key: 'resourcesLimit',
       quota,
-      methods: ['POST'],
     }),
     async (ctx, next) => {
       const { body } = ctx.guard;
@@ -193,7 +192,6 @@ export default function resourceRoutes<T extends ManagementApiRouter>(
     koaReportSubscriptionUpdates({
       key: 'resourcesLimit',
       quota,
-      methods: ['DELETE'],
     }),
     async (ctx, next) => {
       const { id } = ctx.guard.params;

packages/core/src/routes/sso-connector/index.ts

@@ -83,7 +83,6 @@ export default function singleSignOnConnectorsRoutes<T extends ManagementApiRout
     koaReportSubscriptionUpdates({
       quota,
       key: 'enterpriseSsoLimit',
-      methods: ['POST'],
     }),
     async (ctx, next) => {
       const { body } = ctx.guard;
@@ -233,7 +232,6 @@ export default function singleSignOnConnectorsRoutes<T extends ManagementApiRout
     koaReportSubscriptionUpdates({
       quota,
       key: 'enterpriseSsoLimit',
-      methods: ['DELETE'],
     }),
     async (ctx, next) => {
       const { id } = ctx.guard.params;

packages/core/src/utils/SchemaRouter.test.ts

@@ -1,4 +1,5 @@
 import { type GeneratedSchema } from '@logto/schemas';
+import type { Middleware } from 'koa';
 import { z } from 'zod';
 
 import RequestError from '#src/errors/RequestError/index.js';
@@ -132,6 +133,120 @@ describe('SchemaRouter', () => {
     });
   });
 
+  describe('middlewares', () => {
+    it('should allow status codes declared in middleware guard config with scope and method', async () => {
+      // eslint-disable-next-line unicorn/consistent-function-scoping
+      const throwingMiddleware: Middleware = async () => {
+        throw new RequestError({ code: 'entity.not_found', status: 403 });
+      };
+
+      const schemaRouterWithMiddleware = new SchemaRouter(schema, queries, {
+        middlewares: [
+          {
+            middleware: throwingMiddleware,
+            scope: 'native',
+            method: ['get'],
+            status: [403],
+          },
+        ],
+      });
+      const requestWithMiddleware = createRequester({
+        authedRoutes: (router) => router.use(schemaRouterWithMiddleware.routes()),
+      });
+
+      const response = await requestWithMiddleware.get(baseRoute);
+
+      expect(response.status).toEqual(403);
+    });
+
+    it('should apply middleware to all methods when method is not specified', async () => {
+      const callTracker = jest.fn();
+
+      const trackingMiddleware: Middleware = async (ctx, next) => {
+        callTracker(ctx.method);
+        return next();
+      };
+
+      const schemaRouterWithMiddleware = new SchemaRouter(schema, queries, {
+        middlewares: [
+          {
+            middleware: trackingMiddleware,
+            scope: 'native',
+          },
+        ],
+      });
+      const requestWithMiddleware = createRequester({
+        authedRoutes: (router) => router.use(schemaRouterWithMiddleware.routes()),
+      });
+
+      await requestWithMiddleware.get(baseRoute);
+      await requestWithMiddleware.post(baseRoute).send({});
+      await requestWithMiddleware.patch(`${baseRoute}/test`).send({ name: 'test' });
+
+      expect(callTracker).toHaveBeenCalledTimes(3);
+      expect(callTracker).toHaveBeenCalledWith('GET');
+      expect(callTracker).toHaveBeenCalledWith('POST');
+      expect(callTracker).toHaveBeenCalledWith('PATCH');
+    });
+
+    it('should only apply middleware to specified methods', async () => {
+      const callTracker = jest.fn();
+
+      const trackingMiddleware: Middleware = async (ctx, next) => {
+        callTracker(ctx.method);
+        return next();
+      };
+
+      const schemaRouterWithMiddleware = new SchemaRouter(schema, queries, {
+        middlewares: [
+          {
+            middleware: trackingMiddleware,
+            scope: 'native',
+            method: ['get', 'post'],
+          },
+        ],
+      });
+      const requestWithMiddleware = createRequester({
+        authedRoutes: (router) => router.use(schemaRouterWithMiddleware.routes()),
+      });
+
+      await requestWithMiddleware.get(baseRoute);
+      await requestWithMiddleware.post(baseRoute).send({});
+      await requestWithMiddleware.patch(`${baseRoute}/test`).send({ name: 'test' });
+
+      expect(callTracker).toHaveBeenCalledTimes(2);
+      expect(callTracker).toHaveBeenCalledWith('GET');
+      expect(callTracker).toHaveBeenCalledWith('POST');
+      expect(callTracker).not.toHaveBeenCalledWith('PATCH');
+    });
+
+    it('should apply middleware without scope to all routes', async () => {
+      const callTracker = jest.fn();
+
+      const trackingMiddleware: Middleware = async (ctx, next) => {
+        callTracker();
+        return next();
+      };
+
+      const schemaRouterWithMiddleware = new SchemaRouter(schema, queries, {
+        middlewares: [
+          {
+            middleware: trackingMiddleware,
+          },
+        ],
+      });
+      const requestWithMiddleware = createRequester({
+        authedRoutes: (router) => router.use(schemaRouterWithMiddleware.routes()),
+      });
+
+      await requestWithMiddleware.get(baseRoute);
+      await requestWithMiddleware.post(baseRoute).send({});
+
+      // Global middleware should be called for all routes
+      expect(callTracker).toHaveBeenCalledTimes(2);
+    });
+  });
+
   describe('disable routes', () => {
     it('should be able to disable routes', async () => {
       const disabledSchemaRouter = new SchemaRouter(schema, queries, {