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.

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' });

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 (
     <>

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,
   };
 

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

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;

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' });

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

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

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;
 };

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 }) => {

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>
   );

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',