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

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

@@ -25,13 +25,18 @@ export const createAffineDownloadProvider = (
         );
         return;
       }
-      affineApis.downloadWorkspace(id, false).then(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: () => {

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

@@ -15,6 +15,7 @@ const logger = new DebugLogger('auth-token');
 const revalidate = async () => {
   const storage = getLoginStorage();
   if (storage) {
+    try {
       const tokenMessage = parseIdToken(storage.token);
       logger.debug('revalidate affine user');
       if (isExpired(tokenMessage)) {
@@ -25,19 +26,26 @@ const revalidate = async () => {
           storageChangeSlot.emit();
         }
       }
+    } catch (e) {
+      return false;
     }
     return true;
+  } else {
+    return false;
+  }
 };
 
 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')} />}
           >
+            <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);

packages/env/src/constant.ts

@@ -1,3 +1,4 @@
+export const AFFINE_STORAGE_KEY = 'affine-local-storage-v2';
 export const DEFAULT_WORKSPACE_NAME = 'Demo Workspace';
 export const UNTITLED_WORKSPACE_NAME = 'Untitled';
 export const DEFAULT_HELLO_WORLD_PAGE_ID = 'hello-world';

packages/workspace/src/affine/api/index.ts

@@ -403,7 +403,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
           },
         })
           .then(r =>
-            r.status === 403
+            !r.ok
               ? Promise.reject(new RequestError(MessageCode.noPermission))
               : r
           )

packages/workspace/src/affine/channel.ts

@@ -4,6 +4,8 @@ import {
   isExpired,
   parseIdToken,
 } from '@affine/workspace/affine/login';
+import { WorkspaceFlavour } from '@affine/workspace/type';
+import { cleanupWorkspace } from '@affine/workspace/utils';
 import { assertExists } from '@blocksuite/global/utils';
 import * as url from 'lib0/url';
 import * as websocket from 'lib0/websocket';
@@ -28,6 +30,10 @@ export class WebsocketClient {
 
   public connect(callback: (message: any) => void) {
     const loginResponse = getLoginStorage();
+    if (!loginResponse || isExpired(parseIdToken(loginResponse.token))) {
+      cleanupWorkspace(WorkspaceFlavour.AFFINE);
+      return;
+    }
     assertExists(loginResponse, 'loginResponse is null');
     const encodedParams = url.encodeQueryParams({
       token: loginResponse.token,

packages/workspace/src/utils.ts

@@ -1,4 +1,5 @@
 import type { createWorkspaceApis } from '@affine/workspace/affine/api';
+import { rootStore, rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import { createAffineBlobStorage } from '@affine/workspace/blob';
 import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
 import type { Generator, StoreOptions } from '@blocksuite/store';
@@ -7,6 +8,12 @@ import { createIndexeddbStorage, Workspace } from '@blocksuite/store';
 import { createSQLiteStorage } from './blob/sqlite-blob-storage';
 import { WorkspaceFlavour } from './type';
 
+export function cleanupWorkspace(flavour: WorkspaceFlavour) {
+  rootStore.set(rootWorkspacesMetadataAtom, metas =>
+    metas.filter(meta => meta.flavour !== flavour)
+  );
+}
+
 const hashMap = new Map<string, Workspace>();
 
 export function createEmptyBlockSuiteWorkspace(

tests/libs/load-page.ts

@@ -2,8 +2,10 @@ import type { Page } from '@playwright/test';
 
 import { getMetas } from './utils';
 
+export const webUrl = 'http://localhost:8080';
+
 export async function openHomePage(page: Page) {
-  await page.goto('http://localhost:8080');
+  await page.goto(webUrl);
 }
 
 export async function initHomePageWithPinboard(page: Page) {

tests/libs/utils.ts

@@ -131,10 +131,6 @@ export async function loginUser(
   }, token);
 }
 
-export async function openHomePage(page: Page) {
-  return page.goto('http://localhost:8080');
-}
-
 export async function getMetas(page: Page): Promise<PageMeta[]> {
   return page.evaluate(
     () => globalThis.currentWorkspace.blockSuiteWorkspace.meta.pageMetas ?? []

tests/libs/workspace.ts

@@ -1,6 +1,9 @@
 import type { Page } from '@playwright/test';
 import { expect } from '@playwright/test';
 
+import { clickCollaborationPanel } from './setting';
+import { clickSideBarSettingButton } from './sidebar';
+
 interface CreateWorkspaceParams {
   name: string;
 }
@@ -34,3 +37,14 @@ export async function assertCurrentWorkspaceFlavour(
   const actual = await page.evaluate(() => globalThis.currentWorkspace.flavour);
   expect(actual).toBe(flavour);
 }
+
+export async function enableAffineCloudWorkspace(page: Page) {
+  await clickSideBarSettingButton(page);
+  await page.waitForTimeout(50);
+  await clickCollaborationPanel(page);
+  await page.getByTestId('local-workspace-enable-cloud-button').click();
+  await page.getByTestId('confirm-enable-cloud-button').click();
+  await page.waitForSelector("[data-testid='member-length']", {
+    timeout: 20000,
+  });
+}

tests/parallels/affine/affine-built-in-workspace.spec.ts

@@ -1,12 +1,13 @@
 import { expect } from '@playwright/test';
 
+import { openHomePage } from '../../libs/load-page';
 import { waitMarkdownImported } from '../../libs/page-logic';
 import { test } from '../../libs/playwright';
 import {
   clickNewPageButton,
   clickSideBarCurrentWorkspaceBanner,
 } from '../../libs/sidebar';
-import { getBuiltInUser, loginUser, openHomePage } from '../../libs/utils';
+import { getBuiltInUser, loginUser } from '../../libs/utils';
 
 test('collaborative', async ({ page, browser }) => {
   await openHomePage(page);

tests/parallels/affine/affine-lost-auth.spec.ts

@@ -0,0 +1,18 @@
+import { openHomePage } from '../../libs/load-page';
+import { waitMarkdownImported } from '../../libs/page-logic';
+import { test } from '../../libs/playwright';
+import { clickSideBarAllPageButton } from '../../libs/sidebar';
+import { createFakeUser, loginUser } from '../../libs/utils';
+import { enableAffineCloudWorkspace } from '../../libs/workspace';
+
+test('authorization expired', async ({ page }) => {
+  await openHomePage(page);
+  await waitMarkdownImported(page);
+  const [a] = await createFakeUser();
+  await loginUser(page, a);
+  await enableAffineCloudWorkspace(page);
+  await clickSideBarAllPageButton(page);
+  await page.evaluate(() => localStorage.removeItem('affine-login-v2'));
+  await openHomePage(page);
+  await waitMarkdownImported(page);
+});

tests/parallels/affine/affine-public-single-page.spec.ts

@@ -1,9 +1,10 @@
 import { expect } from '@playwright/test';
 
+import { openHomePage } from '../../libs/load-page';
 import { waitMarkdownImported } from '../../libs/page-logic';
 import { test } from '../../libs/playwright';
 import { clickNewPageButton } from '../../libs/sidebar';
-import { createFakeUser, loginUser, openHomePage } from '../../libs/utils';
+import { createFakeUser, loginUser } from '../../libs/utils';
 import { createWorkspace } from '../../libs/workspace';
 
 test('public single page', async ({ page, browser }) => {

tests/parallels/affine/affine-public-workspace.spec.ts

@@ -1,5 +1,6 @@
 import { expect } from '@playwright/test';
 
+import { openHomePage } from '../../libs/load-page';
 import { waitMarkdownImported } from '../../libs/page-logic';
 import { test } from '../../libs/playwright';
 import { clickPublishPanel } from '../../libs/setting';
@@ -7,7 +8,7 @@ import {
   clickSideBarAllPageButton,
   clickSideBarSettingButton,
 } from '../../libs/sidebar';
-import { createFakeUser, loginUser, openHomePage } from '../../libs/utils';
+import { createFakeUser, loginUser } from '../../libs/utils';
 import { createWorkspace } from '../../libs/workspace';
 
 test('enable public workspace', async ({ page, context }) => {

tests/parallels/affine/affine-workspace.spec.ts

@@ -1,5 +1,6 @@
 import { expect } from '@playwright/test';
 
+import { openHomePage } from '../../libs/load-page';
 import { waitMarkdownImported } from '../../libs/page-logic';
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
@@ -7,17 +8,15 @@ const userA = require('../../fixtures/userA.json');
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const userB = require('../../fixtures/userB.json');
 import { test } from '../../libs/playwright';
-import { clickCollaborationPanel } from '../../libs/setting';
 import {
   clickNewPageButton,
-  clickSideBarAllPageButton,
   clickSideBarCurrentWorkspaceBanner,
-  clickSideBarSettingButton,
 } from '../../libs/sidebar';
-import { createFakeUser, loginUser, openHomePage } from '../../libs/utils';
+import { createFakeUser, loginUser } from '../../libs/utils';
 import {
   assertCurrentWorkspaceFlavour,
   createWorkspace,
+  enableAffineCloudWorkspace,
   openWorkspaceListModal,
 } from '../../libs/workspace';
 
@@ -41,15 +40,7 @@ test('should enable affine workspace successfully', async ({ page }) => {
   const name = `test-${Date.now()}`;
   await createWorkspace({ name }, page);
   await page.waitForTimeout(50);
-  await clickSideBarSettingButton(page);
-  await page.waitForTimeout(50);
-  await clickCollaborationPanel(page);
-  await page.getByTestId('local-workspace-enable-cloud-button').click();
-  await page.getByTestId('confirm-enable-cloud-button').click();
-  await page.waitForSelector("[data-testid='member-length']", {
-    timeout: 20000,
-  });
-  await clickSideBarAllPageButton(page);
+  await enableAffineCloudWorkspace(page);
   await clickNewPageButton(page);
   await page.locator('[data-block-is-title="true"]').type('Hello, world!', {
     delay: 50,

tests/parallels/router.spec.ts

@@ -0,0 +1,25 @@
+import { expect } from '@playwright/test';
+
+import { openHomePage, webUrl } from '../libs/load-page';
+import { waitMarkdownImported } from '../libs/page-logic';
+import { test } from '../libs/playwright';
+import { clickSideBarAllPageButton } from '../libs/sidebar';
+
+test('goto not found page', async ({ page }) => {
+  await openHomePage(page);
+  await waitMarkdownImported(page);
+  const currentUrl = page.url();
+  const invalidUrl = currentUrl.replace(/\/$/, '') + '/invalid';
+  await page.goto(invalidUrl);
+  expect(await page.getByTestId('notFound').isVisible()).toBeTruthy();
+});
+
+test('goto not found workspace', async ({ page }) => {
+  await openHomePage(page);
+  await waitMarkdownImported(page);
+  await clickSideBarAllPageButton(page);
+  const currentUrl = page.url();
+  await page.goto(new URL('/workspace/invalid/all', webUrl).toString());
+  await clickSideBarAllPageButton(page);
+  expect(page.url()).toEqual(currentUrl);
+});