From 36de0bb25514c4724b573ff6810ba1b0203b022b Mon Sep 17 00:00:00 2001
From: Hui Zhao <zhohz@amazon.com>
Date: Wed, 8 Nov 2023 21:34:14 -0800
Subject: [PATCH] fix(auth): refreshToken unawaited async operation caused race
 condition

---
 .../tokenProvider/tokenOrchestrator.test.ts   | 76 +++++++++++++++++++
 .../tokenProvider/TokenOrchestrator.ts        |  2 +-
 2 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 packages/auth/__tests__/providers/cognito/tokenProvider/tokenOrchestrator.test.ts

diff --git a/packages/auth/__tests__/providers/cognito/tokenProvider/tokenOrchestrator.test.ts b/packages/auth/__tests__/providers/cognito/tokenProvider/tokenOrchestrator.test.ts
new file mode 100644
index 00000000000..a9bdbcaca73
--- /dev/null
+++ b/packages/auth/__tests__/providers/cognito/tokenProvider/tokenOrchestrator.test.ts
@@ -0,0 +1,76 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { Hub } from '@aws-amplify/core';
+import { assertTokenProviderConfig } from '@aws-amplify/core/internals/utils';
+import { tokenOrchestrator } from '../../../../src/providers/cognito/tokenProvider';
+import { CognitoAuthTokens } from '../../../../src/providers/cognito/tokenProvider/types';
+
+jest.mock('@aws-amplify/core/internals/utils');
+jest.mock('@aws-amplify/core', () => ({
+	...jest.requireActual('@aws-amplify/core'),
+	Hub: {
+		dispatch: jest.fn(),
+	},
+}));
+
+const mockAssertTokenProviderConfig = assertTokenProviderConfig as jest.Mock;
+
+describe('tokenOrchestrator', () => {
+	const mockTokenRefresher = jest.fn();
+	const mockTokenStore = {
+		storeTokens: jest.fn(),
+		getLastAuthUser: jest.fn(),
+		loadTokens: jest.fn(),
+		clearTokens: jest.fn(),
+		setKeyValueStorage: jest.fn(),
+		getDeviceMetadata: jest.fn(),
+		clearDeviceMetadata: jest.fn(),
+	};
+
+	beforeAll(() => {
+		tokenOrchestrator.setTokenRefresher(mockTokenRefresher);
+		tokenOrchestrator.setAuthTokenStore(mockTokenStore);
+	});
+
+	describe('refreshTokens method', () => {
+		it('calls the set tokenRefresher, tokenStore and Hub while refreshing tokens', async () => {
+			const testUsername = 'username';
+			const testInputTokens = {
+				accessToken: {
+					payload: {},
+				},
+				clockDrift: 400000,
+				username: testUsername,
+			};
+
+			const mockTokens: CognitoAuthTokens = {
+				accessToken: {
+					payload: {},
+				},
+				clockDrift: 300,
+				username: testUsername,
+			};
+			mockTokenRefresher.mockResolvedValueOnce(mockTokens);
+			mockTokenStore.storeTokens.mockResolvedValue(void 0);
+			const newTokens = await tokenOrchestrator['refreshTokens']({
+				tokens: testInputTokens,
+				username: testUsername,
+			});
+
+			// ensure the underlying async operations to be completed
+			// async #1
+			expect(mockTokenRefresher).toHaveBeenCalledWith(
+				expect.objectContaining({
+					tokens: testInputTokens,
+					username: testUsername,
+				})
+			);
+			// async #2
+			expect(mockTokenStore.storeTokens).toHaveBeenCalledWith(mockTokens);
+
+			// ensure the result is correct
+			expect(newTokens).toEqual(mockTokens);
+		});
+	});
+});
diff --git a/packages/auth/src/providers/cognito/tokenProvider/TokenOrchestrator.ts b/packages/auth/src/providers/cognito/tokenProvider/TokenOrchestrator.ts
index ac08d7731de..58b3b4dcad5 100644
--- a/packages/auth/src/providers/cognito/tokenProvider/TokenOrchestrator.ts
+++ b/packages/auth/src/providers/cognito/tokenProvider/TokenOrchestrator.ts
@@ -124,7 +124,7 @@ export class TokenOrchestrator implements AuthTokenOrchestrator {
 				username,
 			});
 
-			this.setTokens({ tokens: newTokens });
+			await this.setTokens({ tokens: newTokens });
 			Hub.dispatch('auth', { event: 'tokenRefresh' }, 'Auth', AMPLIFY_SYMBOL);
 
 			return newTokens;