feat(gateway): extend app selector as middleware (#2761)

* feat: extend app selector as middleware

* fix: test

* chore: prevent duplicate middleware additions

packages/core/server/src/__tests__/gateway.test.ts

@@ -17,6 +17,43 @@ describe('gateway', () => {
     await AppSupervisor.getInstance().destroy();
   });
 
+  describe('app selector', () => {
+    it('should get app as default main app', async () => {
+      expect(
+        await gateway.getRequestHandleAppName({
+          url: '/test',
+          headers: {},
+        }),
+      ).toBe('main');
+    });
+
+    it('should add middleware into app selector', async () => {
+      gateway.addAppSelectorMiddleware(async (ctx, next) => {
+        ctx.resolvedAppName = 'test';
+        await next();
+      });
+
+      expect(
+        await gateway.getRequestHandleAppName({
+          url: '/test',
+          headers: {},
+        }),
+      ).toEqual('test');
+    });
+
+    it('should add same middleware into app selector once', async () => {
+      const fn = async (ctx, next) => {
+        ctx.resolvedAppName = 'test';
+        await next();
+      };
+
+      gateway.addAppSelectorMiddleware(fn);
+      gateway.addAppSelectorMiddleware(fn);
+
+      expect(gateway.getAppSelectorMiddlewares().nodes.length).toBe(2);
+    });
+  });
+
   describe('http api', () => {
     it('should return error when app not found', async () => {
       const res = await supertest.agent(gateway.getCallback()).get('/api/app:getInfo');

packages/core/server/src/gateway/index.ts

@@ -1,4 +1,4 @@
-import { uid } from '@nocobase/utils';
+import { Toposort, ToposortOptions, uid } from '@nocobase/utils';
 import { createStoragePluginsSymlink } from '@nocobase/utils/plugin-symlink';
 import { Command } from 'commander';
 import compression from 'compression';
@@ -18,6 +18,7 @@ import { applyErrorWithArgs, getErrorWithCode } from './errors';
 import { IPCSocketClient } from './ipc-socket-client';
 import { IPCSocketServer } from './ipc-socket-server';
 import { WSServer } from './ws-server';
+import compose from 'koa-compose';
 
 const compress = promisify(compression());
 
@@ -27,6 +28,7 @@ export interface IncomingRequest {
 }
 
 export type AppSelector = (req: IncomingRequest) => string | Promise<string>;
+export type AppSelectorMiddleware = (ctx: AppSelectorMiddlewareContext, next: () => Promise<void>) => void;
 
 interface StartHttpServerOptions {
   port: number;
@@ -38,12 +40,18 @@ interface RunOptions {
   mainAppOptions: ApplicationOptions;
 }
 
+export interface AppSelectorMiddlewareContext {
+  req: IncomingRequest;
+  resolvedAppName: string | null;
+}
+
 export class Gateway extends EventEmitter {
   private static instance: Gateway;
   /**
    * use main app as default app to handle request
    */
-  appSelector: AppSelector;
+  selectorMiddlewares: Toposort<AppSelectorMiddleware> = new Toposort<AppSelectorMiddleware>();
+
   public server: http.Server | null = null;
   public ipcSocketServer: IPCSocketServer | null = null;
   private port: number = process.env.APP_PORT ? parseInt(process.env.APP_PORT) : null;
@@ -73,18 +81,28 @@ export class Gateway extends EventEmitter {
   }
 
   public reset() {
-    this.setAppSelector(async (req) => {
-      const appName = qs.parse(parse(req.url).query)?.__appName;
-      if (appName) {
-        return appName;
-      }
+    this.selectorMiddlewares = new Toposort<AppSelectorMiddleware>();
 
-      if (req.headers['x-app']) {
-        return req.headers['x-app'];
-      }
+    this.addAppSelectorMiddleware(
+      async (ctx: AppSelectorMiddlewareContext, next) => {
+        const { req } = ctx;
+        const appName = qs.parse(parse(req.url).query)?.__appName as string | null;
 
-      return null;
-    });
+        if (appName) {
+          ctx.resolvedAppName = appName;
+        }
+
+        if (req.headers['x-app']) {
+          ctx.resolvedAppName = req.headers['x-app'];
+        }
+
+        await next();
+      },
+      {
+        tag: 'core',
+        group: 'core',
+      },
+    );
 
     if (this.server) {
       this.server.close();
@@ -97,8 +115,12 @@ export class Gateway extends EventEmitter {
     }
   }
 
-  setAppSelector(selector: AppSelector) {
-    this.appSelector = selector;
+  addAppSelectorMiddleware(middleware: AppSelectorMiddleware, options?: ToposortOptions) {
+    if (this.selectorMiddlewares.nodes.some((existingFunc) => existingFunc.toString() === middleware.toString())) {
+      return;
+    }
+
+    this.selectorMiddlewares.add(middleware, options);
     this.emit('appSelectorChanged');
   }
 
@@ -194,8 +216,25 @@ export class Gateway extends EventEmitter {
     app.callback()(req, res);
   }
 
+  getAppSelectorMiddlewares() {
+    return this.selectorMiddlewares;
+  }
+
   async getRequestHandleAppName(req: IncomingRequest) {
-    return (await this.appSelector(req)) || 'main';
+    const appSelectorMiddlewares = this.selectorMiddlewares.sort();
+
+    const ctx: AppSelectorMiddlewareContext = {
+      req,
+      resolvedAppName: null,
+    };
+
+    await compose(appSelectorMiddlewares)(ctx);
+
+    if (!ctx.resolvedAppName) {
+      ctx.resolvedAppName = 'main';
+    }
+
+    return ctx.resolvedAppName;
   }
 
   getCallback() {

packages/plugins/@nocobase/plugin-multi-app-manager/src/server/__tests__/multiple-apps.test.ts

@@ -190,7 +190,7 @@ describe('multiple apps', () => {
 
     expect(AppSupervisor.getInstance().hasApp(name)).toBeFalsy();
 
-    Gateway.getInstance().appSelector = () => name;
+    Gateway.getInstance().addAppSelectorMiddleware((ctx) => (ctx.resolvedAppName = name));
 
     await AppSupervisor.getInstance().getApp(name);
 

packages/plugins/@nocobase/plugin-multi-app-manager/src/server/server.ts

@@ -224,21 +224,16 @@ export class PluginMultiAppManager extends Plugin {
 
     AppSupervisor.getInstance().setAppBootstrapper(LazyLoadApplication);
 
-    Gateway.getInstance().setAppSelector(async (req) => {
-      const appName = qs.parse(parse(req.url).query)?.__appName;
-      if (appName) {
-        return appName;
-      }
+    Gateway.getInstance().addAppSelectorMiddleware(async (ctx, next) => {
+      const { req } = ctx;
 
-      if (req.headers['x-app']) {
-        return req.headers['x-app'];
-      }
-
-      if (req.headers['x-hostname']) {
+      if (!ctx.resolvedAppName && req.headers['x-hostname']) {
         const repository = this.db.getRepository('applications');
         if (!repository) {
-          return null;
+          await next();
+          return;
         }
+
         const appInstance = await repository.findOne({
           filter: {
             cname: req.headers['x-hostname'],
@@ -246,11 +241,11 @@ export class PluginMultiAppManager extends Plugin {
         });
 
         if (appInstance) {
-          return appInstance.name;
+          ctx.resolvedAppName = appInstance.name;
         }
       }
 
-      return null;
+      await next();
     });
 
     this.app.on('afterStart', async (app) => {
@@ -258,8 +253,9 @@ export class PluginMultiAppManager extends Plugin {
       const appSupervisor = AppSupervisor.getInstance();
 
       this.app.setMaintainingMessage('starting sub applications...');
+
       if (appSupervisor.runningMode == 'single') {
-        Gateway.getInstance().setAppSelector(() => appSupervisor.singleAppName);
+        Gateway.getInstance().addAppSelectorMiddleware((ctx) => (ctx.resolvedAppName = appSupervisor.singleAppName));
 
         // If the sub application is running in single mode, register the application automatically
         try {