From fd3ce431fe7fb7787253a658bb2ff91b70cac26a Mon Sep 17 00:00:00 2001
From: fengmk2 <fengmk2@gmail.com>
Date: Wed, 12 Mar 2025 08:42:35 +0000
Subject: [PATCH] fix(core): assert app schema url on open-app (#10687)

---
 .../core/src/desktop/pages/open-app/index.tsx   | 17 +++++++++++------
 packages/frontend/core/src/utils/channel.ts     | 12 ++++++++++++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/packages/frontend/core/src/desktop/pages/open-app/index.tsx b/packages/frontend/core/src/desktop/pages/open-app/index.tsx
index 8a71d7b2b1..4978aaedff 100644
--- a/packages/frontend/core/src/desktop/pages/open-app/index.tsx
+++ b/packages/frontend/core/src/desktop/pages/open-app/index.tsx
@@ -1,7 +1,11 @@
 import { useNavigateHelper } from '@affine/core/components/hooks/use-navigate-helper';
 import { GraphQLService } from '@affine/core/modules/cloud';
 import { OpenInAppPage } from '@affine/core/modules/open-in-app/views/open-in-app-page';
-import { appSchemes, channelToScheme } from '@affine/core/utils/channel';
+import {
+  appSchemaUrl,
+  appSchemes,
+  channelToScheme,
+} from '@affine/core/utils/channel';
 import type { GetCurrentUserQuery } from '@affine/graphql';
 import { getCurrentUserQuery } from '@affine/graphql';
 import { useService } from '@toeverything/infra';
@@ -23,12 +27,13 @@ const OpenUrl = () => {
     [navigateHelper]
   );
 
-  if (!urlToOpen) {
+  const parsed = appSchemaUrl.safeParse(urlToOpen);
+  if (!parsed.success) {
+    console.error(parsed.error);
     return null;
   }
 
-  const urlObj = new URL(urlToOpen || '');
-
+  const urlObj = new URL(parsed.data);
   params.forEach((v, k) => {
     if (k === 'url') {
       return;
@@ -55,7 +60,7 @@ const OpenOAuthJwt = () => {
   const scheme = maybeScheme.success
     ? maybeScheme.data
     : channelToScheme[BUILD_CONFIG.appBuildType];
-  const next = params.get('next');
+  const next = params.get('next') || '';
 
   useEffect(() => {
     graphqlService
@@ -74,7 +79,7 @@ const OpenOAuthJwt = () => {
 
   const urlToOpen = `${scheme}://signin-redirect?token=${
     currentUser.token.sessionToken
-  }&next=${next || ''}`;
+  }&next=${next}`;
 
   return <OpenInAppPage urlToOpen={urlToOpen} />;
 };
diff --git a/packages/frontend/core/src/utils/channel.ts b/packages/frontend/core/src/utils/channel.ts
index 56a5762bdf..d1a16ae639 100644
--- a/packages/frontend/core/src/utils/channel.ts
+++ b/packages/frontend/core/src/utils/channel.ts
@@ -42,3 +42,15 @@ export const appNames = {
   beta: 'AFFiNE Beta',
   internal: 'AFFiNE Internal',
 } satisfies Record<Channel, string>;
+
+export const appSchemaUrl = z.custom<string>(
+  (url: string) => {
+    try {
+      return appSchemes.safeParse(new URL(url).protocol.replace(':', ''))
+        .success;
+    } catch {
+      return false;
+    }
+  },
+  { message: 'Invalid URL or protocol' }
+);