fix(auth): oAuthStore is used before a valid OAuth config is confirmed

aws-amplify · Dec 23, 2023 · 3ce6701 · 3ce6701
1 parent 94d1fb2
commit 3ce6701
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 50 deletions.
diff --git a/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts b/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts
@@ -21,9 +21,7 @@ import {
 	completeOAuthFlow,
 } from '../../../src/providers/cognito/utils/oauth';
 import { attemptCompleteOAuthFlow } from '../../../src/providers/cognito/utils/oauth/attemptCompleteOAuthFlow';
-import { addInflightPromise } from '../../../src/providers/cognito/utils/oauth/inflightPromise';
 import { createOAuthError } from '../../../src/providers/cognito/utils/oauth/createOAuthError';
-import { cognitoUserPoolsTokenProvider } from '../../../src/providers/cognito/tokenProvider/tokenProvider';
 
 import { signInWithRedirect } from '../../../src/providers/cognito/apis/signInWithRedirect';
 
@@ -72,18 +70,8 @@ jest.mock('../../../src/providers/cognito/utils/oauth/oAuthStore', () => ({
 		clearOAuthInflightData: jest.fn(),
 	} as OAuthStore,
 }));
-jest.mock('../../../src/providers/cognito/utils/oauth/inflightPromise', () => ({
-	addInflightPromise: jest.fn(resolver => {
-		resolver();
-	}),
-}));
 jest.mock('../../../src/providers/cognito/utils/oauth/createOAuthError');
 jest.mock('../../../src/utils');
-jest.mock('../../../src/providers/cognito/tokenProvider/tokenProvider', () => ({
-	cognitoUserPoolsTokenProvider: {
-		setWaitForInflightOAuth: jest.fn(),
-	},
-}));
 
 const mockAssertOAuthConfig = assertOAuthConfig as jest.Mock;
 const mockAssertTokenProviderConfig = assertTokenProviderConfig as jest.Mock;
@@ -93,7 +81,6 @@ const mockOpenAuthSession = openAuthSession as jest.Mock;
 const mockGenerateCodeVerifier = generateCodeVerifier as jest.Mock;
 const mockGenerateState = generateState as jest.Mock;
 const mockIsBrowser = isBrowser as jest.Mock;
-const mockAddInflightPromise = addInflightPromise as jest.Mock;
 
 const mockCompleteOAuthFlow = completeOAuthFlow as jest.Mock;
 const mockGetAuthUserAgentValue = getAuthUserAgentValue as jest.Mock;
@@ -199,31 +186,6 @@ describe('signInWithRedirect', () => {
 				expect(Amplify[ADD_OAUTH_LISTENER]).toHaveBeenCalledWith(
 					attemptCompleteOAuthFlow
 				);
-				expect(
-					cognitoUserPoolsTokenProvider.setWaitForInflightOAuth
-				).toHaveBeenCalledTimes(1);
-
-				const callback = (
-					cognitoUserPoolsTokenProvider.setWaitForInflightOAuth as jest.Mock
-				).mock.calls[0][0];
-
-				await callback();
-
-				expect(oAuthStore.loadOAuthInFlight).toHaveBeenCalledTimes(1);
-			});
-
-			it('adds a promise to block auth process when there is an inflight oauth process', async () => {
-				(oAuthStore.loadOAuthInFlight as jest.Mock).mockResolvedValueOnce(true);
-				expect(
-					cognitoUserPoolsTokenProvider.setWaitForInflightOAuth
-				).toHaveBeenCalledTimes(1);
-
-				const callback = (
-					cognitoUserPoolsTokenProvider.setWaitForInflightOAuth as jest.Mock
-				).mock.calls[0][0];
-
-				await callback();
-				expect(mockAddInflightPromise).toHaveBeenCalledTimes(1);
 			});
 		});
 	});

diff --git a/packages/auth/__tests__/providers/cognito/utils/oauth/attemptCompleteOAuthFlow.test.ts b/packages/auth/__tests__/providers/cognito/utils/oauth/attemptCompleteOAuthFlow.test.ts
@@ -10,6 +10,8 @@ import { attemptCompleteOAuthFlow } from '../../../../../src/providers/cognito/u
 import { completeOAuthFlow } from '../../../../../src/providers/cognito/utils/oauth/completeOAuthFlow';
 import { getRedirectUrl } from '../../../../../src/providers/cognito/utils/oauth/getRedirectUrl';
 import { oAuthStore } from '../../../../../src/providers/cognito/utils/oauth/oAuthStore';
+import { addInflightPromise } from '../../../../../src/providers/cognito/utils/oauth/inflightPromise';
+import { cognitoUserPoolsTokenProvider } from '../../../../../src/providers/cognito/tokenProvider/tokenProvider';
 import { mockAuthConfigWithOAuth } from '../../../../mockData';
 
 import type { OAuthStore } from '../../../../../src/providers/cognito/utils/types';
@@ -35,11 +37,28 @@ jest.mock(
 		} as OAuthStore,
 	})
 );
+jest.mock(
+	'../../../../../src/providers/cognito/tokenProvider/tokenProvider',
+	() => ({
+		cognitoUserPoolsTokenProvider: {
+			setWaitForInflightOAuth: jest.fn(),
+		},
+	})
+);
+jest.mock(
+	'../../../../../src/providers/cognito/utils/oauth/inflightPromise',
+	() => ({
+		addInflightPromise: jest.fn(resolver => {
+			resolver();
+		}),
+	})
+);
 
 const mockAssertOAuthConfig = assertOAuthConfig as jest.Mock;
 const mockAssertTokenProviderConfig = assertTokenProviderConfig as jest.Mock;
 const mockCompleteOAuthFlow = completeOAuthFlow as jest.Mock;
 const mockGetRedirectUrl = getRedirectUrl as jest.Mock;
+const mockAddInflightPromise = addInflightPromise as jest.Mock;
 
 describe('attemptCompleteOAuthFlow', () => {
 	let windowSpy = jest.spyOn(window, 'window', 'get');
@@ -90,7 +109,7 @@ describe('attemptCompleteOAuthFlow', () => {
 		expect(mockCompleteOAuthFlow).not.toHaveBeenCalled();
 	});
 
-	it('invokes `completeOAuthFlow` to complete an inflight oauth process', async () => {
+	it('sets inflight oauth promise and invokes `completeOAuthFlow` to complete an inflight oauth process', async () => {
 		(oAuthStore.loadOAuthInFlight as jest.Mock).mockResolvedValueOnce(true);
 
 		await attemptCompleteOAuthFlow(mockAuthConfigWithOAuth.Auth.Cognito);
@@ -101,6 +120,17 @@ describe('attemptCompleteOAuthFlow', () => {
 				redirectUri: 'http://localhost:3000/',
 			})
 		);
+
+		expect(
+			cognitoUserPoolsTokenProvider.setWaitForInflightOAuth
+		).toHaveBeenCalledTimes(1);
+
+		const callback = (
+			cognitoUserPoolsTokenProvider.setWaitForInflightOAuth as jest.Mock
+		).mock.calls[0][0];
+
+		// the blocking promise should resolve
+		expect(callback()).resolves.toBeUndefined();
 	});
 
 	test.each([

diff --git a/packages/auth/src/providers/cognito/utils/oauth/attemptCompleteOAuthFlow.ts b/packages/auth/src/providers/cognito/utils/oauth/attemptCompleteOAuthFlow.ts
@@ -12,6 +12,8 @@ import { completeOAuthFlow } from './completeOAuthFlow';
 import { getAuthUserAgentValue } from '../../../../utils';
 import { getRedirectUrl } from './getRedirectUrl';
 import { handleFailure } from './handleFailure';
+import { cognitoUserPoolsTokenProvider } from '../../tokenProvider';
+import { addInflightPromise } from './inflightPromise';
 
 export const attemptCompleteOAuthFlow = async (
 	authConfig: AuthConfig['Cognito']
@@ -32,6 +34,16 @@ export const attemptCompleteOAuthFlow = async (
 		return;
 	}
 
+	// when there is valid oauth config and there is an inflight oauth flow, try
+	// to block async calls that require fetching tokens before the oauth flow completes
+	// e.g. getCurrentUser, fetchAuthSession etc.
+	cognitoUserPoolsTokenProvider.setWaitForInflightOAuth(
+		() =>
+			new Promise(async (res, _rej) => {
+				addInflightPromise(res);
+			})
+	);
+
 	try {
 		const currentUrl = window.location.href;
 		const { loginWith, userPoolClientId } = authConfig;

diff --git a/packages/auth/src/providers/cognito/utils/oauth/enableOAuthListener.ts b/packages/auth/src/providers/cognito/utils/oauth/enableOAuthListener.ts
@@ -17,17 +17,6 @@ isBrowser() &&
 	(() => {
 		// add the listener to the singleton for triggering
 		Amplify[ADD_OAUTH_LISTENER](attemptCompleteOAuthFlow);
-
-		cognitoUserPoolsTokenProvider.setWaitForInflightOAuth(
-			() =>
-				new Promise(async (res, _rej) => {
-					if (!(await oAuthStore.loadOAuthInFlight())) {
-						res();
-					} else {
-						addInflightPromise(res);
-					}
-				})
-		);
 	})();
 
 // required to present for module loaders