fix: store user info locally

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',

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<BlocksuiteWorkspace, WebsocketProvider> = new Map();
   private _apis: Apis;
   private _channel?: WebsocketClient;
+  private _refreshToken?: string;
   // private _idbMap: Map<string, IndexedDBProvider> = new Map();
   private _workspaceLoadingQueue: Set<string> = 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<User | undefined> {
     const user = this._apis.token.user;
     return user
@@ -441,11 +434,10 @@ export class AffineProvider extends BaseProvider {
   }
 
   public override async logout(): Promise<void> {
-    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) {

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<typeof workspace, 'WorkspaceType' | 'PermissionType'> & {
-    signInWithGoogle: ReturnType<typeof getAuthorizer>[0];
-    onAuthStateChanged: ReturnType<typeof getAuthorizer>[1];
-  } & { token: typeof token };
+// See https://twitter.com/mattpocockuk/status/1622730173446557697
+// TODO: move to ts utils?
+type Prettify<T> = {
+  [K in keyof T]: T[K];
+  // eslint-disable-next-line @typescript-eslint/ban-types
+} & {};
+
+export type Apis = Prettify<
+  typeof user &
+    Omit<typeof workspace, 'WorkspaceType' | 'PermissionType'> & {
+      signInWithGoogle: ReturnType<typeof getAuthorizer>[0];
+      onAuthStateChanged: ReturnType<typeof getAuthorizer>[1];
+    } & { token: typeof token }
+>;
 
 export const getApis = (): Apis => {
   const [signInWithGoogle, onAuthStateChanged] = getAuthorizer();

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);
       },

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<LoginResponse> =>
+// 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<LoginResponse> =>
   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<LoginResponse>;
 
   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);
   }
 }
 

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;