From 8f694aceb7b048794eb3b6ae2efe105ebe61153f Mon Sep 17 00:00:00 2001
From: pengx17 <pengxiao@outlook.com>
Date: Fri, 25 Oct 2024 03:43:00 +0000
Subject: [PATCH] fix(core): redirect to old page after login via 404 page
 (#8588)

fix AF-1487

When visit a cloud worskapce page without login, we should allow the user to redirect back to that page after login.

The logic is only added for this case, including login with email magin link + google login. Login with password is already working without changes.
---
 .../server/src/plugins/oauth/controller.ts    |  1 -
 .../affine/auth/after-sign-in-send-email.tsx  | 10 ++++-
 .../affine/auth/after-sign-up-send-email.tsx  | 11 ++++--
 .../core/src/components/affine/auth/index.tsx | 10 ++++-
 .../core/src/components/affine/auth/oauth.tsx | 38 ++++++++++++++-----
 .../src/components/affine/auth/send-email.tsx |  1 +
 .../affine/auth/sign-in-with-password.tsx     | 18 ++++++++-
 .../src/components/affine/auth/sign-in.tsx    | 18 +++++++--
 .../frontend/core/src/desktop/pages/404.tsx   |  6 ++-
 .../src/desktop/pages/auth/magic-link.tsx     | 18 +++++++--
 .../src/desktop/pages/auth/oauth-callback.tsx | 10 ++++-
 .../core/src/desktop/pages/auth/sign-in.tsx   | 16 +++++---
 .../core/src/modules/cloud/services/auth.ts   | 13 ++++++-
 13 files changed, 135 insertions(+), 35 deletions(-)

diff --git a/packages/backend/server/src/plugins/oauth/controller.ts b/packages/backend/server/src/plugins/oauth/controller.ts
index 999c34ef3d..d553aedf22 100644
--- a/packages/backend/server/src/plugins/oauth/controller.ts
+++ b/packages/backend/server/src/plugins/oauth/controller.ts
@@ -111,7 +111,6 @@ export class OAuthController {
     await this.auth.setCookies(req, res, user.id);
     res.send({
       id: user.id,
-      /* @deprecated */
       redirectUri: state.redirectUri,
     });
   }
diff --git a/packages/frontend/core/src/components/affine/auth/after-sign-in-send-email.tsx b/packages/frontend/core/src/components/affine/auth/after-sign-in-send-email.tsx
index 99916c6424..d2e4b68d41 100644
--- a/packages/frontend/core/src/components/affine/auth/after-sign-in-send-email.tsx
+++ b/packages/frontend/core/src/components/affine/auth/after-sign-in-send-email.tsx
@@ -19,6 +19,7 @@ import { Captcha, useCaptcha } from './use-captcha';
 export const AfterSignInSendEmail = ({
   setAuthData: setAuth,
   email,
+  redirectUrl,
 }: AuthPanelProps<'afterSignInSendEmail'>) => {
   const [resendCountDown, setResendCountDown] = useState(60);
 
@@ -44,7 +45,12 @@ export const AfterSignInSendEmail = ({
     try {
       if (verifyToken) {
         setResendCountDown(60);
-        await authService.sendEmailMagicLink(email, verifyToken, challenge);
+        await authService.sendEmailMagicLink(
+          email,
+          verifyToken,
+          challenge,
+          redirectUrl
+        );
       }
     } catch (err) {
       console.error(err);
@@ -53,7 +59,7 @@ export const AfterSignInSendEmail = ({
       });
     }
     setIsSending(false);
-  }, [authService, challenge, email, verifyToken]);
+  }, [authService, challenge, email, redirectUrl, verifyToken]);
 
   const onSignInWithPasswordClick = useCallback(() => {
     setAuth({ state: 'signInWithPassword' });
diff --git a/packages/frontend/core/src/components/affine/auth/after-sign-up-send-email.tsx b/packages/frontend/core/src/components/affine/auth/after-sign-up-send-email.tsx
index be1e0c0b96..04c9c257a2 100644
--- a/packages/frontend/core/src/components/affine/auth/after-sign-up-send-email.tsx
+++ b/packages/frontend/core/src/components/affine/auth/after-sign-up-send-email.tsx
@@ -19,7 +19,7 @@ import { Captcha, useCaptcha } from './use-captcha';
 
 export const AfterSignUpSendEmail: FC<
   AuthPanelProps<'afterSignUpSendEmail'>
-> = ({ setAuthData, email }) => {
+> = ({ setAuthData, email, redirectUrl }) => {
   const [resendCountDown, setResendCountDown] = useState(60);
 
   useEffect(() => {
@@ -42,7 +42,12 @@ export const AfterSignUpSendEmail: FC<
     setIsSending(true);
     try {
       if (verifyToken) {
-        await authService.sendEmailMagicLink(email, verifyToken, challenge);
+        await authService.sendEmailMagicLink(
+          email,
+          verifyToken,
+          challenge,
+          redirectUrl
+        );
       }
       setResendCountDown(60);
     } catch (err) {
@@ -52,7 +57,7 @@ export const AfterSignUpSendEmail: FC<
       });
     }
     setIsSending(false);
-  }, [authService, challenge, email, verifyToken]);
+  }, [authService, challenge, email, redirectUrl, verifyToken]);
 
   return (
     <>
diff --git a/packages/frontend/core/src/components/affine/auth/index.tsx b/packages/frontend/core/src/components/affine/auth/index.tsx
index bf4a2642fe..b6ec7edfe2 100644
--- a/packages/frontend/core/src/components/affine/auth/index.tsx
+++ b/packages/frontend/core/src/components/affine/auth/index.tsx
@@ -30,6 +30,7 @@ export type AuthPanelProps<State extends AuthAtomData['state']> = {
     updates: { state: T } & Difference<AuthAtomType<State>, AuthAtomType<T>>
   ) => void;
   onSkip?: () => void;
+  redirectUrl?: string;
 } & Extract<AuthAtomData, { state: State }>;
 
 const config: {
@@ -58,7 +59,13 @@ export function AuthModal() {
   );
 }
 
-export function AuthPanel({ onSkip }: { onSkip?: () => void }) {
+export function AuthPanel({
+  onSkip,
+  redirectUrl,
+}: {
+  onSkip?: () => void;
+  redirectUrl?: string | null;
+}) {
   const t = useI18n();
   const [authAtomValue, setAuthAtom] = useAtom(authAtom);
   const authService = useService(AuthService);
@@ -98,6 +105,7 @@ export function AuthPanel({ onSkip }: { onSkip?: () => void }) {
   const props = {
     ...authAtomValue,
     onSkip,
+    redirectUrl,
     setAuthData,
   };
 
diff --git a/packages/frontend/core/src/components/affine/auth/oauth.tsx b/packages/frontend/core/src/components/affine/auth/oauth.tsx
index 8f765678ae..f464e160ed 100644
--- a/packages/frontend/core/src/components/affine/auth/oauth.tsx
+++ b/packages/frontend/core/src/components/affine/auth/oauth.tsx
@@ -29,7 +29,7 @@ const OAuthProviderMap: Record<
   },
 };
 
-export function OAuth() {
+export function OAuth({ redirectUrl }: { redirectUrl?: string }) {
   const serverConfig = useService(ServerConfigService).serverConfig;
   const oauth = useLiveData(serverConfig.features$.map(r => r?.oauth));
   const oauthProviders = useLiveData(
@@ -41,25 +41,43 @@ export function OAuth() {
   }
 
   return oauthProviders?.map(provider => (
-    <OAuthProvider key={provider} provider={provider} />
+    <OAuthProvider
+      key={provider}
+      provider={provider}
+      redirectUrl={redirectUrl}
+    />
   ));
 }
 
-function OAuthProvider({ provider }: { provider: OAuthProviderType }) {
+function OAuthProvider({
+  provider,
+  redirectUrl,
+}: {
+  provider: OAuthProviderType;
+  redirectUrl?: string;
+}) {
   const { icon } = OAuthProviderMap[provider];
 
   const onClick = useCallback(() => {
-    let oauthUrl =
-      (BUILD_CONFIG.isElectron || BUILD_CONFIG.isIOS || BUILD_CONFIG.isAndroid
-        ? BUILD_CONFIG.serverUrlPrefix
-        : '') + `/oauth/login?provider=${provider}`;
+    const params = new URLSearchParams();
 
-    if (BUILD_CONFIG.isElectron) {
-      oauthUrl += `&client=${appInfo?.schema}`;
+    params.set('provider', provider);
+
+    if (redirectUrl) {
+      params.set('redirect_uri', redirectUrl);
     }
 
+    if (BUILD_CONFIG.isElectron && appInfo) {
+      params.set('client', appInfo.schema);
+    }
+
+    const oauthUrl =
+      (BUILD_CONFIG.isElectron || BUILD_CONFIG.isIOS || BUILD_CONFIG.isAndroid
+        ? BUILD_CONFIG.serverUrlPrefix
+        : '') + `/oauth/login?${params.toString()}`;
+
     popupWindow(oauthUrl);
-  }, [provider]);
+  }, [provider, redirectUrl]);
 
   return (
     <Button
diff --git a/packages/frontend/core/src/components/affine/auth/send-email.tsx b/packages/frontend/core/src/components/affine/auth/send-email.tsx
index cea9e1bbdb..b42d60b7e0 100644
--- a/packages/frontend/core/src/components/affine/auth/send-email.tsx
+++ b/packages/frontend/core/src/components/affine/auth/send-email.tsx
@@ -141,6 +141,7 @@ export const SendEmail = ({
   setAuthData,
   email,
   emailType,
+  // todo(@pengx17): impl redirectUrl for sendEmail?
 }: AuthPanelProps<'sendEmail'>) => {
   const t = useI18n();
   const serverConfig = useService(ServerConfigService).serverConfig;
diff --git a/packages/frontend/core/src/components/affine/auth/sign-in-with-password.tsx b/packages/frontend/core/src/components/affine/auth/sign-in-with-password.tsx
index ff2505625a..cad1014991 100644
--- a/packages/frontend/core/src/components/affine/auth/sign-in-with-password.tsx
+++ b/packages/frontend/core/src/components/affine/auth/sign-in-with-password.tsx
@@ -19,6 +19,7 @@ import { useCaptcha } from './use-captcha';
 export const SignInWithPassword: FC<AuthPanelProps<'signInWithPassword'>> = ({
   setAuthData,
   email,
+  redirectUrl,
 }) => {
   const t = useI18n();
   const authService = useService(AuthService);
@@ -62,7 +63,12 @@ export const SignInWithPassword: FC<AuthPanelProps<'signInWithPassword'>> = ({
     setSendingEmail(true);
     try {
       if (verifyToken) {
-        await authService.sendEmailMagicLink(email, verifyToken, challenge);
+        await authService.sendEmailMagicLink(
+          email,
+          verifyToken,
+          challenge,
+          redirectUrl
+        );
         setAuthData({ state: 'afterSignInSendEmail' });
       }
     } catch (err) {
@@ -73,7 +79,15 @@ export const SignInWithPassword: FC<AuthPanelProps<'signInWithPassword'>> = ({
       // TODO(@eyhn): handle error better
     }
     setSendingEmail(false);
-  }, [sendingEmail, verifyToken, authService, email, challenge, setAuthData]);
+  }, [
+    sendingEmail,
+    verifyToken,
+    authService,
+    email,
+    challenge,
+    redirectUrl,
+    setAuthData,
+  ]);
 
   const sendChangePasswordEmail = useCallback(() => {
     setAuthData({ state: 'sendEmail', emailType: 'changePassword' });
diff --git a/packages/frontend/core/src/components/affine/auth/sign-in.tsx b/packages/frontend/core/src/components/affine/auth/sign-in.tsx
index 1d11642271..9d339f2dc8 100644
--- a/packages/frontend/core/src/components/affine/auth/sign-in.tsx
+++ b/packages/frontend/core/src/components/affine/auth/sign-in.tsx
@@ -24,6 +24,7 @@ function validateEmail(email: string) {
 export const SignIn: FC<AuthPanelProps<'signIn'>> = ({
   setAuthData: setAuthState,
   onSkip,
+  redirectUrl,
 }) => {
   const t = useI18n();
   const authService = useService(AuthService);
@@ -59,14 +60,24 @@ export const SignIn: FC<AuthPanelProps<'signIn'>> = ({
               email,
             });
           } else {
-            await authService.sendEmailMagicLink(email, verifyToken, challenge);
+            await authService.sendEmailMagicLink(
+              email,
+              verifyToken,
+              challenge,
+              redirectUrl
+            );
             setAuthState({
               state: 'afterSignInSendEmail',
               email,
             });
           }
         } else {
-          await authService.sendEmailMagicLink(email, verifyToken, challenge);
+          await authService.sendEmailMagicLink(
+            email,
+            verifyToken,
+            challenge,
+            redirectUrl
+          );
           setAuthState({
             state: 'afterSignUpSendEmail',
             email,
@@ -87,6 +98,7 @@ export const SignIn: FC<AuthPanelProps<'signIn'>> = ({
     authService,
     challenge,
     email,
+    redirectUrl,
     refreshChallenge,
     setAuthState,
     verifyToken,
@@ -99,7 +111,7 @@ export const SignIn: FC<AuthPanelProps<'signIn'>> = ({
         subTitle={t['com.affine.brand.affineCloud']()}
       />
 
-      <OAuth />
+      <OAuth redirectUrl={redirectUrl} />
 
       <div className={style.authModalContent}>
         <AuthInput
diff --git a/packages/frontend/core/src/desktop/pages/404.tsx b/packages/frontend/core/src/desktop/pages/404.tsx
index 01a61399a5..a8aedd57e0 100644
--- a/packages/frontend/core/src/desktop/pages/404.tsx
+++ b/packages/frontend/core/src/desktop/pages/404.tsx
@@ -44,6 +44,10 @@ export const PageNotFound = ({
     apis?.ui.pingAppLayoutReady().catch(console.error);
   }, []);
 
+  // not using workbench location or router location deliberately
+  // strip the origin
+  const currentUrl = window.location.href.replace(window.location.origin, '');
+
   return (
     <>
       {noPermission ? (
@@ -51,7 +55,7 @@ export const PageNotFound = ({
           user={account}
           onBack={handleBackButtonClick}
           onSignOut={handleOpenSignOutModal}
-          signInComponent={<SignIn />}
+          signInComponent={<SignIn redirectUrl={currentUrl} />}
         />
       ) : (
         <NotFoundPage
diff --git a/packages/frontend/core/src/desktop/pages/auth/magic-link.tsx b/packages/frontend/core/src/desktop/pages/auth/magic-link.tsx
index 85e1b7cb63..ce8ba88e1c 100644
--- a/packages/frontend/core/src/desktop/pages/auth/magic-link.tsx
+++ b/packages/frontend/core/src/desktop/pages/auth/magic-link.tsx
@@ -1,5 +1,5 @@
 import { useService } from '@toeverything/infra';
-import { useEffect } from 'react';
+import { useEffect, useRef } from 'react';
 import {
   type LoaderFunction,
   redirect,
@@ -59,17 +59,29 @@ export const Component = () => {
   const data = useLoaderData() as LoaderData;
 
   const nav = useNavigate();
+  // loader data from useLoaderData is not reactive, so that we can safely
+  // assume the effect below is only triggered once
+  const triggeredRef = useRef(false);
 
   useEffect(() => {
+    if (triggeredRef.current) {
+      return;
+    }
+    triggeredRef.current = true;
     auth
       .signInMagicLink(data.email, data.token)
       .then(() => {
-        nav(data.redirectUri ?? '/');
+        const subscription = auth.session.status$.subscribe(status => {
+          if (status === 'authenticated') {
+            nav(data.redirectUri ?? '/');
+            subscription?.unsubscribe();
+          }
+        });
       })
       .catch(e => {
         nav(`/sign-in?error=${encodeURIComponent(e.message)}`);
       });
-  }, [data, auth, nav]);
+  }, [auth, data, data.email, data.redirectUri, data.token, nav]);
 
   return null;
 };
diff --git a/packages/frontend/core/src/desktop/pages/auth/oauth-callback.tsx b/packages/frontend/core/src/desktop/pages/auth/oauth-callback.tsx
index c886447cd9..0cd92e6e97 100644
--- a/packages/frontend/core/src/desktop/pages/auth/oauth-callback.tsx
+++ b/packages/frontend/core/src/desktop/pages/auth/oauth-callback.tsx
@@ -1,5 +1,5 @@
 import { useService } from '@toeverything/infra';
-import { useEffect } from 'react';
+import { useEffect, useRef } from 'react';
 import {
   type LoaderFunction,
   redirect,
@@ -62,9 +62,17 @@ export const Component = () => {
   const auth = useService(AuthService);
   const data = useLoaderData() as LoaderData;
 
+  // loader data from useLoaderData is not reactive, so that we can safely
+  // assume the effect below is only triggered once
+  const triggeredRef = useRef(false);
+
   const nav = useNavigate();
 
   useEffect(() => {
+    if (triggeredRef.current) {
+      return;
+    }
+    triggeredRef.current = true;
     auth
       .signInOauth(data.code, data.state, data.provider)
       .then(({ redirectUri }) => {
diff --git a/packages/frontend/core/src/desktop/pages/auth/sign-in.tsx b/packages/frontend/core/src/desktop/pages/auth/sign-in.tsx
index 89c4213f77..ecebc15cd3 100644
--- a/packages/frontend/core/src/desktop/pages/auth/sign-in.tsx
+++ b/packages/frontend/core/src/desktop/pages/auth/sign-in.tsx
@@ -12,7 +12,11 @@ import {
   useNavigateHelper,
 } from '../../../components/hooks/use-navigate-helper';
 
-export const SignIn = () => {
+export const SignIn = ({
+  redirectUrl: redirectUrlFromProps,
+}: {
+  redirectUrl?: string;
+}) => {
   const session = useService(AuthService).session;
   const status = useLiveData(session.status$);
   const isRevalidating = useLiveData(session.isRevalidating$);
@@ -20,12 +24,12 @@ export const SignIn = () => {
   const { jumpToIndex } = useNavigateHelper();
   const [searchParams] = useSearchParams();
   const isLoggedIn = status === 'authenticated' && !isRevalidating;
+  const redirectUrl = redirectUrlFromProps ?? searchParams.get('redirect_uri');
 
   useEffect(() => {
     if (isLoggedIn) {
-      const redirectUri = searchParams.get('redirect_uri');
-      if (redirectUri) {
-        navigate(redirectUri, {
+      if (redirectUrl) {
+        navigate(redirectUrl, {
           replace: true,
         });
       } else {
@@ -34,12 +38,12 @@ export const SignIn = () => {
         });
       }
     }
-  }, [jumpToIndex, navigate, isLoggedIn, searchParams]);
+  }, [jumpToIndex, navigate, isLoggedIn, redirectUrl, searchParams]);
 
   return (
     <SignInPageContainer>
       <div style={{ maxWidth: '400px', width: '100%' }}>
-        <AuthPanel onSkip={jumpToIndex} />
+        <AuthPanel onSkip={jumpToIndex} redirectUrl={redirectUrl} />
       </div>
     </SignInPageContainer>
   );
diff --git a/packages/frontend/core/src/modules/cloud/services/auth.ts b/packages/frontend/core/src/modules/cloud/services/auth.ts
index b8118c2be3..c756fb67ba 100644
--- a/packages/frontend/core/src/modules/cloud/services/auth.ts
+++ b/packages/frontend/core/src/modules/cloud/services/auth.ts
@@ -112,17 +112,26 @@ export class AuthService extends Service {
   async sendEmailMagicLink(
     email: string,
     verifyToken: string,
-    challenge?: string
+    challenge?: string,
+    redirectUrl?: string // url to redirect to after signed-in
   ) {
     track.$.$.auth.signIn({ method: 'magic-link' });
     try {
+      const magicLinkUrlParams = new URLSearchParams();
+      if (redirectUrl) {
+        magicLinkUrlParams.set('redirect_uri', redirectUrl);
+      }
+      magicLinkUrlParams.set(
+        'client',
+        BUILD_CONFIG.isElectron && appInfo ? appInfo.schema : 'web'
+      );
       await this.fetchService.fetch('/api/auth/sign-in', {
         method: 'POST',
         body: JSON.stringify({
           email,
           // we call it [callbackUrl] instead of [redirect_uri]
           // to make it clear the url is used to finish the sign-in process instead of redirect after signed-in
-          callbackUrl: `/magic-link?client=${BUILD_CONFIG.isElectron ? appInfo?.schema : 'web'}`,
+          callbackUrl: `/magic-link?${magicLinkUrlParams.toString()}`,
         }),
         headers: {
           'content-type': 'application/json',