From 191298fff489616dfea2b0175bc275ad6669ef57 Mon Sep 17 00:00:00 2001
From: Charles Zhao <charleszhao@silverhand.io>
Date: Tue, 16 Jan 2024 22:46:12 +0800
Subject: [PATCH] fix(core): fix quota guard for schema routers (#5242)

* fix(core): fix quota guard for schema routers

* chore: add comments
---
 packages/core/src/routes/organization/index.ts  | 3 +--
 packages/core/src/routes/organization/roles.ts  | 3 +--
 packages/core/src/routes/organization/scopes.ts | 3 +--
 packages/core/src/utils/SchemaRouter.ts         | 7 +++++++
 4 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/packages/core/src/routes/organization/index.ts b/packages/core/src/routes/organization/index.ts
index 9efc735ee..4b2dc024a 100644
--- a/packages/core/src/routes/organization/index.ts
+++ b/packages/core/src/routes/organization/index.ts
@@ -34,6 +34,7 @@ export default function organizationRoutes<T extends AuthedRouter>(...args: Rout
   ] = args;
 
   const router = new SchemaRouter(Organizations, organizations, {
+    middlewares: [koaQuotaGuard({ key: 'organizationsEnabled', quota, methods: ['POST', 'PUT'] })],
     errorHandler,
     searchFields: ['name'],
     disabled: { get: true },
@@ -243,8 +244,6 @@ export default function organizationRoutes<T extends AuthedRouter>(...args: Rout
     organizationInvitationRoutes(...args);
   }
 
-  router.use(koaQuotaGuard({ key: 'organizationsEnabled', quota, methods: ['POST', 'PUT'] }));
-
   // Add routes to the router
   originalRouter.use(router.routes());
 }
diff --git a/packages/core/src/routes/organization/roles.ts b/packages/core/src/routes/organization/roles.ts
index 502e56f91..a4e7b08ad 100644
--- a/packages/core/src/routes/organization/roles.ts
+++ b/packages/core/src/routes/organization/roles.ts
@@ -30,6 +30,7 @@ export default function organizationRoleRoutes<T extends AuthedRouter>(
   ]: RouterInitArgs<T>
 ) {
   const router = new SchemaRouter(OrganizationRoles, roles, {
+    middlewares: [koaQuotaGuard({ key: 'organizationsEnabled', quota, methods: ['POST', 'PUT'] })],
     disabled: { get: true, post: true },
     errorHandler,
     searchFields: ['name'],
@@ -89,7 +90,5 @@ export default function organizationRoleRoutes<T extends AuthedRouter>(
 
   router.addRelationRoutes(rolesScopes, 'scopes');
 
-  router.use(koaQuotaGuard({ key: 'organizationsEnabled', quota, methods: ['POST', 'PUT'] }));
-
   originalRouter.use(router.routes());
 }
diff --git a/packages/core/src/routes/organization/scopes.ts b/packages/core/src/routes/organization/scopes.ts
index be58caf35..3cab9b2fd 100644
--- a/packages/core/src/routes/organization/scopes.ts
+++ b/packages/core/src/routes/organization/scopes.ts
@@ -19,11 +19,10 @@ export default function organizationScopeRoutes<T extends AuthedRouter>(
   ]: RouterInitArgs<T>
 ) {
   const router = new SchemaRouter(OrganizationScopes, scopes, {
+    middlewares: [koaQuotaGuard({ key: 'organizationsEnabled', quota, methods: ['POST', 'PUT'] })],
     errorHandler,
     searchFields: ['name'],
   });
 
-  router.use(koaQuotaGuard({ key: 'organizationsEnabled', quota, methods: ['POST', 'PUT'] }));
-
   originalRouter.use(router.routes());
 }
diff --git a/packages/core/src/utils/SchemaRouter.ts b/packages/core/src/utils/SchemaRouter.ts
index ad3aaa095..f2a284dc5 100644
--- a/packages/core/src/utils/SchemaRouter.ts
+++ b/packages/core/src/utils/SchemaRouter.ts
@@ -3,6 +3,7 @@ import { generateStandardId } from '@logto/shared';
 import { type DeepPartial } from '@silverhand/essentials';
 import camelcase from 'camelcase';
 import deepmerge from 'deepmerge';
+import { type MiddlewareType } from 'koa';
 import Router, { type IRouterParamContext } from 'koa-router';
 import { z } from 'zod';
 
@@ -50,6 +51,8 @@ type SchemaRouterConfig<Key extends string> = {
     /** Disable `DELETE /:id` route. */
     deleteById: boolean;
   };
+  /** Middlewares that are used before creating API routes */
+  middlewares?: MiddlewareType[];
   /** A custom error handler for the router before throwing the error. */
   errorHandler?: (error: unknown) => void;
   /** The fields that can be searched for the `GET /` route. */
@@ -113,6 +116,10 @@ export default class SchemaRouter<
       config
     );
 
+    if (this.config.middlewares?.length) {
+      this.use(...this.config.middlewares);
+    }
+
     if (this.config.errorHandler) {
       this.use(async (_, next) => {
         try {