From 33a589a8ba01146a59cb362e79a0da0498eb19c4 Mon Sep 17 00:00:00 2001
From: DarkSky <darksky.graphite@toeverything.info>
Date: Tue, 19 Dec 2023 13:54:43 +0000
Subject: [PATCH] feat: onboarding electron redirect (#5327)

---
 .../auth-components/onboarding-page.tsx       | 38 ++++++++++++++++++-
 packages/frontend/core/src/pages/open-app.tsx |  5 ++-
 .../frontend/electron/src/main/deep-link.ts   |  4 ++
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/packages/frontend/component/src/components/auth-components/onboarding-page.tsx b/packages/frontend/component/src/components/auth-components/onboarding-page.tsx
index c49af66595fcd..a2c9ef9629e7f 100644
--- a/packages/frontend/component/src/components/auth-components/onboarding-page.tsx
+++ b/packages/frontend/component/src/components/auth-components/onboarding-page.tsx
@@ -2,6 +2,8 @@ import { fetchWithTraceReport } from '@affine/graphql';
 import { ArrowRightSmallIcon } from '@blocksuite/icons';
 import clsx from 'clsx';
 import { useMemo, useState } from 'react';
+// eslint-disable-next-line @typescript-eslint/no-restricted-imports
+import { type Location, useLocation, useNavigate } from 'react-router-dom';
 import useSWR from 'swr';
 
 import { Button } from '../../ui/button';
@@ -30,6 +32,23 @@ type QuestionnaireAnswer = {
   answer: string[];
 };
 
+function getCallbackUrl(location: Location) {
+  try {
+    const url =
+      location.state?.callbackURL ||
+      new URLSearchParams(location.search).get('callbackUrl');
+    if (typeof url === 'string' && url) {
+      if (!url.startsWith('http:') && !url.startsWith('https:')) {
+        return url;
+      }
+      // we will ignore host to avoid redirect hack
+      const parsedUrl = new URL(url);
+      return parsedUrl.pathname + parsedUrl.search;
+    }
+  } catch (_) {}
+  return null;
+}
+
 export const ScrollableLayout = ({
   children,
 }: {
@@ -69,6 +88,8 @@ export const OnboardingPage = ({
   user: User;
   onOpenAffine: () => void;
 }) => {
+  const location = useLocation();
+  const navigate = useNavigate();
   const [questionIdx, setQuestionIdx] = useState(0);
   const { data: questions } = useSWR<Question[]>(
     '/api/worker/questionnaire',
@@ -78,6 +99,7 @@ export const OnboardingPage = ({
   const [options, setOptions] = useState(new Set<string>());
   const [inputs, setInputs] = useState<Record<string, string>>({});
 
+  const callbackUrl = useMemo(() => getCallbackUrl(location), [location]);
   const question = useMemo(
     () => questions?.[questionIdx],
     [questionIdx, questions]
@@ -87,6 +109,14 @@ export const OnboardingPage = ({
     return null;
   }
 
+  if (callbackUrl?.startsWith('/open-app/signin-redirect')) {
+    const url = new URL(callbackUrl, window.location.origin);
+    url.searchParams.set('next', 'onboarding');
+    console.log('redirect to', url.toString());
+    window.location.assign(url.toString());
+    return null;
+  }
+
   if (question) {
     return (
       <ScrollableLayout>
@@ -198,7 +228,13 @@ export const OnboardingPage = ({
           className={clsx(styles.button, styles.openAFFiNEButton)}
           type="primary"
           size="extraLarge"
-          onClick={onOpenAffine}
+          onClick={() => {
+            if (callbackUrl) {
+              navigate(callbackUrl);
+            } else {
+              onOpenAffine();
+            }
+          }}
           iconPosition="end"
           icon={<ArrowRightSmallIcon />}
         >
diff --git a/packages/frontend/core/src/pages/open-app.tsx b/packages/frontend/core/src/pages/open-app.tsx
index 6941a45289e35..a6e849cf320b2 100644
--- a/packages/frontend/core/src/pages/open-app.tsx
+++ b/packages/frontend/core/src/pages/open-app.tsx
@@ -169,13 +169,16 @@ const OpenOAuthJwt = () => {
     const maybeSchema = appSchemas.safeParse(params.get('schema'));
     return maybeSchema.success ? maybeSchema.data : 'affine';
   }, [params]);
+  const next = useMemo(() => params.get('next'), [params]);
   const channel = schemaToChanel[schema as Schema];
 
   if (!currentUser || !currentUser?.token?.sessionToken) {
     return null;
   }
 
-  const urlToOpen = `${schema}://signin-redirect?token=${currentUser.token.sessionToken}`;
+  const urlToOpen = `${schema}://signin-redirect?token=${
+    currentUser.token.sessionToken
+  }&next=${next || ''}`;
 
   return <OpenAppImpl urlToOpen={urlToOpen} channel={channel} />;
 };
diff --git a/packages/frontend/electron/src/main/deep-link.ts b/packages/frontend/electron/src/main/deep-link.ts
index 93ee4739b0a06..a098bfb8615e0 100644
--- a/packages/frontend/electron/src/main/deep-link.ts
+++ b/packages/frontend/electron/src/main/deep-link.ts
@@ -108,6 +108,10 @@ async function handleOauthJwt(url: string) {
 
       ipcMain.once('affine:login', () => {
         hiddenWindow?.destroy();
+        if (urlObj.searchParams.get('next') === 'onboarding') {
+          // eslint-disable-next-line @typescript-eslint/no-floating-promises
+          mainWindow.loadURL(mainWindowOrigin + '/auth/onboarding');
+        }
       });
 
       // hacks to refresh auth state in the main window