From 862d27dee4436f267651ab55a92d9442b66a0612 Mon Sep 17 00:00:00 2001
From: simeng-li <simeng@silverhand.io>
Date: Tue, 10 Sep 2024 14:37:01 +0800
Subject: [PATCH] fix(core): fix custom UI not triggered bug (#6563)

fix custom UI not triggered on legacy-experience bug
---
 .../core/src/middleware/koa-spa-proxy.test.ts | 68 ++++++++++++-------
 packages/core/src/middleware/koa-spa-proxy.ts | 16 ++++-
 2 files changed, 58 insertions(+), 26 deletions(-)

diff --git a/packages/core/src/middleware/koa-spa-proxy.test.ts b/packages/core/src/middleware/koa-spa-proxy.test.ts
index 23f003c41..026aa1301 100644
--- a/packages/core/src/middleware/koa-spa-proxy.test.ts
+++ b/packages/core/src/middleware/koa-spa-proxy.test.ts
@@ -1,3 +1,5 @@
+import path from 'node:path';
+
 import { pickDefault, createMockUtils } from '@logto/shared/esm';
 import Sinon from 'sinon';
 
@@ -12,13 +14,14 @@ const mockProxyMiddleware = jest.fn();
 const mockStaticMiddleware = jest.fn();
 const mockCustomUiAssetsMiddleware = jest.fn();
 const mountedApps = Object.values(UserApps);
+const mockStaticMiddlewareFactory = jest.fn(() => mockStaticMiddleware);
 
 mockEsmDefault('node:fs/promises', () => ({
   readdir: jest.fn().mockResolvedValue(['sign-in']),
 }));
 
 mockEsmDefault('koa-proxies', () => jest.fn(() => mockProxyMiddleware));
-mockEsmDefault('#src/middleware/koa-serve-static.js', () => jest.fn(() => mockStaticMiddleware));
+mockEsmDefault('#src/middleware/koa-serve-static.js', () => mockStaticMiddlewareFactory);
 mockEsmDefault('#src/middleware/koa-serve-custom-ui-assets.js', () =>
   jest.fn(() => mockCustomUiAssetsMiddleware)
 );
@@ -83,32 +86,49 @@ describe('koaSpaProxy middleware', () => {
     stub.restore();
   });
 
-  it('production env should call the static middleware if path hit the ui file directory', async () => {
-    const stub = Sinon.stub(EnvSet, 'values').value({
-      ...EnvSet.values,
-      isProduction: true,
-    });
+  it.each([true, false])(
+    'production env should call the proxy middleware if path does not hit the ui file directory: devFeatureEnabled %p',
+    async (isDevFeaturesEnabled) => {
+      const stub = Sinon.stub(EnvSet, 'values').value({
+        ...EnvSet.values,
+        isProduction: true,
+        isDevFeaturesEnabled,
+      });
 
-    const ctx = createContextWithRouteParameters({
-      url: '/sign-in',
-    });
+      const ctx = createContextWithRouteParameters({
+        url: '/sign-in',
+      });
 
-    await koaSpaProxy({ mountedApps, queries })(ctx, next);
-    expect(mockStaticMiddleware).toBeCalled();
-    stub.restore();
-  });
+      await koaSpaProxy({ mountedApps, queries })(ctx, next);
 
-  it('should serve custom UI assets if user uploaded them', async () => {
-    const customUiAssets = { id: 'custom-ui-assets', createdAt: Date.now() };
-    mockFindDefaultSignInExperience.mockResolvedValue({ customUiAssets });
+      const packagePath = isDevFeaturesEnabled ? 'experience' : 'experience-legacy';
+      const distributionPath = path.join('node_modules/@logto', packagePath, 'dist');
 
-    const ctx = createContextWithRouteParameters({
-      url: '/sign-in',
-    });
+      expect(mockStaticMiddlewareFactory).toBeCalledWith(distributionPath);
+      stub.restore();
+    }
+  );
 
-    await koaSpaProxy({ mountedApps, queries })(ctx, next);
-    expect(mockCustomUiAssetsMiddleware).toBeCalled();
-    expect(mockStaticMiddleware).not.toBeCalled();
-    expect(mockProxyMiddleware).not.toBeCalled();
-  });
+  it.each([true, false])(
+    'should serve custom UI assets if user uploaded them: : devFeatureEnabled %p',
+    async (isDevFeaturesEnabled) => {
+      const stub = Sinon.stub(EnvSet, 'values').value({
+        ...EnvSet.values,
+        isDevFeaturesEnabled,
+      });
+
+      const customUiAssets = { id: 'custom-ui-assets', createdAt: Date.now() };
+      mockFindDefaultSignInExperience.mockResolvedValue({ customUiAssets });
+
+      const ctx = createContextWithRouteParameters({
+        url: '/sign-in',
+      });
+
+      await koaSpaProxy({ mountedApps, queries })(ctx, next);
+      expect(mockCustomUiAssetsMiddleware).toBeCalled();
+      expect(mockStaticMiddleware).not.toBeCalled();
+      expect(mockProxyMiddleware).not.toBeCalled();
+      stub.restore();
+    }
+  );
 });
diff --git a/packages/core/src/middleware/koa-spa-proxy.ts b/packages/core/src/middleware/koa-spa-proxy.ts
index 71b65058b..25314d205 100644
--- a/packages/core/src/middleware/koa-spa-proxy.ts
+++ b/packages/core/src/middleware/koa-spa-proxy.ts
@@ -20,16 +20,27 @@ type Properties = {
   readonly prefix?: string;
 };
 
+const getDistributionPath = (packagePath: string) => {
+  if (packagePath === 'experience') {
+    // Use the new experience package if dev features are enabled
+    const moduleName = EnvSet.values.isDevFeaturesEnabled ? 'experience' : 'experience-legacy';
+
+    return path.join('node_modules/@logto', moduleName, 'dist');
+  }
+
+  return path.join('node_modules/@logto', packagePath, 'dist');
+};
+
 export default function koaSpaProxy<StateT, ContextT extends IRouterParamContext, ResponseBodyT>({
   mountedApps,
-  packagePath = EnvSet.values.isDevFeaturesEnabled ? 'experience' : 'experience-legacy',
+  packagePath = 'experience',
   port = 5001,
   prefix = '',
   queries,
 }: Properties): MiddlewareType<StateT, ContextT, ResponseBodyT> {
   type Middleware = MiddlewareType<StateT, ContextT, ResponseBodyT>;
 
-  const distributionPath = path.join('node_modules/@logto', packagePath, 'dist');
+  const distributionPath = getDistributionPath(packagePath);
 
   const spaProxy: Middleware = EnvSet.values.isProduction
     ? serveStatic(distributionPath)
@@ -58,6 +69,7 @@ export default function koaSpaProxy<StateT, ContextT extends IRouterParamContext
     }
 
     const { customUiAssets } = await queries.signInExperiences.findDefaultSignInExperience();
+
     // If user has uploaded custom UI assets, serve them instead of native experience UI
     if (customUiAssets && packagePath === 'experience') {
       const serve = serveCustomUiAssets(customUiAssets.id);