fix: revalidate user token with no refresh page (#1842)

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

@@ -4,6 +4,7 @@ import {
   isExpired,
   parseIdToken,
   setLoginStorage,
+  storageChangeSlot,
 } from '@affine/workspace/affine/login';
 import useSWR from 'swr';
 
@@ -21,10 +22,7 @@ const revalidate = async () => {
       const response = await affineAuth.refreshToken(storage);
       if (response) {
         setLoginStorage(response);
-
-        // todo: need to notify the app that the token has been refreshed
-        //  this is a hack to force a reload
-        window.location.reload();
+        storageChangeSlot.emit();
       }
     }
   }
@@ -38,5 +36,8 @@ export function useAffineRefreshAuthToken(
   useSWR('autoRefreshToken', {
     fetcher: revalidate,
     refreshInterval,
+    revalidateOnFocus: true,
+    revalidateOnReconnect: true,
+    revalidateOnMount: true,
   });
 }

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

@@ -40,10 +40,11 @@ describe('AFFiNE workspace', () => {
       // but refresh is still valid
       refresh: data.refresh,
     });
-    renderHook(() => useAffineRefreshAuthToken(1));
+    const hook = renderHook(() => useAffineRefreshAuthToken(1));
     await new Promise(resolve => setTimeout(resolve, 3000));
     const userData = parseIdToken(getLoginStorage()?.token as string);
     expect(userData).not.toBeNull();
     expect(isExpired(userData)).toBe(false);
+    hook.unmount();
   });
 });

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

@@ -1,8 +1,7 @@
 import { MessageCode, Messages } from '@affine/env';
-import { assertExists } from '@blocksuite/global/utils';
 import { z } from 'zod';
 
-import { getLoginStorage } from '../login';
+import { checkLoginStorage } from '../login';
 
 export class RequestError extends Error {
   public readonly code: MessageCode;
@@ -51,8 +50,7 @@ export type UsageResponse = z.infer<typeof usageResponseSchema>;
 export function createUserApis(prefixUrl = '/') {
   return {
     getUsage: async (): Promise<UsageResponse> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + 'api/resource/usage', {
         method: 'GET',
         headers: {
@@ -63,8 +61,7 @@ export function createUserApis(prefixUrl = '/') {
     getUserByEmail: async (
       params: GetUserByEmailParams
     ): Promise<User[] | null> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       const target = new URL(prefixUrl + 'api/user', window.location.origin);
       target.searchParams.append('email', params.email);
       target.searchParams.append('workspace_id', params.workspace_id);
@@ -187,8 +184,7 @@ export const createWorkspaceResponseSchema = z.object({
 export function createWorkspaceApis(prefixUrl = '/') {
   return {
     getWorkspaces: async (): Promise<Workspace[]> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + 'api/workspace', {
         method: 'GET',
         headers: {
@@ -204,8 +200,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
     getWorkspaceDetail: async (
       params: GetWorkspaceDetailParams
     ): Promise<WorkspaceDetail | null> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + `api/workspace/${params.id}`, {
         method: 'GET',
         headers: {
@@ -220,8 +215,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
     getWorkspaceMembers: async (
       params: GetWorkspaceDetailParams
     ): Promise<Member[]> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + `api/workspace/${params.id}/permission`, {
         method: 'GET',
         headers: {
@@ -236,8 +230,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
     createWorkspace: async (
       encodedYDoc: ArrayBuffer
     ): Promise<{ id: string }> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage();
       return fetch(prefixUrl + 'api/workspace', {
         method: 'POST',
         body: encodedYDoc,
@@ -254,8 +247,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
     updateWorkspace: async (
       params: UpdateWorkspaceParams
     ): Promise<{ public: boolean | null }> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + `api/workspace/${params.id}`, {
         method: 'POST',
         body: JSON.stringify({
@@ -274,8 +266,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
     deleteWorkspace: async (
       params: DeleteWorkspaceParams
     ): Promise<boolean> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + `api/workspace/${params.id}`, {
         method: 'DELETE',
         headers: {
@@ -292,8 +283,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
      * Notice: Only support normal(contrast to private) workspace.
      */
     inviteMember: async (params: InviteMemberParams): Promise<void> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + `api/workspace/${params.id}/permission`, {
         method: 'POST',
         body: JSON.stringify({
@@ -310,8 +300,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
         });
     },
     removeMember: async (params: RemoveMemberParams): Promise<void> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + `api/permission/${params.permissionId}`, {
         method: 'DELETE',
         headers: {
@@ -339,8 +328,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
       arrayBuffer: ArrayBuffer,
       type: string
     ): Promise<string> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       const mb = arrayBuffer.byteLength / 1048576;
       if (mb > 10) {
         throw new RequestError(MessageCode.blobTooLarge);
@@ -358,8 +346,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
       workspaceId: string,
       blobId: string
     ): Promise<ArrayBuffer> => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + `api/workspace/${workspaceId}/blob/${blobId}`, {
         method: 'GET',
         headers: {
@@ -372,8 +359,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
         });
     },
     leaveWorkspace: async ({ id }: LeaveWorkspaceParams) => {
-      const auth = getLoginStorage();
-      assertExists(auth);
+      const auth = await checkLoginStorage(prefixUrl);
       return fetch(prefixUrl + `api/workspace/${id}/permission`, {
         method: 'DELETE',
         headers: {
@@ -405,8 +391,7 @@ export function createWorkspaceApis(prefixUrl = '/') {
           method: 'GET',
         }).then(r => r.arrayBuffer());
       } else {
-        const auth = getLoginStorage();
-        assertExists(auth);
+        const auth = await checkLoginStorage(prefixUrl);
         return fetch(prefixUrl + `api/workspace/${workspaceId}/doc`, {
           method: 'GET',
           headers: {

packages/workspace/src/affine/login.ts

@@ -1,4 +1,6 @@
 import { DebugLogger } from '@affine/debug';
+import { assertExists } from '@blocksuite/global/utils';
+import { Slot } from '@blocksuite/store';
 import { initializeApp } from 'firebase/app';
 import type { AuthProvider } from 'firebase/auth';
 import {
@@ -78,6 +80,32 @@ export const getLoginStorage = (): LoginResponse | null => {
   return null;
 };
 
+export const storageChangeSlot = new Slot();
+
+export const checkLoginStorage = async (
+  prefixUrl = '/'
+): Promise<LoginResponse> => {
+  const storage = getLoginStorage();
+  assertExists(storage, 'Login token is not set');
+  if (isExpired(parseIdToken(storage.token), 0)) {
+    logger.debug('refresh token needed');
+    const response: LoginResponse = await fetch(prefixUrl + 'api/user/token', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({
+        type: 'Refresh',
+        token: storage.refresh,
+      }),
+    }).then(r => r.json());
+    setLoginStorage(response);
+    logger.debug('refresh token emit');
+    storageChangeSlot.emit();
+  }
+  return getLoginStorage() as LoginResponse;
+};
+
 export const enum SignMethod {
   Google = 'Google',
   GitHub = 'GitHub',

packages/workspace/src/affine/sync.ts

@@ -1,14 +1,19 @@
+import { DebugLogger } from '@affine/debug';
 import {
   workspaceDetailSchema,
   workspaceSchema,
 } from '@affine/workspace/affine/api';
 import { WebsocketClient } from '@affine/workspace/affine/channel';
+import { storageChangeSlot } from '@affine/workspace/affine/login';
 import { jotaiStore, jotaiWorkspacesAtom } from '@affine/workspace/atom';
 import type { WorkspaceCRUD } from '@affine/workspace/type';
 import type { WorkspaceFlavour } from '@affine/workspace/type';
 import { assertExists } from '@blocksuite/global/utils';
+import type { Disposable } from '@blocksuite/store';
 import { z } from 'zod';
 
+const logger = new DebugLogger('affine-sync');
+
 const channelMessageSchema = z.object({
   ws_list: z.array(workspaceSchema),
   ws_details: z.record(workspaceDetailSchema),
@@ -28,7 +33,7 @@ export function createAffineGlobalChannel(
   let client: WebsocketClient | null;
 
   async function handleMessage(channelMessage: ChannelMessage) {
-    console.log('channelMessage', channelMessage);
+    logger.debug('channelMessage', channelMessage);
     const parseResult = channelMessageSchema.safeParse(channelMessage);
     if (!parseResult.success) {
       console.error(
@@ -53,8 +58,8 @@ export function createAffineGlobalChannel(
       }
     }
   }
-
-  return {
+  let dispose: Disposable | undefined = undefined;
+  const apis = {
     connect: () => {
       client = new WebsocketClient(
         `${window.location.protocol === 'https:' ? 'wss' : 'ws'}://${
@@ -62,11 +67,18 @@ export function createAffineGlobalChannel(
         }/api/global/sync`
       );
       client.connect(handleMessage);
+      dispose = storageChangeSlot.on(() => {
+        apis.disconnect();
+        apis.connect();
+      });
     },
     disconnect: () => {
       assertExists(client, 'client is null');
       client.disconnect();
+      dispose?.dispose();
       client = null;
     },
   };
+
+  return apis;
 }

packages/workspace/src/providers/index.ts

@@ -1,12 +1,18 @@
 import { config } from '@affine/env';
 import { KeckProvider } from '@affine/workspace/affine/keck';
-import { getLoginStorage } from '@affine/workspace/affine/login';
+import {
+  getLoginStorage,
+  storageChangeSlot,
+} from '@affine/workspace/affine/login';
 import type { Provider } from '@affine/workspace/type';
 import type {
   AffineWebSocketProvider,
   LocalIndexedDBProvider,
 } from '@affine/workspace/type';
-import type { Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
+import type {
+  Disposable,
+  Workspace as BlockSuiteWorkspace,
+} from '@blocksuite/store';
 import { assertExists } from '@blocksuite/store';
 import {
   createIndexedDBProvider as create,
@@ -20,15 +26,21 @@ const createAffineWebSocketProvider = (
   blockSuiteWorkspace: BlockSuiteWorkspace
 ): AffineWebSocketProvider => {
   let webSocketProvider: KeckProvider | null = null;
-  return {
+  let dispose: Disposable | undefined = undefined;
+  const apis: AffineWebSocketProvider = {
     flavour: 'affine-websocket',
     background: false,
     cleanup: () => {
       assertExists(webSocketProvider);
       webSocketProvider.destroy();
       webSocketProvider = null;
+      dispose?.dispose();
     },
     connect: () => {
+      dispose = storageChangeSlot.on(() => {
+        apis.disconnect();
+        apis.connect();
+      });
       const wsUrl = `${
         window.location.protocol === 'https:' ? 'wss' : 'ws'
       }://${window.location.host}/api/sync/`;
@@ -53,8 +65,11 @@ const createAffineWebSocketProvider = (
       localProviderLogger.info('disconnect', webSocketProvider.url);
       webSocketProvider.destroy();
       webSocketProvider = null;
+      dispose?.dispose();
     },
   };
+
+  return apis;
 };
 
 const createIndexedDBProvider = (