feat: add route info into unstable_middleware context (#6605)

zllkjc · caohuilin · web-flow · commit 3d2bf551ee95 · 2024-12-05T01:57:49.000Z
Co-authored-by: Belinda Cao &lt;caohuilin@bytedance.com&gt;
diff --git a/.changeset/quick-cows-reply.md b/.changeset/quick-cows-reply.md
@@ -0,0 +1,8 @@
+---
+'@modern-js/plugin-server': patch
+'@modern-js/types': patch
+'@modern-js/server-core': patch
+---
+
+feat: add route info into unstable_middleware context
+feat: 添加路由信息到 unstable_middleware 上下文中
diff --git a/packages/server/core/src/plugins/customServer/index.ts b/packages/server/core/src/plugins/customServer/index.ts
@@ -218,7 +218,7 @@ export function getServerMidFromUnstableMid(
 ): Array<Middleware<ServerNodeEnv & ServerEnv>> {
   return serverMiddleware.map(middleware => {
     return async (c, next) => {
-      const context = await createMiddlewareContextFromHono(c as Context);
+      const context = await createMiddlewareContextFromHono(c);
 
       return middleware(context, next);
     };
@@ -230,7 +230,7 @@ function isRedirect(headers: Headers, code?: number) {
 }
 
 async function createMiddlewareContextFromHono(
-  c: Context<ServerNodeEnv>,
+  c: Context<ServerNodeEnv & ServerEnv>,
 ): Promise<UnstableMiddlewareContext> {
   const loaderContext = getLoaderCtx(c as Context);
 
@@ -276,6 +276,10 @@ async function createMiddlewareContextFromHono(
       c.res = newRes;
     },
 
+    get route() {
+      return c.get('route');
+    },
+
     get(key) {
       return loaderContext.get(key as string);
     },
@@ -295,3 +299,15 @@ async function createMiddlewareContextFromHono(
     redirect: c.redirect.bind(c),
   };
 }
+
+export function injectRoute(route: {
+  entryName: string;
+}): Middleware<ServerEnv> {
+  return async (c, next) => {
+    if (route && !c.get('route')) {
+      c.set('route', route);
+    }
+
+    await next();
+  };
+}
diff --git a/packages/server/core/src/plugins/render/index.ts b/packages/server/core/src/plugins/render/index.ts
@@ -10,7 +10,11 @@ import type {
   ServerPlugin,
 } from '../../types';
 import { sortRoutes } from '../../utils';
-import { CustomServer, getServerMidFromUnstableMid } from '../customServer';
+import {
+  CustomServer,
+  getServerMidFromUnstableMid,
+  injectRoute,
+} from '../customServer';
 import { initReporter } from '../monitors';
 
 export * from './inject';
@@ -44,21 +48,26 @@ export const renderPlugin = (): ServerPlugin => ({
         const pageRoutes = getPageRoutes(routes);
 
         for (const route of pageRoutes) {
-          const { urlPath: originUrlPath, entryName } = route;
+          const { urlPath: originUrlPath, entryName = MAIN_ENTRY_NAME } = route;
           const urlPath = originUrlPath.endsWith('/')
             ? `${originUrlPath}*`
             : `${originUrlPath}/*`;
 
           middlewares.push({
             name: 'init-reporter',
-            handler: initReporter(entryName || MAIN_ENTRY_NAME),
+            handler: initReporter(entryName),
           });
 
           const customServerHookMiddleware = customServer.getHookMiddleware(
-            entryName || 'main',
+            entryName,
             routes,
           );
 
+          middlewares.push({
+            name: 'inject-route-info',
+            handler: injectRoute({ entryName }),
+          });
+
           middlewares.push({
             name: 'custom-server-hook',
             path: urlPath,
diff --git a/packages/server/core/src/types/server.ts b/packages/server/core/src/types/server.ts
@@ -70,6 +70,11 @@ type ServerVariables = {
    * Custom by ssrRuntime.
    */
   locals?: Record<string, any>;
+
+  /**
+   * The current matched route, now only expose entryName field.
+   */
+  route: Required<Pick<ServerRoute, 'entryName'>>;
 };
 
 export type ServerEnv = {
diff --git a/packages/server/plugin-server/src/server.ts b/packages/server/plugin-server/src/server.ts
@@ -85,6 +85,7 @@ export default (): ServerPlugin => ({
     const pwd = isProd() ? distDirectory : appDirectory;
 
     const loadMod = async () => {
+      // TODO：future, Modern.js removes the Hook concept, and `_middleware` file conventions can also be removed
       const { middleware: unstableMiddleware } = await loadMiddleware(pwd);
       const {
         defaultExports,
diff --git a/packages/toolkit/types/server/middleware.d.ts b/packages/toolkit/types/server/middleware.d.ts
@@ -1,4 +1,5 @@
 import type { RequestPayload } from './context';
+import type { ServerRoute } from './route';
 
 type Set<V extends Record<string>> = <Key extends keyof V>(
   key: Key,
@@ -16,6 +17,7 @@ export type UnstableMiddlewareContext<
 > = {
   request: Request;
   response: Response;
+  route: Required<Pick<ServerRoute, 'entryName'>>;
   get: Get<V & RequestPayload>;
   set: Set<V & RequestPayload>;
   header: (name: string, value: string, options?: { append?: boolean }) => void;
diff --git a/tests/integration/server-hook/new-middleware/server/index.ts b/tests/integration/server-hook/new-middleware/server/index.ts
@@ -104,10 +104,19 @@ function modifyRequest(): UnstableMiddleware {
   };
 }
 
+function getRoute(): UnstableMiddleware {
+  return async (c, next) => {
+    c.response.headers.set('x-matched-route', c.route.entryName);
+
+    await next();
+  };
+}
+
 export const unstableMiddleware: UnstableMiddleware[] = [
   time(),
   modifyRequest(),
   injectRequestBody(),
   injectMessage(),
   auth() as unknown as UnstableMiddleware,
+  getRoute(),
 ];
diff --git a/tests/integration/server-hook/new-middleware/tests/index.test.ts b/tests/integration/server-hook/new-middleware/tests/index.test.ts
@@ -50,6 +50,8 @@ describe('test new middleware run correctly', () => {
     expect(body).toMatch(/lang="en"/);
 
     expect(headers).toHaveProperty('x-custom-value', 'modern');
+
+    expect(headers).toHaveProperty('x-matched-route', 'main');
   });
 
   test('should redirect corretly', async () => {
