From b483800e6a7d88205ec06b9a5c265149cc9554b2 Mon Sep 17 00:00:00 2001
From: simeng-li <simeng@silverhand.io>
Date: Mon, 11 Nov 2024 10:44:03 +0800
Subject: [PATCH] feat(core): redirect unknown session to app fallback uri

redirect unknown session to app fallback uri
---
 .../middleware/koa-spa-session-guard.test.ts  | 25 +++++++++++-
 .../src/middleware/koa-spa-session-guard.ts   | 40 ++++++++++++++++++-
 packages/core/src/oidc/utils.ts               |  6 ++-
 3 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/packages/core/src/middleware/koa-spa-session-guard.test.ts b/packages/core/src/middleware/koa-spa-session-guard.test.ts
index 9a68c290e..9ac01e995 100644
--- a/packages/core/src/middleware/koa-spa-session-guard.test.ts
+++ b/packages/core/src/middleware/koa-spa-session-guard.test.ts
@@ -6,6 +6,8 @@ import { EnvSet, UserApps } from '#src/env-set/index.js';
 import { MockQueries } from '#src/test-utils/tenant.js';
 import { createContextWithRouteParameters } from '#src/utils/test-utils.js';
 
+import { LoginQueryParamsKey } from '../oidc/utils.js';
+
 const { jest } = import.meta;
 
 const { mockEsmWithActual } = createMockUtils(jest);
@@ -25,7 +27,11 @@ describe('koaSpaSessionGuard', () => {
   const provider = new Provider('https://logto.test');
   const interactionDetails = jest.spyOn(provider, 'interactionDetails');
   const getRowsByKeys = jest.fn().mockResolvedValue({ rows: [] });
-  const queries = new MockQueries({ logtoConfigs: { getRowsByKeys } });
+  const findApplicationById = jest.fn();
+  const queries = new MockQueries({
+    logtoConfigs: { getRowsByKeys },
+    applications: { findApplicationById },
+  });
 
   beforeEach(() => {
     process.env = { ...envBackup };
@@ -115,4 +121,21 @@ describe('koaSpaSessionGuard', () => {
     expect(ctx.redirect).toBeCalledWith('https://test.com/unknown-session');
     stub.restore();
   });
+
+  it('should redirect to configured application fallback URL if session not found', async () => {
+    const unknownSessionFallbackUri = 'https://foo.bar';
+    findApplicationById.mockResolvedValueOnce({
+      unknownSessionFallbackUri,
+    });
+
+    const appId = '123';
+    interactionDetails.mockRejectedValue(new Error('session not found'));
+    const ctx = createContextWithRouteParameters({
+      url: `${guardedPath[0]!}?${LoginQueryParamsKey.AppId}=${appId}`,
+    });
+
+    await koaSpaSessionGuard(provider, queries)(ctx, next);
+    expect(findApplicationById).toBeCalledWith(appId);
+    expect(ctx.redirect).toBeCalledWith(unknownSessionFallbackUri);
+  });
 });
diff --git a/packages/core/src/middleware/koa-spa-session-guard.ts b/packages/core/src/middleware/koa-spa-session-guard.ts
index 66ebf5623..6e2f11185 100644
--- a/packages/core/src/middleware/koa-spa-session-guard.ts
+++ b/packages/core/src/middleware/koa-spa-session-guard.ts
@@ -1,6 +1,11 @@
-import { logtoConfigGuards, LogtoTenantConfigKey } from '@logto/schemas';
+import {
+  logtoConfigGuards,
+  logtoCookieKey,
+  LogtoTenantConfigKey,
+  logtoUiCookieGuard,
+} from '@logto/schemas';
 import { appendPath, trySafe } from '@silverhand/essentials';
-import type { MiddlewareType } from 'koa';
+import type { MiddlewareType, ParameterizedContext } from 'koa';
 import type { IRouterParamContext } from 'koa-router';
 import type Provider from 'oidc-provider';
 
@@ -9,6 +14,8 @@ import RequestError from '#src/errors/RequestError/index.js';
 import type Queries from '#src/tenants/Queries.js';
 import { getTenantId } from '#src/utils/tenant.js';
 
+import { LoginQueryParamsKey } from '../oidc/utils.js';
+
 // Need To Align With UI
 export const sessionNotFoundPath = '/unknown-session';
 
@@ -21,6 +28,24 @@ export const guardedPath = [
   '/forgot-password',
 ];
 
+/**
+ * Retrieve the appId from ctx for find the fallback url
+ * First check the query param appId
+ * If not found, check the logto ui cookie
+ */
+const getAppIdFromContext = (ctx: ParameterizedContext) => {
+  const appId = ctx.request.URL.searchParams.get(LoginQueryParamsKey.AppId);
+
+  if (appId) {
+    return appId;
+  }
+
+  const cookie = ctx.cookies.get(logtoCookieKey);
+  const parsed = trySafe(() => logtoUiCookieGuard.parse(cookie));
+
+  return parsed?.appId;
+};
+
 export default function koaSpaSessionGuard<
   StateT,
   ContextT extends IRouterParamContext,
@@ -37,6 +62,17 @@ export default function koaSpaSessionGuard<
       try {
         await provider.interactionDetails(ctx.req, ctx.res);
       } catch {
+        // Try to find the fallback url for the application
+        const appId = getAppIdFromContext(ctx);
+        const application = appId
+          ? await trySafe(async () => queries.applications.findApplicationById(appId))
+          : undefined;
+
+        if (application?.unknownSessionFallbackUri) {
+          ctx.redirect(application.unknownSessionFallbackUri);
+          return;
+        }
+
         const {
           rows: [data],
         } = await queries.logtoConfigs.getRowsByKeys([
diff --git a/packages/core/src/oidc/utils.ts b/packages/core/src/oidc/utils.ts
index e5faf35bc..11a17265e 100644
--- a/packages/core/src/oidc/utils.ts
+++ b/packages/core/src/oidc/utils.ts
@@ -96,6 +96,10 @@ const firstScreenRouteMapping: Record<FirstScreen, keyof typeof experience.route
   [FirstScreen.SignInDeprecated]: 'signIn',
 };
 
+export enum LoginQueryParamsKey {
+  AppId = 'app_id',
+}
+
 // Note: this eslint comment can be removed once the dev feature flag is removed
 
 export const buildLoginPromptUrl = (params: ExtraParamsObject, appId?: unknown): string => {
@@ -114,7 +118,7 @@ export const buildLoginPromptUrl = (params: ExtraParamsObject, appId?: unknown):
   const getSearchParamString = () => (searchParams.size > 0 ? `?${searchParams.toString()}` : '');
 
   if (appId) {
-    searchParams.append('app_id', String(appId));
+    searchParams.append(LoginQueryParamsKey.AppId, String(appId));
   }
 
   if (params[ExtraParamsKey.OrganizationId]) {