fix(oauth): Move keyFetchToken and unwrapBKey to sensitiveDataClient

Because:
* It's better practice to store these in memory, and this causes a problem on session restart in oauth desktop

This commit:
* Moves keyFetchToken and unwrapBKey off of location state and into sensitiveDataClient

closes FXA-10709, closes FXA-10529

packages/fxa-settings/src/lib/oauth/hooks.tsx

@@ -9,6 +9,7 @@ import {
   OAuthIntegration,
   isOAuthNativeIntegrationSync,
   isOAuthIntegration,
+  isSyncDesktopV3Integration,
 } from '../../models';
 import { createEncryptedBundle } from '../crypto/scoped-keys';
 import { Constants } from '../constants';
@@ -283,3 +284,31 @@ export function useFinishOAuthFlowHandler(
   }
   return { oAuthDataError: null, finishOAuthFlowHandler };
 }
+
+/**
+ * If we don't have `keyFetchToken` or `unwrapBKey` and it's required, the user needs to
+ * restart the flow. We only store these values in memory, so they don't persist across page
+ * refreshes or browser session restarts. We can't redirect to /signin and reprompt for
+ * password in a browser session restart/crash because the browser stores the flow state in
+ * memory and the query params won't match after a browser session is restored.
+ *
+ * TODO: Can we check session storage for if the user just refreshed so we can redirect them
+ * to /signin instead of showing an error component? FXA-10707
+ */
+export function useOAuthKeysCheck(
+  integration: Pick<Integration, 'type' | 'wantsKeys'>,
+  keyFetchToken?: hexstring,
+  unwrapBKey?: hexstring
+) {
+  if (
+    (isOAuthIntegration(integration) ||
+      isSyncDesktopV3Integration(integration)) &&
+    integration.wantsKeys() &&
+    (!keyFetchToken || !unwrapBKey)
+  ) {
+    return {
+      oAuthKeysCheckError: OAUTH_ERRORS.TRY_AGAIN,
+    };
+  }
+  return { oAuthKeysCheckError: null };
+}

packages/fxa-settings/src/lib/sensitive-data-client.ts

@@ -6,6 +6,20 @@
 
 export const AUTH_DATA_KEY = 'auth';
 
+export type SensitiveDataClientData = {
+  [AUTH_DATA_KEY]: {
+    emailForAuth?: string;
+    authPW?: string;
+    keyFetchToken?: hexstring;
+    unwrapBKey?: hexstring;
+  };
+};
+
+export type SensitiveDataClientAuthKeys = Pick<
+  SensitiveDataClientData[typeof AUTH_DATA_KEY],
+  'keyFetchToken' | 'unwrapBKey'
+>;
+
 /**
  * Class representing a client for handling sensitive data.
  *

packages/fxa-settings/src/models/integrations/oauth-native-integration.ts

@@ -20,6 +20,9 @@ export function isOAuthNativeIntegration(integration: {
 
 export type OAuthIntegration = OAuthWebIntegration | OAuthNativeIntegration;
 
+/**
+ * Check if the integration is OAuthWeb or OAuthNative
+ */
 export function isOAuthIntegration(integration: {
   type: IntegrationType;
 }): integration is OAuthIntegration {

packages/fxa-settings/src/pages/InlineRecoveryKeySetup/container.test.tsx

@@ -55,6 +55,7 @@ function mockModelsModule() {
   mockSensitiveDataClient.getData = jest.fn().mockReturnValue({
     emailForAuth: '[email protected]',
     authPW: MOCK_AUTH_PW,
+    unwrapBKey: MOCK_UNWRAP_BKEY,
   });
 }
 

packages/fxa-settings/src/pages/InlineRecoveryKeySetup/container.tsx

@@ -28,13 +28,16 @@ export const InlineRecoveryKeySetupContainer = (_: RouteComponentProps) => {
   const location = useLocation() as ReturnType<typeof useLocation> & {
     state?: SigninLocationState;
   };
-  const { email, uid, sessionToken, unwrapBKey } = location.state || {};
+  const { email, uid, sessionToken } = location.state || {};
 
   const sensitiveDataClient = useSensitiveDataClient();
   const sensitiveData = sensitiveDataClient.getData(AUTH_DATA_KEY);
-  const { authPW, emailForAuth } =
-    (sensitiveData as unknown as { emailForAuth: string; authPW: string }) ||
-    {};
+  const { authPW, emailForAuth, unwrapBKey } =
+    (sensitiveData as {
+      emailForAuth?: string;
+      authPW?: string;
+      unwrapBKey?: hexstring;
+    }) || {};
 
   const navigateForward = useCallback(() => {
     setCurrentStep(currentStep + 1);
@@ -44,7 +47,9 @@ export const InlineRecoveryKeySetupContainer = (_: RouteComponentProps) => {
     (
         uid: string,
         sessionToken: string,
-        unwrapBKey: string
+        unwrapBKey: string,
+        emailForAuth: string,
+        authPW: string
       ): (() => Promise<CreateRecoveryKeyHandler>) =>
       async () => {
         try {
@@ -93,7 +98,7 @@ export const InlineRecoveryKeySetupContainer = (_: RouteComponentProps) => {
           };
         }
       },
-    [authClient, emailForAuth, authPW, navigateForward, ftlMsgResolver]
+    [authClient, navigateForward, ftlMsgResolver]
   );
 
   const updateRecoveryHint = useCallback(
@@ -124,7 +129,9 @@ export const InlineRecoveryKeySetupContainer = (_: RouteComponentProps) => {
   const createRecoveryKeyHandler = createRecoveryKey(
     uid,
     sessionToken,
-    unwrapBKey
+    unwrapBKey,
+    emailForAuth,
+    authPW
   );
   const updateRecoveryHintHandler = updateRecoveryHint(sessionToken);
 

packages/fxa-settings/src/pages/InlineRecoverySetup/container.test.tsx

@@ -11,7 +11,8 @@ import { LocationProvider } from '@reach/router';
 import { renderWithLocalizationProvider } from 'fxa-react/lib/test-utils/localizationProvider';
 import { AuthUiError, AuthUiErrors } from '../../lib/auth-errors/auth-errors';
 import { MozServices } from '../../lib/types';
-import { OAuthIntegration } from '../../models';
+import { OAuthIntegration, useSensitiveDataClient } from '../../models';
+import { mockSensitiveDataClient as createMockSensitiveDataClient } from '../../models/mocks';
 import {
   MOCK_QUERY_PARAMS,
   MOCK_SIGNIN_LOCATION_STATE,
@@ -27,7 +28,11 @@ import {
   MOCK_OAUTH_FLOW_HANDLER_RESPONSE,
   MOCK_TOTP_STATUS_VERIFIED,
 } from '../Signin/mocks';
-import { useFinishOAuthFlowHandler } from '../../lib/oauth/hooks';
+import {
+  useFinishOAuthFlowHandler,
+  useOAuthKeysCheck,
+} from '../../lib/oauth/hooks';
+import { AUTH_DATA_KEY } from '../../lib/sensitive-data-client';
 
 let mockLocationState = {};
 const mockSearch = '?' + new URLSearchParams(MOCK_QUERY_PARAMS);
@@ -52,9 +57,12 @@ jest.mock('../../lib/oauth/hooks.tsx', () => {
   return {
     __esModule: true,
     useFinishOAuthFlowHandler: jest.fn(),
+    useOAuthKeysCheck: jest.fn(),
   };
 });
 
+const mockSensitiveDataClient = createMockSensitiveDataClient();
+mockSensitiveDataClient.getData = jest.fn();
 const mockAuthClient = new AuthClient('http://localhost:9000', {
   keyStretchVersion: 1,
 });
@@ -64,6 +72,7 @@ jest.mock('../../models', () => {
     ...jest.requireActual('../../models'),
     useSession: jest.fn(() => mockSessionHook()),
     useAuthClient: jest.fn(() => mockAuthClient),
+    useSensitiveDataClient: jest.fn(),
   };
 });
 
@@ -119,6 +128,12 @@ function setMocks() {
       .mockReturnValueOnce(MOCK_OAUTH_FLOW_HANDLER_RESPONSE),
     oAuthDataError: null,
   }));
+  (useSensitiveDataClient as jest.Mock).mockImplementation(
+    () => mockSensitiveDataClient
+  );
+  (useOAuthKeysCheck as jest.Mock).mockImplementation(() => ({
+    oAuthKeysCheckError: null,
+  }));
 }
 
 const defaultProps = {
@@ -201,6 +216,13 @@ describe('InlineRecoverySetupContainer', () => {
       });
     });
 
+    it('reads data from sensitive data client', () => {
+      render();
+      expect(mockSensitiveDataClient.getData).toHaveBeenCalledWith(
+        AUTH_DATA_KEY
+      );
+    });
+
     describe('callbacks', () => {
       describe('cancelSetupHandler', () => {
         it('redirects when returnOnError is true', async () => {

packages/fxa-settings/src/pages/InlineRecoverySetup/container.tsx

@@ -11,10 +11,15 @@ import { AuthUiErrors } from '../../lib/auth-errors/auth-errors';
 import {
   FinishOAuthFlowHandlerResult,
   useFinishOAuthFlowHandler,
+  useOAuthKeysCheck,
 } from '../../lib/oauth/hooks';
 import { getCode } from '../../lib/totp';
 import { MozServices } from '../../lib/types';
-import { OAuthIntegration, useAuthClient } from '../../models';
+import {
+  OAuthIntegration,
+  useAuthClient,
+  useSensitiveDataClient,
+} from '../../models';
 import { VERIFY_TOTP_MUTATION } from './gql';
 import InlineRecoverySetup from './index';
 import { hardNavigate } from 'fxa-react/lib/utils';
@@ -23,6 +28,10 @@ import { TotpStatusResponse } from '../Signin/SigninTokenCode/interfaces';
 import { GET_TOTP_STATUS } from '../../components/App/gql';
 import OAuthDataError from '../../components/OAuthDataError';
 import { isFirefoxService } from '../../models/integrations/utils';
+import {
+  AUTH_DATA_KEY,
+  SensitiveDataClientAuthKeys,
+} from '../../lib/sensitive-data-client';
 
 export const InlineRecoverySetupContainer = ({
   isSignedIn,
@@ -46,6 +55,16 @@ export const InlineRecoverySetupContainer = ({
   };
   const signinRecoveryLocationState = location.state;
   const { totp, ...signinLocationState } = signinRecoveryLocationState || {};
+  const sensitiveDataClient = useSensitiveDataClient();
+  const sensitiveData = sensitiveDataClient.getData(AUTH_DATA_KEY);
+  const { keyFetchToken, unwrapBKey } =
+    (sensitiveData as SensitiveDataClientAuthKeys) || {};
+
+  const { oAuthKeysCheckError } = useOAuthKeysCheck(
+    integration,
+    keyFetchToken,
+    unwrapBKey
+  );
 
   const [recoveryCodes, setRecoveryCodes] = useState<string[]>();
   const [verifyTotp] = useMutation<{ verifyTotp: { success: boolean } }>(
@@ -78,15 +97,20 @@ export const InlineRecoverySetupContainer = ({
     const { redirect, error } = await finishOAuthFlowHandler(
       signinRecoveryLocationState!.uid,
       signinRecoveryLocationState!.sessionToken,
-      signinRecoveryLocationState!.keyFetchToken,
-      signinRecoveryLocationState!.unwrapBKey
+      keyFetchToken,
+      unwrapBKey
     );
     if (error) {
       setOAuthError(error);
       return;
     }
     hardNavigate(redirect);
-  }, [signinRecoveryLocationState, finishOAuthFlowHandler]);
+  }, [
+    signinRecoveryLocationState,
+    finishOAuthFlowHandler,
+    keyFetchToken,
+    unwrapBKey,
+  ]);
 
   const cancelSetupHandler = useCallback(() => {
     const error = AuthUiErrors.TOTP_REQUIRED;
@@ -125,6 +149,11 @@ export const InlineRecoverySetupContainer = ({
   if (oAuthDataError) {
     return <OAuthDataError error={oAuthDataError} />;
   }
+  // Note that we don't currently need this check on this page right now since AMO is the only
+  // RP requiring 2FA and it doesn't require keys. However it's here for consistency.
+  if (oAuthKeysCheckError) {
+    return <OAuthDataError error={oAuthKeysCheckError} />;
+  }
 
   return (
     <InlineRecoverySetup

packages/fxa-settings/src/pages/Signin/SigninPushCode/container.test.tsx

@@ -7,16 +7,22 @@ import { SigninPushCodeProps } from './interfaces';
 import * as ReactUtils from 'fxa-react/lib/utils';
 import * as CacheModule from '../../../lib/cache';
 
-import { Integration } from '../../../models';
+import { Integration, useSensitiveDataClient } from '../../../models';
 import { renderWithLocalizationProvider } from 'fxa-react/lib/test-utils/localizationProvider';
 import { LocationProvider } from '@reach/router';
 import SigninPushCodeContainer from './container';
 import { waitFor } from '@testing-library/react';
-import { MOCK_EMAIL, MOCK_STORED_ACCOUNT } from '../../mocks';
+import {
+  MOCK_EMAIL,
+  MOCK_KEY_FETCH_TOKEN,
+  MOCK_STORED_ACCOUNT,
+  MOCK_UNWRAP_BKEY,
+} from '../../mocks';
 import {
   createMockSigninLocationState,
   createMockSyncIntegration,
 } from './mocks';
+import { mockSensitiveDataClient as createMockSensitiveDataClient } from '../../../models/mocks';
 
 import { MozServices } from '../../../lib/types';
 
@@ -35,8 +41,11 @@ function applyDefaultMocks() {
 
   mockSigninPushCodeModule();
   mockCurrentAccount();
+  resetMockSensitiveDataClient();
 }
 
+const mockSensitiveDataClient = createMockSensitiveDataClient();
+mockSensitiveDataClient.setData = jest.fn();
 let mockHasTotpAuthClient = false;
 let mockSessionStatus = 'verified';
 let mockSendLoginPushRequest = jest.fn().mockResolvedValue({});
@@ -54,6 +63,7 @@ jest.mock('../../../models', () => {
         sendLoginPushRequest: mockSendLoginPushRequest,
       };
     },
+    useSensitiveDataClient: jest.fn(),
   };
 });
 
@@ -96,6 +106,16 @@ function mockCurrentAccount(storedAccount = { uid: '123' }) {
   jest.spyOn(CacheModule, 'currentAccount').mockReturnValue(storedAccount);
 }
 
+function resetMockSensitiveDataClient() {
+  (useSensitiveDataClient as jest.Mock).mockImplementation(
+    () => mockSensitiveDataClient
+  );
+  mockSensitiveDataClient.getData = jest.fn().mockReturnValue({
+    keyFetchToken: MOCK_KEY_FETCH_TOKEN,
+    unwrapBKey: MOCK_UNWRAP_BKEY,
+  });
+}
+
 async function render() {
   renderWithLocalizationProvider(
     <LocationProvider>

packages/fxa-settings/src/pages/Signin/SigninPushCode/container.tsx

@@ -7,14 +7,23 @@ import SigninPushCode from '.';
 import { MozServices } from '../../../lib/types';
 import { getSigninState, handleNavigation } from '../utils';
 import { SigninLocationState } from '../interfaces';
-import { Integration, isWebIntegration, useAuthClient } from '../../../models';
+import {
+  Integration,
+  isWebIntegration,
+  useAuthClient,
+  useSensitiveDataClient,
+} from '../../../models';
 import { useFinishOAuthFlowHandler } from '../../../lib/oauth/hooks';
 import { hardNavigate } from 'fxa-react/lib/utils';
 import LoadingSpinner from 'fxa-react/components/LoadingSpinner';
 import OAuthDataError from '../../../components/OAuthDataError';
 import { useWebRedirect } from '../../../lib/hooks/useWebRedirect';
 import { useEffect, useState } from 'react';
 import { useNavigateWithQuery as useNavigate } from '../../../lib/hooks/useNavigateWithQuery';
+import {
+  AUTH_DATA_KEY,
+  SensitiveDataClientAuthKeys,
+} from '../../../lib/sensitive-data-client';
 
 export type SigninPushCodeContainerProps = {
   integration: Integration;
@@ -35,8 +44,10 @@ export const SigninPushCodeContainer = ({
   const location = useLocation() as ReturnType<typeof useLocation> & {
     state: SigninLocationState;
   };
-
   const signinState = getSigninState(location.state);
+  const sensitiveDataClient = useSensitiveDataClient();
+  const sensitiveData = sensitiveDataClient.getData(AUTH_DATA_KEY);
+  const { unwrapBKey } = (sensitiveData as SensitiveDataClientAuthKeys) || {};
 
   const webRedirectCheck = useWebRedirect(integration.data.redirectTo);
 
@@ -82,7 +93,7 @@ export const SigninPushCodeContainer = ({
     const navigationOptions = {
       email: signinState.email,
       signinData: { ...signinState, verified: true },
-      unwrapBKey: signinState.unwrapBKey,
+      unwrapBKey,
       integration,
       finishOAuthFlowHandler,
       queryParams: location.search,