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.
toeverything · Oct 25, 2024 · 8f694ac · 8f694ac
1 parent 08319bc
commit 8f694ac
Show file tree

Hide file tree

Showing 13 changed files with 135 additions and 35 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -44,14 +44,18 @@ 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 ? (
         <NoPermissionOrNotFound
           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
@@ -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
@@ -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
@@ -12,20 +12,24 @@ 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$);
   const navigate = useNavigate();
   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>
   );