#5799 Use custom idp token for enterprise server authentication if special jwt is stored in local store

extension/chrome/settings/setup/setup-key-manager-autogen.ts

@@ -55,7 +55,7 @@ export class SetupWithEmailKeyManagerModule {
     /* eslint-enable @typescript-eslint/naming-convention */
     try {
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      const { privateKeys } = await this.view.keyManager!.getPrivateKeys(this.view.idToken!);
+      const { privateKeys } = await this.view.keyManager!.getPrivateKeys(this.view.acctEmail);
       if (privateKeys.length) {
         // keys already exist on keyserver, auto-import
         try {
@@ -115,7 +115,7 @@ export class SetupWithEmailKeyManagerModule {
     }
     const storePrvOnKm = async () => {
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      await this.view.keyManager!.storePrivateKey(this.view.idToken!, KeyUtil.armor(decryptablePrv));
+      await this.view.keyManager!.storePrivateKey(this.view.acctEmail, KeyUtil.armor(decryptablePrv));
     };
     await Settings.retryUntilSuccessful(
       storePrvOnKm,

extension/js/common/api/account-servers/external-service.ts

@@ -5,16 +5,16 @@ import { Api, ProgressCb, ProgressCbs } from '../shared/api.js';
 import { AcctStore } from '../../platform/store/acct-store.js';
 import { Dict, Str } from '../../core/common.js';
 import { ErrorReport } from '../../platform/catch.js';
-import { ApiErr, BackendAuthErr } from '../shared/api-error.js';
-import { FLAVOR, InMemoryStoreKeys } from '../../core/const.js';
+import { ApiErr } from '../shared/api-error.js';
+import { FLAVOR } from '../../core/const.js';
 import { Attachment } from '../../core/attachment.js';
 import { ParsedRecipients } from '../email-provider/email-provider-api.js';
 import { Buf } from '../../core/buf.js';
 import { ClientConfigurationError, ClientConfigurationJson } from '../../client-configuration.js';
-import { InMemoryStore } from '../../platform/store/in-memory-store.js';
 import { Serializable } from '../../platform/store/abstract-store.js';
 import { AuthenticationConfiguration } from '../../authentication-configuration.js';
 import { Xss } from '../../platform/xss.js';
+import { ConfiguredIdpOAuth } from '../authentication/configured-idp-oauth.js';
 
 // todo - decide which tags to use
 type EventTag = 'compose' | 'decrypt' | 'setup' | 'settings' | 'import-pub' | 'import-prv';
@@ -160,15 +160,6 @@ export class ExternalService extends Api {
     });
   };
 
-  private authHdr = async (): Promise<{ authorization: string }> => {
-    const idToken = await InMemoryStore.getUntilAvailable(this.acctEmail, InMemoryStoreKeys.ID_TOKEN);
-    if (idToken) {
-      return { authorization: `Bearer ${idToken}` };
-    }
-    // user will not actually see this message, they'll see a generic login prompt
-    throw new BackendAuthErr('Missing id token, please re-authenticate');
-  };
-
   private request = async <RT>(
     path: string,
     vals?:
@@ -192,6 +183,6 @@ export class ExternalService extends Api {
           method: 'POST',
         }
       : undefined;
-    return await ExternalService.apiCall(this.url, path, values, progress, await this.authHdr(), 'json');
+    return await ExternalService.apiCall(this.url, path, values, progress, await ConfiguredIdpOAuth.authHdr(this.acctEmail), 'json');
   };
 }

extension/js/common/api/authentication/configured-idp-oauth.ts

@@ -12,6 +12,7 @@ import { Catch } from '../../platform/catch.js';
 import { InMemoryStoreKeys } from '../../core/const.js';
 import { InMemoryStore } from '../../platform/store/in-memory-store.js';
 import { AcctStore } from '../../platform/store/acct-store.js';
+import { BackendAuthErr } from '../shared/api-error.js';
 export class ConfiguredIdpOAuth extends OAuth {
   public static newAuthPopupForEnterpriseServerAuthenticationIfNeeded = async (authRes: AuthRes) => {
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
@@ -24,6 +25,24 @@ export class ConfiguredIdpOAuth extends OAuth {
     return authRes;
   };
 
+  public static authHdr = async (acctEmail: string, shouldThrowErrorForEmptyIdToken = true): Promise<{ authorization: string } | undefined> => {
+    let idToken = await InMemoryStore.getUntilAvailable(acctEmail, InMemoryStoreKeys.ID_TOKEN);
+    if (idToken) {
+      const customIDPIdToken = await InMemoryStore.get(acctEmail, InMemoryStoreKeys.CUSTOM_IDP_ID_TOKEN);
+      // if special JWT is stored in local store, it should be used for Enterprise Server authentication instead of Google JWT
+      // https://github.com/FlowCrypt/flowcrypt-browser/issues/5799
+      if (customIDPIdToken) {
+        idToken = customIDPIdToken;
+      }
+      return { authorization: `Bearer ${idToken}` };
+    }
+    if (shouldThrowErrorForEmptyIdToken) {
+      // user will not actually see this message, they'll see a generic login prompt
+      throw new BackendAuthErr('Missing id token, please re-authenticate');
+    }
+    return undefined;
+  };
+
   public static async newAuthPopup(acctEmail: string, authConf: AuthenticationConfiguration): Promise<AuthRes> {
     acctEmail = acctEmail.toLowerCase();
     const authRequest = this.newAuthRequest(acctEmail, this.OAUTH_REQUEST_SCOPES);

extension/js/common/api/key-server/key-manager.ts

@@ -4,6 +4,7 @@
 
 import { Api } from './../shared/api.js';
 import { Url } from '../../core/common.js';
+import { ConfiguredIdpOAuth } from '../authentication/configured-idp-oauth.js';
 
 type LoadPrvRes = { privateKeys: { decryptedPrivateKey: string }[] };
 
@@ -15,17 +16,17 @@ export class KeyManager extends Api {
     this.url = Url.removeTrailingSlash(url);
   }
 
-  public getPrivateKeys = async (idToken: string): Promise<LoadPrvRes> => {
-    return await Api.apiCall(this.url, '/v1/keys/private', undefined, undefined, idToken ? { authorization: `Bearer ${idToken}` } : undefined, 'json');
+  public getPrivateKeys = async (acctEmail: string): Promise<LoadPrvRes> => {
+    return await Api.apiCall(this.url, '/v1/keys/private', undefined, undefined, await ConfiguredIdpOAuth.authHdr(acctEmail, false), 'json');
   };
 
-  public storePrivateKey = async (idToken: string, privateKey: string): Promise<void> => {
+  public storePrivateKey = async (acctEmail: string, privateKey: string): Promise<void> => {
     await Api.apiCall(
       this.url,
       '/v1/keys/private',
       { data: { privateKey }, fmt: 'JSON', method: 'PUT' },
       undefined,
-      idToken ? { authorization: `Bearer ${idToken}` } : undefined
+      await ConfiguredIdpOAuth.authHdr(acctEmail, false)
     );
   };
 }

extension/js/content_scripts/webmail/generic/setup-webmail-content-script.ts

@@ -349,7 +349,7 @@ export const contentScriptSetupIfVacant = async (webmailSpecific: WebmailSpecifi
         const keyManager = new KeyManager(clientConfiguration.getKeyManagerUrlForPrivateKeys()!);
         Catch.setHandledTimeout(async () => {
           try {
-            const { privateKeys } = await keyManager.getPrivateKeys(idToken);
+            const { privateKeys } = await keyManager.getPrivateKeys(acctEmail);
             await processKeysFromEkm(
               acctEmail,
               privateKeys.map(entry => entry.decryptedPrivateKey),

test/source/mock/fes/customer-url-fes-endpoints.ts

@@ -11,11 +11,13 @@ import { FesConfig } from './shared-tenant-fes-endpoints';
 const standardFesUrl = (port: string) => {
   return `fes.standardsubdomainfes.localhost:${port}`;
 };
-const issuedAccessTokens: string[] = [];
+export const issuedGoogleIDPIdTokens: string[] = [];
+export const issuedCustomIDPIdTokens: string[] = [];
 
 // eslint-disable-next-line @typescript-eslint/naming-convention
 export const standardSubDomainFesClientConfiguration = { flags: [], disallow_attester_search_for_domains: ['[email protected]'] };
 export const getMockCustomerUrlFesEndpoints = (config: FesConfig | undefined): HandlersDefinition => {
+  const isCustomIDPUsed = !!config?.authenticationConfiguration;
   return {
     // standard fes location at https://fes.domain.com
     '/api/': async ({}, req) => {
@@ -57,7 +59,7 @@ export const getMockCustomerUrlFesEndpoints = (config: FesConfig | undefined): H
     },
     '/api/v1/message/new-reply-token': async ({}, req) => {
       if (parseAuthority(req) === standardFesUrl(parsePort(req)) && req.method === 'POST') {
-        authenticate(req, 'oidc');
+        authenticate(req, isCustomIDPUsed);
         return { replyToken: 'mock-fes-reply-token' };
       }
       throw new HttpClientErr('Not Found', 404);
@@ -67,7 +69,7 @@ export const getMockCustomerUrlFesEndpoints = (config: FesConfig | undefined): H
       const fesUrl = standardFesUrl(port);
       // body is a mime-multipart string, we're doing a few smoke checks here without parsing it
       if (parseAuthority(req) === fesUrl && req.method === 'POST' && typeof body === 'string') {
-        authenticate(req, 'oidc');
+        authenticate(req, isCustomIDPUsed);
         if (config?.messagePostValidator) {
           return await config.messagePostValidator(body, fesUrl);
         }
@@ -79,7 +81,7 @@ export const getMockCustomerUrlFesEndpoints = (config: FesConfig | undefined): H
       const port = parsePort(req);
       if (parseAuthority(req) === standardFesUrl(port) && req.method === 'POST') {
         // test: `compose - [email protected]:8001 - PWD encrypted message with FES web portal`
-        authenticate(req, 'oidc');
+        authenticate(req, isCustomIDPUsed);
         expect(body).to.match(messageIdRegex(port));
         return {};
       }
@@ -89,7 +91,7 @@ export const getMockCustomerUrlFesEndpoints = (config: FesConfig | undefined): H
       const port = parsePort(req);
       if (parseAuthority(req) === standardFesUrl(port) && req.method === 'POST') {
         // test: `compose - [email protected]:8001 - PWD encrypted message with FES - Reply rendering`
-        authenticate(req, 'oidc');
+        authenticate(req, isCustomIDPUsed);
         expect(body).to.match(messageIdRegex(port));
         return {};
       }
@@ -102,7 +104,7 @@ export const getMockCustomerUrlFesEndpoints = (config: FesConfig | undefined): H
         // test: `compose - [email protected]:8001 - PWD encrypted message with FES - Reply rendering`
         // test: `compose - [email protected]:8001 - PWD encrypted message with FES web portal - pubkey recipient in bcc`
         // test: `compose - [email protected]:8001 - PWD encrypted message with FES web portal - some sends fail with BadRequest error`
-        authenticate(req, 'oidc');
+        authenticate(req, isCustomIDPUsed);
         expect(body).to.match(messageIdRegex(port));
         return {};
       }
@@ -112,7 +114,7 @@ export const getMockCustomerUrlFesEndpoints = (config: FesConfig | undefined): H
       const port = parsePort(req);
       if (parseAuthority(req) === standardFesUrl(port) && req.method === 'POST') {
         // test: `compose - [email protected]:8001 - PWD encrypted message with FES web portal`
-        authenticate(req, 'oidc');
+        authenticate(req, isCustomIDPUsed);
         expect(body).to.match(messageIdRegex(port));
         return {};
       }
@@ -125,20 +127,15 @@ export const getMockCustomerUrlFesEndpoints = (config: FesConfig | undefined): H
   };
 };
 
-const authenticate = (req: { headers: IncomingHttpHeaders }, type: 'oidc' | 'fes'): string => {
+const authenticate = (req: { headers: IncomingHttpHeaders }, isCustomIDPUsed: boolean): string => {
   const jwt = (req.headers.authorization || '').replace('Bearer ', '');
   if (!jwt) {
     throw new Error('Mock FES missing authorization header');
   }
-  if (type === 'oidc') {
-    if (issuedAccessTokens.includes(jwt)) {
-      throw new Error('Mock FES access-token call wrongly with FES token');
-    }
-  } else {
-    // fes
-    if (!issuedAccessTokens.includes(jwt)) {
-      throw new HttpClientErr('FES mock received access token it didnt issue', 401);
-    }
+  const issuedTokens = isCustomIDPUsed ? issuedCustomIDPIdTokens : issuedGoogleIDPIdTokens;
+
+  if (!issuedTokens.includes(jwt)) {
+    throw new HttpClientErr('FES mock received access token it didnt issue', 401);
   }
   return MockJwt.parseEmail(jwt);
 };

test/source/mock/google/google-endpoints.ts

@@ -148,15 +148,15 @@ export const getMockGoogleEndpoints = (oauth: OauthMock, config: GoogleConfig |
       if (isPost(req)) {
         if (grant_type === 'authorization_code' && code && client_id === oauth.clientId) {
           // auth code from auth screen gets exchanged for access and refresh tokens
-          return oauth.getRefreshTokenResponse(code);
+          return oauth.getRefreshTokenResponse(code, false);
         } else if (grant_type === 'refresh_token' && refreshToken && client_id === oauth.clientId) {
           // here also later refresh token gets exchanged for access token
-          return oauth.getTokenResponse(refreshToken);
+          return oauth.getTokenResponse(refreshToken, false);
         }
         const parsedBody = body as OAuthTokenRequestModel;
         // Above is for Google OAuth and this is for normal OAuth
         if (parsedBody.grant_type === 'authorization_code' && parsedBody.code && parsedBody.client_id === OauthMock.customIDPClientId) {
-          return oauth.getRefreshTokenResponse(parsedBody.code);
+          return oauth.getRefreshTokenResponse(parsedBody.code, true);
         }
       }
       throw new Error(`Method not implemented for ${req.url}: ${req.method}`);

test/source/mock/lib/oauth.ts

@@ -4,6 +4,7 @@ import { HttpClientErr, Status } from './api';
 
 import { Buf } from '../../core/buf';
 import { Str } from '../../core/common';
+import { issuedCustomIDPIdTokens, issuedGoogleIDPIdTokens } from '../fes/customer-url-fes-endpoints';
 
 const authURL = 'https://localhost:8001';
 
@@ -21,6 +22,15 @@ export class OauthMock {
   private acctByIdToken: { [acct: string]: string } = {};
   private issuedIdTokensByAcct: { [acct: string]: string[] } = {};
 
+  public static getCustomIDPOAuthConfig = (port: number | undefined) => {
+    return {
+      clientId: OauthMock.customIDPClientId,
+      clientSecret: OauthMock.customIDPClientSecret,
+      redirectUrl: `custom-redirect-url`, // This won't be used as we use our https://{id}.chromiumapp.org with chrome.identity.getRedirectURL
+      authCodeUrl: `https://localhost:${port}/o/oauth2/auth`,
+      tokensUrl: `https://localhost:${port}/token`,
+    };
+  };
   public renderText = (text: string) => {
     return this.htmlPage(text, text);
   };
@@ -43,12 +53,12 @@ export class OauthMock {
     return url.href;
   };
 
-  public getRefreshTokenResponse = (code: string) => {
+  public getRefreshTokenResponse = (code: string, isCustomIDPAuth: boolean) => {
     /* eslint-disable @typescript-eslint/naming-convention */
     const refresh_token = this.refreshTokenByAuthCode[code];
     const access_token = this.getAccessToken(refresh_token);
     const acct = this.acctByAccessToken[access_token];
-    const id_token = this.generateIdToken(acct);
+    const id_token = this.generateIdToken(acct, isCustomIDPAuth);
     return { access_token, refresh_token, expires_in: this.expiresIn, id_token, token_type: 'refresh_token' }; // guessed the token_type
     /* eslint-enable @typescript-eslint/naming-convention */
   };
@@ -62,12 +72,12 @@ export class OauthMock {
     };
   };
 
-  public getTokenResponse = (refreshToken: string) => {
+  public getTokenResponse = (refreshToken: string, isCustomIDPAuth: boolean) => {
     try {
       /* eslint-disable @typescript-eslint/naming-convention */
       const access_token = this.getAccessToken(refreshToken);
       const acct = this.acctByAccessToken[access_token];
-      const id_token = this.generateIdToken(acct);
+      const id_token = this.generateIdToken(acct, isCustomIDPAuth);
       return { access_token, expires_in: this.expiresIn, id_token, token_type: 'Bearer' };
       /* eslint-enable @typescript-eslint/naming-convention */
     } catch (e) {
@@ -141,13 +151,18 @@ export class OauthMock {
 
   // -- private
 
-  private generateIdToken = (email: string): string => {
+  private generateIdToken = (email: string, isCustomIDPAuth: boolean): string => {
     const newIdToken = MockJwt.new(email, this.expiresIn);
     if (!this.issuedIdTokensByAcct[email]) {
       this.issuedIdTokensByAcct[email] = [];
     }
     this.issuedIdTokensByAcct[email].push(newIdToken);
     this.acctByIdToken[newIdToken] = email;
+    if (isCustomIDPAuth) {
+      issuedCustomIDPIdTokens.push(newIdToken);
+    } else {
+      issuedGoogleIDPIdTokens.push(newIdToken);
+    }
     return newIdToken;
   };
 

test/source/tests/compose.ts

@@ -43,6 +43,7 @@ import {
 import { Buf } from '../core/buf';
 import { flowcryptCompatibilityAliasList, flowcryptCompatibilityPrimarySignature } from '../mock/google/google-endpoints';
 import { standardSubDomainFesClientConfiguration } from '../mock/fes/customer-url-fes-endpoints';
+import { OauthMock } from '../mock/lib/oauth';
 
 export const defineComposeTests = (testVariant: TestVariant, testWithBrowser: TestWithBrowser) => {
   if (testVariant !== 'CONSUMER-LIVE-GMAIL') {
@@ -3007,18 +3008,21 @@ export const defineComposeTests = (testVariant: TestVariant, testWithBrowser: Te
     test(
       '[email protected]:8001 - PWD encrypted message with FES web portal',
       testWithBrowser(async (t, browser) => {
+        const port = t.context.urls?.port;
         t.context.mockApi!.configProvider = new ConfigurationProvider({
           attester: {
             pubkeyLookup: {},
           },
           fes: {
+            authenticationConfiguration: {
+              oauth: OauthMock.getCustomIDPOAuthConfig(port),
+            },
             messagePostValidator: processMessageFromUser,
             clientConfiguration: standardSubDomainFesClientConfiguration,
           },
         });
-        const port = t.context.urls?.port;
         const acct = `[email protected]:${port}`; // added port to trick extension into calling the mock
-        const settingsPage = await BrowserRecipe.openSettingsLoginApprove(t, browser, acct);
+        const settingsPage = await BrowserRecipe.openSettingsLoginApprove(t, browser, acct, true);
         await SetupPageRecipe.manualEnter(
           settingsPage,
           'flowcrypt.test.key.used.pgp',

test/source/tests/setup.ts

@@ -2526,14 +2526,8 @@ AN8G3r5Htj8olot+jm9mIa5XLXWzMNUZgg==
     test(
       'setup - check custom authentication config from the local store (customer url fes)',
       testWithBrowser(async (t, browser) => {
-        const port = t.context.urls?.port ?? '';
-        const oauthConfig = {
-          clientId: OauthMock.customIDPClientId,
-          clientSecret: OauthMock.customIDPClientSecret,
-          redirectUrl: `custom-redirect-url`, // This won't be used as we use our https://{id}.chromiumapp.org with chrome.identity.getRedirectURL
-          authCodeUrl: `https://localhost:${port}/o/oauth2/auth`,
-          tokensUrl: `https://localhost:${port}/token`,
-        };
+        const port = t.context.urls?.port;
+        const oauthConfig = OauthMock.getCustomIDPOAuthConfig(port);
         t.context.mockApi!.configProvider = new ConfigurationProvider({
           attester: {
             pubkeyLookup: {},