store google_token_access in memory (#4238)

rrrooommmaaa · web-flow · commit daf4f811dbc4 · 2022-01-13T23:12:10.000+03:00
diff --git a/extension/chrome/elements/pgp_block.ts b/extension/chrome/elements/pgp_block.ts
@@ -86,7 +86,7 @@ export class PgpBlockView extends View {
   };
 
   public render = async () => {
-    const storage = await AcctStore.get(this.acctEmail, ['setup_done', 'google_token_scopes']);
+    const storage = await AcctStore.get(this.acctEmail, ['setup_done']);
     this.orgRules = await OrgRules.newInstance(this.acctEmail);
     this.pubLookup = new PubLookup(this.orgRules);
     const scopes = await AcctStore.getScopes(this.acctEmail);
diff --git a/extension/chrome/settings/modules/experimental.ts b/extension/chrome/settings/modules/experimental.ts
@@ -14,6 +14,8 @@ import { Url } from '../../../js/common/core/common.js';
 import { View } from '../../../js/common/view.js';
 import { AcctStore } from '../../../js/common/platform/store/acct-store.js';
 import { Api } from '../../../js/common/api/shared/api.js';
+import { InMemoryStore } from '../../../js/common/platform/store/in-memory-store.js';
+import { InMemoryStoreKeys } from '../../../js/common/core/const.js';
 
 View.run(class ExperimentalView extends View {
 
@@ -82,7 +84,7 @@ View.run(class ExperimentalView extends View {
   };
 
   private makeGoogleAuthTokenUnusableHandler = async () => {
-    await AcctStore.set(this.acctEmail, { google_token_access: 'flowcrypt_test_bad_access_token' });
+    await InMemoryStore.set(this.acctEmail, InMemoryStoreKeys.GOOGLE_TOKEN_ACCESS, 'flowcrypt_test_bad_access_token');
     BrowserMsg.send.reload(this.parentTabId, {});
   };
 
diff --git a/extension/js/background_page/background_page.ts b/extension/js/background_page/background_page.ts
@@ -55,7 +55,7 @@ opgp.initWorker({ path: '/lib/openpgp.worker.js' });
 
   // storage related handlers
   BrowserMsg.bgAddListener('db', (r: Bm.Db) => BgHandlers.dbOperationHandler(db, r));
-  BrowserMsg.bgAddListener('inMemoryStoreSet', async (r: Bm.InMemoryStoreSet) => inMemoryStore.set(emailKeyIndex(r.acctEmail, r.key), r.value));
+  BrowserMsg.bgAddListener('inMemoryStoreSet', async (r: Bm.InMemoryStoreSet) => inMemoryStore.set(emailKeyIndex(r.acctEmail, r.key), r.value, r.expiration));
   BrowserMsg.bgAddListener('inMemoryStoreGet', async (r: Bm.InMemoryStoreGet) => inMemoryStore.get(emailKeyIndex(r.acctEmail, r.key)));
   BrowserMsg.bgAddListener('storeGlobalGet', (r: Bm.StoreGlobalGet) => GlobalStore.get(r.keys));
   BrowserMsg.bgAddListener('storeGlobalSet', (r: Bm.StoreGlobalSet) => GlobalStore.set(r.values));
diff --git a/extension/js/common/api/account-servers/enterprise-server.ts b/extension/js/common/api/account-servers/enterprise-server.ts
@@ -11,7 +11,7 @@ import { BackendRes, ProfileUpdate } from './flowcrypt-com-api.js';
 import { Dict } from '../../core/common.js';
 import { ErrorReport, UnreportableError } from '../../platform/catch.js';
 import { ApiErr, BackendAuthErr } from '../shared/api-error.js';
-import { FLAVOR } from '../../core/const.js';
+import { FLAVOR, InMemoryStoreKeys } from '../../core/const.js';
 import { Attachment } from '../../core/attachment.js';
 import { Recipients } from '../email-provider/email-provider-api.js';
 import { Buf } from '../../core/buf.js';
@@ -141,7 +141,7 @@ export class EnterpriseServer extends Api {
   };
 
   private authHdr = async (): Promise<Dict<string>> => {
-    const idToken = await InMemoryStore.get(this.acctEmail, InMemoryStore.ID_TOKEN_STORAGE_KEY);
+    const idToken = await InMemoryStore.get(this.acctEmail, InMemoryStoreKeys.ID_TOKEN);
     if (idToken) {
       return { Authorization: `Bearer ${idToken}` };
     }
diff --git a/extension/js/common/api/email-provider/gmail/google-auth.ts b/extension/js/common/api/email-provider/gmail/google-auth.ts
@@ -21,6 +21,7 @@ import { AcctStore, AcctStoreDict } from '../../../platform/store/acct-store.js'
 import { AccountServer } from '../../account-server.js';
 import { EnterpriseServer } from '../../account-servers/enterprise-server.js';
 import { InMemoryStore } from '../../../platform/store/in-memory-store.js';
+import { InMemoryStoreKeys } from '../../../core/const.js';
 
 type GoogleAuthTokenInfo = { issued_to: string, audience: string, scope: string, expires_in: number, access_type: 'offline' };
 type GoogleAuthTokensResponse = { access_token: string, expires_in: number, refresh_token?: string, id_token: string, token_type: 'Bearer' };
@@ -80,22 +81,27 @@ export class GoogleAuth {
     if (!acctEmail) {
       throw new Error('missing account_email in api_gmail_call');
     }
-    const storage = await AcctStore.get(acctEmail, ['google_token_access', 'google_token_expires', 'google_token_scopes', 'google_token_refresh']);
-    if (!storage.google_token_access || !storage.google_token_refresh) {
+    const storage = await AcctStore.get(acctEmail, ['google_token_scopes', 'google_token_refresh']);
+    if (!storage.google_token_refresh) {
       throw new GoogleAuthErr(`Account ${acctEmail} not connected to FlowCrypt Browser Extension`);
-    } else if (GoogleAuth.googleApiIsAuthTokenValid(storage) && !forceRefresh) {
-      return `Bearer ${storage.google_token_access}`;
-    } else { // refresh token
-      const refreshTokenRes = await GoogleAuth.googleAuthRefreshToken(storage.google_token_refresh);
+    }
+    if (!forceRefresh) {
+      const googleAccessToken = await InMemoryStore.get(acctEmail, InMemoryStoreKeys.GOOGLE_TOKEN_ACCESS);
+      if (googleAccessToken) {
+        return `Bearer ${googleAccessToken}`;
+      }
+    }
+    // refresh token
+    const refreshTokenRes = await GoogleAuth.googleAuthRefreshToken(storage.google_token_refresh);
+    if (refreshTokenRes.access_token) {
       await GoogleAuth.googleAuthCheckAccessToken(refreshTokenRes.access_token); // https://groups.google.com/forum/#!topic/oauth2-dev/QOFZ4G7Ktzg
       await GoogleAuth.googleAuthSaveTokens(acctEmail, refreshTokenRes, storage.google_token_scopes || []);
-      const auth = await AcctStore.get(acctEmail, ['google_token_access', 'google_token_expires']);
-      if (GoogleAuth.googleApiIsAuthTokenValid(auth)) { // have a valid gmail_api oauth token
-        return `Bearer ${auth.google_token_access}`;
-      } else {
-        throw new GoogleAuthErr(`Could not refresh google auth token - did not become valid (access:${!!auth.google_token_access},expires:${auth.google_token_expires},now:${Date.now()})`);
+      const googleAccessToken = await InMemoryStore.get(acctEmail, InMemoryStoreKeys.GOOGLE_TOKEN_ACCESS);
+      if (googleAccessToken) {
+        return `Bearer ${googleAccessToken}`;
       }
     }
+    throw new GoogleAuthErr(`Could not refresh google auth token - did not become valid (access:${refreshTokenRes.access_token},expires_in:${refreshTokenRes.expires_in},now:${Date.now()})`);
   };
 
   public static apiGoogleCallRetryAuthErrorOneTime = async (acctEmail: string, request: JQuery.AjaxSettings): Promise<any> => {
@@ -282,9 +288,8 @@ export class GoogleAuth {
   private static googleAuthSaveTokens = async (acctEmail: string, tokensObj: GoogleAuthTokensResponse, scopes: string[]) => {
     const parsedOpenId = GoogleAuth.parseIdToken(tokensObj.id_token);
     const { full_name, picture } = await AcctStore.get(acctEmail, ['full_name', 'picture']);
+    const googleTokenExpires = new Date().getTime() + (tokensObj.expires_in as number - 120) * 1000; // let our copy expire 2 minutes beforehand
     const toSave: AcctStoreDict = {
-      google_token_access: tokensObj.access_token,
-      google_token_expires: new Date().getTime() + (tokensObj.expires_in as number) * 1000,
       google_token_scopes: scopes,
       full_name: full_name || parsedOpenId.name,
       picture: picture || parsedOpenId.picture,
@@ -293,7 +298,8 @@ export class GoogleAuth {
       toSave.google_token_refresh = tokensObj.refresh_token;
     }
     await AcctStore.set(acctEmail, toSave);
-    await InMemoryStore.set(acctEmail, InMemoryStore.ID_TOKEN_STORAGE_KEY, tokensObj.id_token);
+    await InMemoryStore.set(acctEmail, InMemoryStoreKeys.ID_TOKEN, tokensObj.id_token);
+    await InMemoryStore.set(acctEmail, InMemoryStoreKeys.GOOGLE_TOKEN_ACCESS, tokensObj.access_token, googleTokenExpires);
   };
 
   private static googleAuthGetTokens = async (code: string) => {
@@ -322,13 +328,6 @@ export class GoogleAuth {
     }, Catch.stackTrace()) as any as GoogleAuthTokenInfo;
   };
 
-  /**
-   * oauth token will be valid for another 2 min
-   */
-  private static googleApiIsAuthTokenValid = (s: AcctStoreDict) => {
-    return s.google_token_access && (!s.google_token_expires || s.google_token_expires > Date.now() + (120 * 1000));
-  };
-
   // todo - would be better to use a TS type guard instead of the type cast when checking OpenId
   // check for things we actually use: photo/name/locale
   private static parseIdToken = (idToken: string): GmailRes.OpenId => {
diff --git a/extension/js/common/browser/browser-msg.ts b/extension/js/common/browser/browser-msg.ts
@@ -50,7 +50,7 @@ export namespace Bm {
   export type StripeResult = { token: string };
   export type PassphraseEntry = { entered: boolean, initiatorFrameId?: string };
   export type Db = { f: string, args: any[] };
-  export type InMemoryStoreSet = { acctEmail: string, key: string, value: string | undefined };
+  export type InMemoryStoreSet = { acctEmail: string, key: string, value: string | undefined, expiration: number | undefined };
   export type InMemoryStoreGet = { acctEmail: string, key: string };
   export type StoreGlobalGet = { keys: GlobalIndex[]; };
   export type StoreGlobalSet = { values: GlobalStoreDict; };
diff --git a/extension/js/common/core/const.ts b/extension/js/common/core/const.ts
@@ -34,3 +34,8 @@ export const gmailBackupSearchQuery = (acctEmail: string) => {
     '-is:trash'
   ].join(' ');
 };
+
+export class InMemoryStoreKeys {
+  public static readonly ID_TOKEN = 'idToken';
+  public static readonly GOOGLE_TOKEN_ACCESS = 'google_token_access';
+}
diff --git a/extension/js/common/core/expiration-cache.ts b/extension/js/common/core/expiration-cache.ts
@@ -9,9 +9,9 @@ export class ExpirationCache {
   constructor(public EXPIRATION_TICKS: number) {
   }
 
-  public set = (key: string, value: string | undefined) => {
+  public set = (key: string, value?: string, expiration?: number) => {
     if (value) {
-      this.cache[key] = { value, expiration: Date.now() + this.EXPIRATION_TICKS };
+      this.cache[key] = { value, expiration: expiration || (Date.now() + this.EXPIRATION_TICKS) };
     } else {
       delete this.cache[key];
     }
diff --git a/extension/js/common/platform/store/acct-store.ts b/extension/js/common/platform/store/acct-store.ts
@@ -26,7 +26,7 @@ export type Scopes = {
   gmail: boolean;
 };
 
-export type AccountIndex = 'keys' | 'notification_setup_needed_dismissed' | 'email_provider' | 'google_token_access' | 'google_token_expires' | 'google_token_scopes' |
+export type AccountIndex = 'keys' | 'notification_setup_needed_dismissed' | 'email_provider' | 'google_token_scopes' |
   'google_token_refresh' | 'hide_message_password' | 'sendAs' |
   'pubkey_sent_to' | 'full_name' | 'cryptup_enabled' | 'setup_done' |
   'successfully_received_at_leat_one_message' | 'notification_setup_done_seen' | 'picture' |
@@ -43,8 +43,6 @@ export type AcctStoreDict = {
   keys?: KeyInfo[];
   notification_setup_needed_dismissed?: boolean;
   email_provider?: EmailProvider;
-  google_token_access?: string;
-  google_token_expires?: number;
   google_token_scopes?: string[]; // these are actuall scope urls the way the provider expects them
   google_token_refresh?: string;
   hide_message_password?: boolean; // is global?
diff --git a/extension/js/common/platform/store/in-memory-store.ts b/extension/js/common/platform/store/in-memory-store.ts
@@ -9,10 +9,8 @@ import { BrowserMsg } from '../../browser/browser-msg.js';
  */
 export class InMemoryStore extends AbstractStore {
 
-  public static ID_TOKEN_STORAGE_KEY = 'idToken';
-
-  public static set = async (acctEmail: string, key: string, value: string | undefined) => {
-    return await BrowserMsg.send.bg.await.inMemoryStoreSet({ acctEmail, key, value });
+  public static set = async (acctEmail: string, key: string, value?: string, expiration?: number) => {
+    return await BrowserMsg.send.bg.await.inMemoryStoreSet({ acctEmail, key, value, expiration });
   };
 
   public static get = async (acctEmail: string, key: string): Promise<string | undefined> => {
diff --git a/extension/js/content_scripts/webmail/webmail.ts b/extension/js/content_scripts/webmail/webmail.ts
@@ -83,7 +83,7 @@ Catch.try(async () => {
 
     const start = async (acctEmail: string, injector: Injector, notifications: Notifications, factory: XssSafeFactory, notifyMurdered: () => void) => {
       hijackGmailHotkeys();
-      const storage = await AcctStore.get(acctEmail, ['sendAs', 'google_token_scopes', 'full_name']);
+      const storage = await AcctStore.get(acctEmail, ['sendAs', 'full_name']);
       const orgRules = await OrgRules.newInstance(acctEmail);
       if (!storage.sendAs) {
         storage.sendAs = {};
diff --git a/test/source/tests/tooling/browser-recipe.ts b/test/source/tests/tooling/browser-recipe.ts
@@ -3,13 +3,15 @@
 import { Config, Util, TestMessage } from '../../util';
 
 import { AvaContext } from '.';
-import { BrowserHandle, ControllablePage } from '../../browser';
+import { BrowserHandle, Controllable, ControllablePage } from '../../browser';
 import { OauthPageRecipe } from './../page-recipe/oauth-page-recipe';
 import { SetupPageRecipe } from './../page-recipe/setup-page-recipe';
 import { TestUrls } from '../../browser/test-urls';
 import { google } from 'googleapis';
 import { testVariant } from '../../test';
 import { testConstants } from './consts';
+import { PageRecipe } from '../page-recipe/abstract-page-recipe';
+import { InMemoryStoreKeys } from '../../core/const';
 
 export class BrowserRecipe {
 
@@ -88,8 +90,17 @@ export class BrowserRecipe {
     }
   };
 
+  public static getGoogleAccessToken = async (controllable: Controllable, acctEmail: string): Promise<string> => {
+    const result = await PageRecipe.sendMessage(controllable, {
+      name: 'inMemoryStoreGet',
+      // tslint:disable-next-line:no-null-keyword
+      data: { bm: { acctEmail, key: InMemoryStoreKeys.GOOGLE_TOKEN_ACCESS }, objUrls: {} }, to: null, uid: '2'
+    });
+    return (result as { result: string }).result;
+  };
+
   public static deleteAllDraftsInGmailAccount = async (settingsPage: ControllablePage): Promise<void> => {
-    const accessToken = (await settingsPage.getFromLocalStorage(['cryptup_citestsgmailflowcryptdev_google_token_access'])).cryptup_citestsgmailflowcryptdev_google_token_access as string;
+    const accessToken = await BrowserRecipe.getGoogleAccessToken(settingsPage, 'ci.tests.gmail@flowcrypt.dev');
     const gmail = google.gmail({ version: 'v1' });
     const list = await gmail.users.drafts.list({ userId: 'me', access_token: accessToken });
     if (list.data.drafts) {

Original file line number	Diff line number	Diff line change
`@@ -9,9 +9,9 @@ export class ExpirationCache {`
`9`	`9`	`constructor(public EXPIRATION_TICKS: number) {`
`10`	`10`	`}`
`11`	`11`
`12`		`- public set = (key: string, value: string \| undefined) => {`
	`12`	`+ public set = (key: string, value?: string, expiration?: number) => {`
`13`	`13`	`if (value) {`
`14`		`- this.cache[key] = { value, expiration: Date.now() + this.EXPIRATION_TICKS };`
	`14`	`+ this.cache[key] = { value, expiration: expiration \|\| (Date.now() + this.EXPIRATION_TICKS) };`
`15`	`15`	`} else {`
`16`	`16`	`delete this.cache[key];`
`17`	`17`	`}`