From a8355eb6a305059a48d445bd324dd327bcf16a85 Mon Sep 17 00:00:00 2001
From: Gao Sun <gao@silverhand.io>
Date: Thu, 2 Mar 2023 16:45:48 +0800
Subject: [PATCH] refactor(core): fix segment matching for path-based

---
 packages/core/src/utils/tenant.test.ts | 19 +++++++++++-
 packages/core/src/utils/tenant.ts      | 43 +++++++++++++++++---------
 2 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/packages/core/src/utils/tenant.test.ts b/packages/core/src/utils/tenant.test.ts
index de1cd2b78..b7e13382b 100644
--- a/packages/core/src/utils/tenant.test.ts
+++ b/packages/core/src/utils/tenant.test.ts
@@ -77,7 +77,7 @@ describe('getTenantId()', () => {
       ADMIN_DISABLE_LOCALHOST: '1',
     };
 
-    expect(getTenantId(new URL('http://localhost:5000/app///asdasd'))).toBe(defaultTenantId);
+    expect(getTenantId(new URL('http://localhost:5000/app///asdasd'))).toBe(undefined);
     expect(getTenantId(new URL('http://localhost:3002/app///asdasd'))).toBe(undefined);
     expect(getTenantId(new URL('https://user.foo.logto.mock/app'))).toBe('foo');
     expect(getTenantId(new URL('https://user.admin.logto.mock/app//'))).toBe(undefined); // Admin endpoint is explicitly set
@@ -92,4 +92,21 @@ describe('getTenantId()', () => {
     };
     expect(getTenantId(new URL('https://user.admin.logto.mock/app//'))).toBe('admin');
   });
+
+  it('should resolve proper tenant ID for path-based multi-tenancy', async () => {
+    process.env = {
+      ...backupEnv,
+      NODE_ENV: 'production',
+      PORT: '5000',
+      ENDPOINT: 'https://user.logto.mock/app',
+      PATH_BASED_MULTI_TENANCY: '1',
+    };
+
+    expect(getTenantId(new URL('http://localhost:5000/app///asdasd'))).toBe('app');
+    expect(getTenantId(new URL('http://localhost:3002///bar///asdasd'))).toBe(adminTenantId);
+    expect(getTenantId(new URL('https://user.foo.logto.mock/app'))).toBe(undefined);
+    expect(getTenantId(new URL('https://user.admin.logto.mock/app//'))).toBe(undefined);
+    expect(getTenantId(new URL('https://user.logto.mock/app'))).toBe(undefined);
+    expect(getTenantId(new URL('https://user.logto.mock/app/admin'))).toBe('admin');
+  });
 });
diff --git a/packages/core/src/utils/tenant.ts b/packages/core/src/utils/tenant.ts
index 5ca1031bf..34c7d4bf6 100644
--- a/packages/core/src/utils/tenant.ts
+++ b/packages/core/src/utils/tenant.ts
@@ -1,6 +1,7 @@
 import { adminTenantId, defaultTenantId } from '@logto/schemas';
 import { conditionalString } from '@silverhand/essentials';
 
+import type UrlSet from '#src/env-set/UrlSet.js';
 import { EnvSet, getTenantEndpoint } from '#src/env-set/index.js';
 
 const normalizePathname = (pathname: string) =>
@@ -14,6 +15,32 @@ const isEndpointOf = (current: URL, endpoint: URL) => {
   );
 };
 
+const matchDomainBasedTenantId = (pattern: URL, url: URL) => {
+  const toMatch = pattern.hostname.replace('*', '([^.]*)');
+  const matchedId = new RegExp(toMatch).exec(url.hostname)?.[1];
+
+  if (!matchedId || matchedId === '*') {
+    return;
+  }
+
+  if (isEndpointOf(url, getTenantEndpoint(matchedId, EnvSet.values))) {
+    return matchedId;
+  }
+};
+
+const matchPathBasedTenantId = (urlSet: UrlSet, url: URL) => {
+  const found = urlSet.deduplicated().find((value) => isEndpointOf(url, value));
+
+  if (!found) {
+    return;
+  }
+
+  const urlSegments = url.pathname.split('/');
+  const endpointSegments = found.pathname.split('/');
+
+  return urlSegments[found.pathname === '/' ? 1 : endpointSegments.length];
+};
+
 export const getTenantId = (url: URL) => {
   const {
     isDomainBasedMultiTenancy,
@@ -40,20 +67,8 @@ export const getTenantId = (url: URL) => {
   }
 
   if (isPathBasedMultiTenancy) {
-    const urlSegments = url.pathname.split('/');
-    const endpointSegments = urlSet.endpoint.pathname.split('/');
-
-    return urlSegments[endpointSegments.length - 1];
+    return matchPathBasedTenantId(urlSet, url);
   }
 
-  const toMatch = urlSet.endpoint.hostname.replace('*', '([^.]*)');
-  const matchedId = new RegExp(toMatch).exec(url.hostname)?.[1];
-
-  if (!matchedId || matchedId === '*') {
-    return;
-  }
-
-  if (isEndpointOf(url, getTenantEndpoint(matchedId, EnvSet.values))) {
-    return matchedId;
-  }
+  return matchDomainBasedTenantId(urlSet.endpoint, url);
 };