diff --git a/packages/data-center/src/provider/affine/__tests__/mock-apis.ts b/packages/data-center/src/provider/affine/__tests__/mock-apis.ts index b05b9cd663..21a36c7690 100644 --- a/packages/data-center/src/provider/affine/__tests__/mock-apis.ts +++ b/packages/data-center/src/provider/affine/__tests__/mock-apis.ts @@ -1,7 +1,7 @@ import type { Apis, AccessTokenMessage } from '../apis'; const user: AccessTokenMessage = { - create_at: Date.now(), + created_at: Date.now(), exp: 100000000, email: 'demo@demo.demo', id: '123', diff --git a/packages/data-center/src/provider/affine/affine.ts b/packages/data-center/src/provider/affine/affine.ts index e8e564e090..d3f1109505 100644 --- a/packages/data-center/src/provider/affine/affine.ts +++ b/packages/data-center/src/provider/affine/affine.ts @@ -5,13 +5,11 @@ import type { } from '../base'; import type { User } from '../../types'; import { Workspace as BlocksuiteWorkspace } from '@blocksuite/store'; -import { storage } from './storage.js'; import assert from 'assert'; import { WebsocketProvider } from './sync.js'; // import { IndexedDBProvider } from '../local/indexeddb'; import { getApis, Workspace } from './apis/index.js'; -import type { Apis, WorkspaceDetail, Callback } from './apis'; -import { token } from './apis/token.js'; +import type { Apis, WorkspaceDetail } from './apis'; import { WebsocketClient } from './channel'; import { loadWorkspaceUnit, @@ -40,10 +38,10 @@ const { export class AffineProvider extends BaseProvider { public id = 'affine'; - private _onTokenRefresh?: Callback = undefined; private _wsMap: Map = new Map(); private _apis: Apis; private _channel?: WebsocketClient; + private _refreshToken?: string; // private _idbMap: Map = new Map(); private _workspaceLoadingQueue: Set = new Set(); @@ -53,40 +51,25 @@ export class AffineProvider extends BaseProvider { } override async init() { - this._onTokenRefresh = () => { - if (this._apis.token.refresh) { - storage.setItem('token', this._apis.token.refresh); + this._apis.token.onChange(() => { + if (this._apis.token.isLogin) { + this._reconnectChannel(); + } else { + this._destroyChannel(); } - }; + }); - this._apis.token.onChange(this._onTokenRefresh); - - // initial login token - if (this._apis.token.isExpired) { - try { - const refreshToken = storage.getItem('token'); - if (!refreshToken) return; - await this._apis.token.refreshToken(refreshToken); - - if (this._apis.token.refresh) { - storage.set('token', this._apis.token.refresh); - } - - assert(this._apis.token.isLogin); - } catch (_) { - // this._logger('Authorization failed, fallback to local mode'); - } - } else { - storage.setItem('token', this._apis.token.refresh); - } - - if (token.isLogin) { - this._connectChannel(); + if (this._apis.token.isExpired && this._apis.token.refresh) { + // do we need to await the following? + this._apis.token.refreshToken(); } } - private _connectChannel() { - if (!this._channel) { + private _reconnectChannel() { + if (this._refreshToken !== this._apis.token.refresh) { + // need to reconnect + this._destroyChannel(); + this._channel = new WebsocketClient( `${window.location.protocol === 'https:' ? 'wss' : 'ws'}://${ window.location.host @@ -98,10 +81,21 @@ export class AffineProvider extends BaseProvider { }, } ); + + this._channel.on('message', (msg: ChannelMessage) => { + this._handlerAffineListMessage(msg); + }); + + this._refreshToken = this._apis.token.refresh; + } + } + + private _destroyChannel() { + if (this._channel) { + this._channel.disconnect(); + this._channel.destroy(); + this._channel = undefined; } - this._channel.on('message', (msg: ChannelMessage) => { - this._handlerAffineListMessage(msg); - }); } private async _handlerAffineListMessage({ @@ -111,7 +105,7 @@ export class AffineProvider extends BaseProvider { this._logger('receive server message'); const newlyCreatedWorkspaces: WorkspaceUnit[] = []; const currentWorkspaceIds = this._workspaces.list().map(w => w.id); - const newlyRemovedWorkspacecIds = currentWorkspaceIds; + const newlyRemovedWorkspaceIds = currentWorkspaceIds; for (const [id, detail] of Object.entries(ws_details)) { const { name, avatar } = metadata[id]; @@ -121,7 +115,7 @@ export class AffineProvider extends BaseProvider { const workspaceIndex = currentWorkspaceIds.indexOf(id); const ifWorkspaceExist = workspaceIndex !== -1; if (ifWorkspaceExist) { - newlyRemovedWorkspacecIds.splice(workspaceIndex, 1); + newlyRemovedWorkspaceIds.splice(workspaceIndex, 1); } /** @@ -163,7 +157,7 @@ export class AffineProvider extends BaseProvider { this._workspaces.add(newlyCreatedWorkspaces); // sync newlyRemoveWorkspaces to context - this._workspaces.remove(newlyRemovedWorkspacecIds); + this._workspaces.remove(newlyRemovedWorkspaceIds); } private _getWebsocketProvider(workspace: BlocksuiteWorkspace) { @@ -181,8 +175,8 @@ export class AffineProvider extends BaseProvider { awareness: workspace.awarenessStore.awareness, }); workspace.awarenessStore.awareness.setLocalStateField('user', { - name: token.user?.name ?? 'other', - id: Number(token.user?.id ?? -1), + name: this._apis.token.user?.name ?? 'other', + id: Number(this._apis.token.user?.id ?? -1), color: '#ffa500', }); @@ -261,23 +255,22 @@ export class AffineProvider extends BaseProvider { } override async auth() { - const refreshToken = await storage.getItem('token'); - if (refreshToken) { - await this._apis.token.refreshToken(refreshToken); + if (this._apis.token.isLogin) { + await this._apis.token.refreshToken(); if (this._apis.token.isLogin && !this._apis.token.isExpired) { // login success return; } } + const user = await this._apis.signInWithGoogle?.(); - if (!this._channel?.connected) { - this._connectChannel(); - } + if (!user) { this._sendMessage(MessageCenter.messageCode.loginError); } } + // TODO: may need to update related workspace attributes on user info change? public override async getUserInfo(): Promise { const user = this._apis.token.user; return user @@ -441,11 +434,10 @@ export class AffineProvider extends BaseProvider { } public override async logout(): Promise { - token.clear(); - this._channel?.disconnect(); + this._apis.token.clear(); + this._destroyChannel(); this._wsMap.forEach(ws => ws.disconnect()); this._workspaces.clear(false); - storage.removeItem('token'); } public override async getWorkspaceMembers(id: string) { diff --git a/packages/data-center/src/provider/affine/apis/index.ts b/packages/data-center/src/provider/affine/apis/index.ts index d3ac780887..27eebb5131 100644 --- a/packages/data-center/src/provider/affine/apis/index.ts +++ b/packages/data-center/src/provider/affine/apis/index.ts @@ -6,11 +6,20 @@ import * as user from './user.js'; import * as workspace from './workspace.js'; import { token } from './token.js'; -export type Apis = typeof user & - Omit & { - signInWithGoogle: ReturnType[0]; - onAuthStateChanged: ReturnType[1]; - } & { token: typeof token }; +// See https://twitter.com/mattpocockuk/status/1622730173446557697 +// TODO: move to ts utils? +type Prettify = { + [K in keyof T]: T[K]; + // eslint-disable-next-line @typescript-eslint/ban-types +} & {}; + +export type Apis = Prettify< + typeof user & + Omit & { + signInWithGoogle: ReturnType[0]; + onAuthStateChanged: ReturnType[1]; + } & { token: typeof token } +>; export const getApis = (): Apis => { const [signInWithGoogle, onAuthStateChanged] = getAuthorizer(); diff --git a/packages/data-center/src/provider/affine/apis/request.ts b/packages/data-center/src/provider/affine/apis/request.ts index 306ac5b9e7..29f90e19df 100644 --- a/packages/data-center/src/provider/affine/apis/request.ts +++ b/packages/data-center/src/provider/affine/apis/request.ts @@ -43,6 +43,7 @@ export const client: KyInstance = bareClient.extend({ beforeRetry: [ async ({ request }) => { + console.log('beforeRetry'); await token.refreshToken(); request.headers.set('Authorization', token.token); }, diff --git a/packages/data-center/src/provider/affine/apis/token.ts b/packages/data-center/src/provider/affine/apis/token.ts index 5b8e165929..60cc8deb39 100644 --- a/packages/data-center/src/provider/affine/apis/token.ts +++ b/packages/data-center/src/provider/affine/apis/token.ts @@ -5,9 +5,10 @@ import { decode } from 'js-base64'; import { getLogger } from '../../../logger.js'; import { bareClient } from './request.js'; +import { storage } from '../storage.js'; export interface AccessTokenMessage { - create_at: number; + created_at: number; exp: number; email: string; id: string; @@ -29,55 +30,85 @@ type LoginResponse = { refresh: string; }; -const login = (params: LoginParams): Promise => +// TODO: organize storage keys in a better way +const AFFINE_LOGIN_STORAGE_KEY = 'affine:login'; + +/** + * Use refresh token to get a new access token (JWT) + * The returned token also contains the user info payload. + */ +const doLogin = (params: LoginParams): Promise => bareClient.post('api/user/token', { json: params }).json(); export class Token { private readonly _logger; - private _accessToken!: string; - private _refreshToken!: string; + private _accessToken = ''; // idtoken (JWT) + private _refreshToken = ''; - private _user!: AccessTokenMessage | null; + private _user: AccessTokenMessage | null = null; private _padding?: Promise; constructor() { this._logger = getLogger('token'); this._logger.enabled = true; - this._setToken(); // fill with default value + this.restoreLogin(); } get user() { return this._user; } - private _setToken(login?: LoginResponse) { - this._accessToken = login?.token || ''; - this._refreshToken = login?.refresh || ''; - + setLogin(login: LoginResponse) { + this._accessToken = login.token; + this._refreshToken = login.refresh; this._user = Token.parse(this._accessToken); - if (login) { - this._logger('set login', login); - this.triggerChange(this._user); - } else { - this._logger('empty login'); + + this.triggerChange(this._user); + this.storeLogin(); + } + + private storeLogin() { + if (this.refresh) { + const { token, refresh } = this; + storage.setItem( + AFFINE_LOGIN_STORAGE_KEY, + JSON.stringify({ token, refresh }) + ); + } + } + + private restoreLogin() { + const loginStr = storage.getItem(AFFINE_LOGIN_STORAGE_KEY); + if (!loginStr) { + return; + } + try { + const login: LoginResponse = JSON.parse(loginStr); + this.setLogin(login); + } catch (err) { + this._logger('Failed to parse login info', err); } } async initToken(token: string) { - const tokens = await login({ token, type: 'Google' }); - this._setToken(tokens); + const res = await doLogin({ token, type: 'Google' }); + this.setLogin(res); return this._user; } - async refreshToken(token?: string) { + async refreshToken(refreshToken?: string) { if (!this._padding) { - this._padding = login({ + this._padding = doLogin({ type: 'Refresh', - token: this._refreshToken || token!, + token: refreshToken || this._refreshToken, }); + this._refreshToken = refreshToken || this._refreshToken; + } + const res = await this._padding; + if (!refreshToken || refreshToken !== this._refreshToken) { + this.setLogin(res); } - this._setToken(await this._padding); this._padding = undefined; } @@ -95,7 +126,7 @@ export class Token { get isExpired() { if (!this._user) return true; - return Date.now() - this._user.create_at > this._user.exp; + return Date.now() > this._user.exp; } static parse(token: string): AccessTokenMessage | null { @@ -128,7 +159,9 @@ export class Token { } clear() { - this._setToken(); + this._accessToken = ''; + this._refreshToken = ''; + storage.removeItem(AFFINE_LOGIN_STORAGE_KEY); } } diff --git a/packages/data-center/src/provider/affine/storage.ts b/packages/data-center/src/provider/affine/storage.ts index 4ac1b74697..95031de655 100644 --- a/packages/data-center/src/provider/affine/storage.ts +++ b/packages/data-center/src/provider/affine/storage.ts @@ -1 +1,3 @@ -export { varStorage as storage } from 'lib0/storage'; +import { varStorage } from 'lib0/storage'; + +export const storage = varStorage as Storage;