fix(workspace): check affine login auth (#2070)

apps/web/src/blocksuite/providers/affine/index.ts

@@ -25,14 +25,19 @@ export const createAffineDownloadProvider = (
         );
         return;
       }
-      affineApis.downloadWorkspace(id, false).then(binary => {
-        hashMap.set(id, binary);
-        providerLogger.debug('applyUpdate');
-        BlockSuiteWorkspace.Y.applyUpdate(
-          blockSuiteWorkspace.doc,
-          new Uint8Array(binary)
-        );
-      });
+      affineApis
+        .downloadWorkspace(id, false)
+        .then(binary => {
+          hashMap.set(id, binary);
+          providerLogger.debug('applyUpdate');
+          BlockSuiteWorkspace.Y.applyUpdate(
+            blockSuiteWorkspace.doc,
+            new Uint8Array(binary)
+          );
+        })
+        .catch(e => {
+          providerLogger.error('downloadWorkspace', e);
+        });
     },
     disconnect: () => {
       providerLogger.info('disconnect download provider', id);

apps/web/src/hooks/affine/use-affine-refresh-auth-token.ts

@@ -15,29 +15,37 @@ const logger = new DebugLogger('auth-token');
 const revalidate = async () => {
   const storage = getLoginStorage();
   if (storage) {
-    const tokenMessage = parseIdToken(storage.token);
-    logger.debug('revalidate affine user');
-    if (isExpired(tokenMessage)) {
-      logger.debug('need to refresh token');
-      const response = await affineAuth.refreshToken(storage);
-      if (response) {
-        setLoginStorage(response);
-        storageChangeSlot.emit();
+    try {
+      const tokenMessage = parseIdToken(storage.token);
+      logger.debug('revalidate affine user');
+      if (isExpired(tokenMessage)) {
+        logger.debug('need to refresh token');
+        const response = await affineAuth.refreshToken(storage);
+        if (response) {
+          setLoginStorage(response);
+          storageChangeSlot.emit();
+        }
       }
+    } catch (e) {
+      return false;
     }
+    return true;
+  } else {
+    return false;
   }
-  return true;
 };
 
 export function useAffineRefreshAuthToken(
   // every 30 seconds, check if the token is expired
   refreshInterval = 30 * 1000
 ) {
-  useSWR('autoRefreshToken', {
+  const { data } = useSWR<boolean>('autoRefreshToken', {
+    suspense: true,
     fetcher: revalidate,
     refreshInterval,
     revalidateOnFocus: true,
     revalidateOnReconnect: true,
     revalidateOnMount: true,
   });
+  return data;
 }

apps/web/src/hooks/use-sync-router-with-current-workspace-id.ts

@@ -20,6 +20,19 @@ export function useSyncRouterWithCurrentWorkspaceId(router: NextRouter) {
       return;
     }
     if (currentWorkspaceId) {
+      if (currentWorkspaceId !== workspaceId) {
+        const target = metadata.find(workspace => workspace.id === workspaceId);
+        if (!target) {
+          // workspaceId is invalid, redirect to currentWorkspaceId
+          void router.push({
+            pathname: router.pathname,
+            query: {
+              ...router.query,
+              workspaceId: currentWorkspaceId,
+            },
+          });
+        }
+      }
       return;
     }
     const targetWorkspace = metadata.find(

apps/web/src/layouts/workspace-layout.tsx

@@ -17,7 +17,14 @@ import { useAtom, useAtomValue, useSetAtom } from 'jotai';
 import Head from 'next/head';
 import { useRouter } from 'next/router';
 import type { FC, PropsWithChildren, ReactElement } from 'react';
-import { lazy, Suspense, useCallback, useEffect, useState } from 'react';
+import {
+  lazy,
+  Suspense,
+  useCallback,
+  useEffect,
+  useMemo,
+  useState,
+} from 'react';
 
 import { openQuickSearchModalAtom, openWorkspacesModalAtom } from '../atoms';
 import {
@@ -176,6 +183,10 @@ export const WorkspaceLayout: FC<PropsWithChildren> =
     }, [i18n]);
     const currentWorkspaceId = useAtomValue(rootCurrentWorkspaceIdAtom);
     const jotaiWorkspaces = useAtomValue(rootWorkspacesMetadataAtom);
+    const meta = useMemo(
+      () => jotaiWorkspaces.find(x => x.id === currentWorkspaceId),
+      [currentWorkspaceId, jotaiWorkspaces]
+    );
     const set = useSetAtom(rootWorkspacesMetadataAtom);
     useEffect(() => {
       logger.info('mount');
@@ -228,6 +239,9 @@ export const WorkspaceLayout: FC<PropsWithChildren> =
         };
       }
     }, [currentWorkspaceId, jotaiWorkspaces]);
+
+    const Provider =
+      (meta && WorkspacePlugins[meta.flavour].UI.Provider) ?? DefaultProvider;
     return (
       <>
         {/* fixme(himself65): don't re-render whole modals */}
@@ -240,7 +254,9 @@ export const WorkspaceLayout: FC<PropsWithChildren> =
           <Suspense
             fallback={<PageLoading text={t('Finding Current Workspace')} />}
           >
-            <WorkspaceLayoutInner>{children}</WorkspaceLayoutInner>
+            <Provider>
+              <WorkspaceLayoutInner>{children}</WorkspaceLayoutInner>
+            </Provider>
           </Suspense>
         </CurrentWorkspaceContext>
       </>
@@ -397,11 +413,8 @@ export const WorkspaceLayoutInner: FC<PropsWithChildren> = ({ children }) => {
     );
   }, [setIsResizing, setSidebarOpen, setSliderWidth]);
 
-  const Provider =
-    WorkspacePlugins[currentWorkspace.flavour].UI.Provider ?? DefaultProvider;
-
   return (
-    <Provider>
+    <>
       <Head>
         <title>{title}</title>
       </Head>
@@ -465,6 +478,6 @@ export const WorkspaceLayoutInner: FC<PropsWithChildren> = ({ children }) => {
         </MainContainerWrapper>
       </StyledPage>
       <QuickSearch />
-    </Provider>
+    </>
   );
 };

apps/web/src/pages/index.tsx

@@ -1,6 +1,5 @@
 import { DebugLogger } from '@affine/debug';
 import { useTranslation } from '@affine/i18n';
-import { WorkspaceFlavour } from '@affine/workspace/type';
 import type { NextPage } from 'next';
 import { useRouter } from 'next/router';
 import { Suspense, useEffect } from 'react';
@@ -27,10 +26,7 @@ const IndexPageInner = () => {
     const targetWorkspace =
       (lastWorkspaceId &&
         workspaces.find(({ id }) => id === lastWorkspaceId)) ||
-      // fixme(himself65):
-      //  when affine workspace is expired and the first workspace is affine,
-      //  the page will crash
-      workspaces.find(({ flavour }) => flavour === WorkspaceFlavour.LOCAL);
+      workspaces.at(0);
 
     if (targetWorkspace) {
       const pageId =

apps/web/src/plugins/affine/index.tsx

@@ -1,10 +1,11 @@
-import { config, prefixUrl } from '@affine/env';
+import { AFFINE_STORAGE_KEY, config, prefixUrl } from '@affine/env';
 import { initPage } from '@affine/env/blocksuite';
 import { currentAffineUserAtom } from '@affine/workspace/affine/atom';
 import {
   clearLoginStorage,
   createAffineAuth,
   getLoginStorage,
+  isExpired,
   parseIdToken,
   setLoginStorage,
   SignMethod,
@@ -12,9 +13,13 @@ import {
 import { rootStore, rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import type { AffineWorkspace } from '@affine/workspace/type';
 import { LoadPriority, WorkspaceFlavour } from '@affine/workspace/type';
-import { createEmptyBlockSuiteWorkspace } from '@affine/workspace/utils';
+import {
+  cleanupWorkspace,
+  createEmptyBlockSuiteWorkspace,
+} from '@affine/workspace/utils';
 import { createJSONStorage } from 'jotai/utils';
-import React from 'react';
+import type { PropsWithChildren, ReactElement } from 'react';
+import { Suspense, useEffect } from 'react';
 import { mutate } from 'swr';
 import { z } from 'zod';
 
@@ -23,8 +28,8 @@ import { PageNotFoundError } from '../../components/affine/affine-error-eoundary
 import { WorkspaceSettingDetail } from '../../components/affine/workspace-setting-detail';
 import { BlockSuitePageList } from '../../components/blocksuite/block-suite-page-list';
 import { PageDetailEditor } from '../../components/page-detail-editor';
+import { PageLoading } from '../../components/pure/loading';
 import { useAffineRefreshAuthToken } from '../../hooks/affine/use-affine-refresh-auth-token';
-import { AffineSWRConfigProvider } from '../../providers/AffineSWRConfigProvider';
 import { BlockSuiteWorkspace } from '../../shared';
 import { affineApis } from '../../shared/apis';
 import { toast } from '../../utils';
@@ -32,7 +37,6 @@ import type { WorkspacePlugin } from '..';
 import { QueryKey } from './fetcher';
 
 const storage = createJSONStorage(() => localStorage);
-const kAffineLocal = 'affine-local-storage-v2';
 const schema = z.object({
   id: z.string(),
   type: z.number(),
@@ -41,7 +45,7 @@ const schema = z.object({
 });
 
 const getPersistenceAllWorkspace = () => {
-  const items = storage.getItem(kAffineLocal);
+  const items = storage.getItem(AFFINE_STORAGE_KEY);
   const allWorkspaces: AffineWorkspace[] = [];
   if (
     Array.isArray(items) &&
@@ -71,6 +75,22 @@ const getPersistenceAllWorkspace = () => {
 
 export const affineAuth = createAffineAuth(prefixUrl);
 
+function AuthContext({ children }: PropsWithChildren): ReactElement {
+  const login = useAffineRefreshAuthToken();
+
+  useEffect(() => {
+    if (!login) {
+      console.warn('No login, redirecting to local workspace page...');
+    }
+  }, [login]);
+  if (!login) {
+    return (
+      <PageLoading text="No login, redirecting to local workspace page..." />
+    );
+  }
+  return <>{children}</>;
+}
+
 export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
   flavour: WorkspaceFlavour.AFFINE,
   loadPriority: LoadPriority.HIGH,
@@ -99,7 +119,7 @@ export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
           workspace => workspace.flavour !== WorkspaceFlavour.AFFINE
         )
       );
-      storage.removeItem(kAffineLocal);
+      storage.removeItem(AFFINE_STORAGE_KEY);
       clearLoginStorage();
       rootStore.set(currentAffineUserAtom, null);
     },
@@ -132,13 +152,13 @@ export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
       return id;
     },
     delete: async workspace => {
-      const items = storage.getItem(kAffineLocal);
+      const items = storage.getItem(AFFINE_STORAGE_KEY);
       if (
         Array.isArray(items) &&
         items.every(item => schema.safeParse(item).success)
       ) {
         storage.setItem(
-          kAffineLocal,
+          AFFINE_STORAGE_KEY,
           items.filter(item => item.id !== workspace.id)
         );
       }
@@ -148,12 +168,17 @@ export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
       await mutate(matcher => matcher === QueryKey.getWorkspaces);
     },
     get: async workspaceId => {
+      // fixme(himself65): rewrite the auth logic
       try {
-        if (!getLoginStorage()) {
-          const workspaces = getPersistenceAllWorkspace();
-          return (
-            workspaces.find(workspace => workspace.id === workspaceId) ?? null
-          );
+        const loginStorage = getLoginStorage();
+        if (
+          loginStorage == null ||
+          isExpired(parseIdToken(loginStorage.token))
+        ) {
+          rootStore.set(currentAffineUserAtom, null);
+          storage.removeItem(AFFINE_STORAGE_KEY);
+          cleanupWorkspace(WorkspaceFlavour.AFFINE);
+          return null;
         }
         const workspaces: AffineWorkspace[] = await AffinePlugin.CRUD.list();
         return (
@@ -168,8 +193,20 @@ export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
     },
     list: async () => {
       const allWorkspaces = getPersistenceAllWorkspace();
-      if (!getLoginStorage()) {
-        return allWorkspaces;
+      const loginStorage = getLoginStorage();
+      // fixme(himself65): rewrite the auth logic
+      try {
+        if (
+          loginStorage == null ||
+          isExpired(parseIdToken(loginStorage.token))
+        ) {
+          rootStore.set(currentAffineUserAtom, null);
+          storage.removeItem(AFFINE_STORAGE_KEY);
+          return [];
+        }
+      } catch (e) {
+        storage.removeItem(AFFINE_STORAGE_KEY);
+        return [];
       }
       try {
         const workspaces = await affineApis.getWorkspaces().then(workspaces => {
@@ -189,7 +226,7 @@ export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
                 permission: workspace.permission,
               } satisfies z.infer<typeof schema>;
             });
-            const old = storage.getItem(kAffineLocal);
+            const old = storage.getItem(AFFINE_STORAGE_KEY);
             if (
               Array.isArray(old) &&
               old.every(item => schema.safeParse(item).success)
@@ -201,7 +238,7 @@ export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
                   data.push(item);
                 }
               });
-              storage.setItem(kAffineLocal, [...data]);
+              storage.setItem(AFFINE_STORAGE_KEY, [...data]);
             }
 
             const affineWorkspace: AffineWorkspace = {
@@ -231,7 +268,7 @@ export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
             permission: workspace.permission,
           } satisfies z.infer<typeof schema>;
         });
-        storage.setItem(kAffineLocal, [...dump]);
+        storage.setItem(AFFINE_STORAGE_KEY, [...dump]);
       } catch (e) {
         console.error('fetch affine workspaces failed', e);
       }
@@ -240,8 +277,11 @@ export const AffinePlugin: WorkspacePlugin<WorkspaceFlavour.AFFINE> = {
   },
   UI: {
     Provider: ({ children }) => {
-      useAffineRefreshAuthToken();
-      return <AffineSWRConfigProvider>{children}</AffineSWRConfigProvider>;
+      return (
+        <Suspense fallback={<PageLoading text="Checking login status..." />}>
+          <AuthContext>{children}</AuthContext>
+        </Suspense>
+      );
     },
     PageDetail: ({ currentWorkspace, currentPageId }) => {
       const page = currentWorkspace.blockSuiteWorkspace.getPage(currentPageId);