From 4f5f749e18bb9c131e7083246b468ff2f3e86cde Mon Sep 17 00:00:00 2001
From: israx <70438514+israx@users.noreply.github.com>
Date: Tue, 24 Oct 2023 15:05:38 -0400
Subject: [PATCH] fix(auth): clear oauth data before initiating a new oauth
 flow (#12409)

* fix: call clearOauthData before starting the oauthf flow

* chore: addess pr feedback

* chore: address feedback
---
 .../providers/cognito/getCurrentUser.test.ts  |  3 --
 .../cognito/signInWithRedirect.test.ts        | 43 ++++++++++++++++++-
 .../cognito/apis/signInWithRedirect.ts        | 14 ++++--
 3 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/packages/auth/__tests__/providers/cognito/getCurrentUser.test.ts b/packages/auth/__tests__/providers/cognito/getCurrentUser.test.ts
index d30a53c24d2..84ea1fc6cd4 100644
--- a/packages/auth/__tests__/providers/cognito/getCurrentUser.test.ts
+++ b/packages/auth/__tests__/providers/cognito/getCurrentUser.test.ts
@@ -5,9 +5,6 @@ import { Amplify } from 'aws-amplify';
 import { decodeJWT } from '@aws-amplify/core/internals/utils';
 import { AuthError } from '../../../src/errors/AuthError';
 import { getCurrentUser } from '../../../src/providers/cognito';
-import { InitiateAuthException } from '../../../src/providers/cognito/types/errors';
-import { fetchTransferHandler } from '@aws-amplify/core/internals/aws-client-utils';
-import { buildMockErrorResponse, mockJsonResponse } from './testUtils/data';
 import { Amplify as AmplifyV6 } from '@aws-amplify/core';
 import { USER_UNAUTHENTICATED_EXCEPTION } from '../../../src/errors/constants';
 jest.mock('@aws-amplify/core/lib/clients/handlers/fetch');
diff --git a/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts b/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts
index c2a25cc6b73..2f4a14f2074 100644
--- a/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts
+++ b/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts
@@ -1,5 +1,44 @@
-describe.skip('signInWithRedirect API', () => {
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { Amplify } from '@aws-amplify/core';
+
+import * as signInWithRedirectModule from '../../../src/providers/cognito/apis/signInWithRedirect';
+Amplify.configure({
+	Auth: {
+		Cognito: {
+			userPoolClientId: '111111-aaaaa-42d8-891d-ee81a1549398',
+			userPoolId: 'us-west-2_zzzzz',
+			loginWith: {
+				oauth: {
+					domain: 'mydomain.com',
+					redirectSignIn: ['localHost:3000'],
+					redirectSignOut: ['localHost:3000'],
+					responseType: 'code',
+					scopes: ['aws.cognito.signin.user.admin'],
+				},
+			},
+		},
+	},
+});
+
+describe('signInWithRedirect API', () => {
+	let oauthSignInSpy;
+	beforeEach(() => {
+		oauthSignInSpy = jest
+			.spyOn(signInWithRedirectModule, 'oauthSignIn')
+			.mockImplementationOnce(async () => {
+				return {};
+			});
+	});
+	afterEach(() => {
+		oauthSignInSpy.mockClear();
+	});
 	it('should pass correct arguments to oauth', () => {
-		// ADD tests
+		// TODO ADD tests
+	});
+
+	it('should try to clear oauth data before starting an oauth flow.', async () => {
+		// TODO: ADD Test: previous test was invalid
 	});
 });
diff --git a/packages/auth/src/providers/cognito/apis/signInWithRedirect.ts b/packages/auth/src/providers/cognito/apis/signInWithRedirect.ts
index 9d7568ef910..cc1fdf946bc 100644
--- a/packages/auth/src/providers/cognito/apis/signInWithRedirect.ts
+++ b/packages/auth/src/providers/cognito/apis/signInWithRedirect.ts
@@ -33,7 +33,8 @@ import { getCurrentUser } from './getCurrentUser';
  *
  * @param input - The SignInWithRedirectInput object, if empty it will redirect to Cognito HostedUI
  *
- * TODO: add config errors
+ * @throws AuthTokenConfigException - Thrown when the userpool config is invalid.
+ * @throws OAuthNotConfigureException - Thrown when the oauth config is invalid.
  */
 export async function signInWithRedirect(
 	input?: SignInWithRedirectInput
@@ -61,9 +62,9 @@ export async function signInWithRedirect(
 	});
 }
 
-const store = new DefaultOAuthStore(defaultStorage);
+export const store = new DefaultOAuthStore(defaultStorage);
 
-async function oauthSignIn({
+export async function oauthSignIn({
 	oauthConfig,
 	provider,
 	clientId,
@@ -162,6 +163,7 @@ async function handleCodeFlow({
 	const code = url.searchParams.get('code');
 
 	if (!code) {
+		await store.clearOAuthData();
 		return;
 	}
 
@@ -258,8 +260,11 @@ async function handleImplicitFlow({
 			tokenType: undefined,
 			expiresIn: undefined,
 		});
+	if (!idToken || !accessToken) {
+		await store.clearOAuthData();
+		return;
+	}
 
-	await store.clearOAuthInflightData();
 	try {
 		validateState(state);
 	} catch (error) {
@@ -290,6 +295,7 @@ async function completeFlow({
 	redirectUri: string;
 	state: string;
 }) {
+	await store.clearOAuthData();
 	await store.storeOAuthSignIn(true, preferPrivateSession);
 	if (isCustomState(state)) {
 		Hub.dispatch(