From 20e56cc474f06c97e8f976a2e14e1b2f1fbda8b6 Mon Sep 17 00:00:00 2001 From: Himself65 Date: Fri, 7 Apr 2023 17:51:51 -0500 Subject: [PATCH] fix: revalidate user token with no refresh page (#1842) --- .../affine/use-affine-refresh-auth-token.ts | 9 ++-- .../plugins/affine/__tests__/index.spec.tsx | 3 +- packages/workspace/src/affine/api/index.ts | 45 +++++++------------ packages/workspace/src/affine/login.ts | 28 ++++++++++++ packages/workspace/src/affine/sync.ts | 18 ++++++-- packages/workspace/src/providers/index.ts | 21 +++++++-- 6 files changed, 83 insertions(+), 41 deletions(-) diff --git a/apps/web/src/hooks/affine/use-affine-refresh-auth-token.ts b/apps/web/src/hooks/affine/use-affine-refresh-auth-token.ts index 80a0674c06..7140d56196 100644 --- a/apps/web/src/hooks/affine/use-affine-refresh-auth-token.ts +++ b/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, }); } diff --git a/apps/web/src/plugins/affine/__tests__/index.spec.tsx b/apps/web/src/plugins/affine/__tests__/index.spec.tsx index 52e224b5be..b81278e459 100644 --- a/apps/web/src/plugins/affine/__tests__/index.spec.tsx +++ b/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(); }); }); diff --git a/packages/workspace/src/affine/api/index.ts b/packages/workspace/src/affine/api/index.ts index 028bfb77dd..e291c2cb0e 100644 --- a/packages/workspace/src/affine/api/index.ts +++ b/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; export function createUserApis(prefixUrl = '/') { return { getUsage: async (): Promise => { - 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 => { - 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 => { - 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 => { - 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 => { - 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 => { - 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 => { - 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 => { - 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 => { - 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 => { - 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: { diff --git a/packages/workspace/src/affine/login.ts b/packages/workspace/src/affine/login.ts index 855b9b000f..445dde1db2 100644 --- a/packages/workspace/src/affine/login.ts +++ b/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 => { + 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', diff --git a/packages/workspace/src/affine/sync.ts b/packages/workspace/src/affine/sync.ts index 9ff58b228b..02cedebca1 100644 --- a/packages/workspace/src/affine/sync.ts +++ b/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; } diff --git a/packages/workspace/src/providers/index.ts b/packages/workspace/src/providers/index.ts index 9378505657..10402a5c0a 100644 --- a/packages/workspace/src/providers/index.ts +++ b/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 = (