chore(auth): add oauth metadata into token orchestrator

packages/auth/__tests__/providers/cognito/signOut.test.ts

@@ -220,6 +220,7 @@ describe('signOut', () => {
 			expect(mockHandleOAuthSignOut).toHaveBeenCalledWith(
 				cognitoConfigWithOauth,
 				mockDefaultOAuthStoreInstance,
+				mockTokenOrchestrator,
 			);
 			// In cases of OAuth, token removal and Hub dispatch should be performed by the OAuth handling since
 			// these actions can be deferred or canceled out of altogether.

packages/auth/__tests__/providers/cognito/tokenOrchestrator.test.ts

@@ -25,6 +25,8 @@ const mockAuthTokenStore = {
 	setKeyValueStorage: jest.fn(),
 	getDeviceMetadata: jest.fn(),
 	clearDeviceMetadata: jest.fn(),
+	setOAuthMetadata: jest.fn(),
+	getOAuthMetadata: jest.fn(),
 };
 const mockTokenRefresher = jest.fn();
 const validAuthConfig: ResourcesConfig = {

packages/auth/__tests__/providers/cognito/tokenProvider/tokenOrchestrator.test.ts

@@ -24,6 +24,8 @@ describe('tokenOrchestrator', () => {
 		setKeyValueStorage: jest.fn(),
 		getDeviceMetadata: jest.fn(),
 		clearDeviceMetadata: jest.fn(),
+		getOAuthMetadata: jest.fn(),
+		setOAuthMetadata: jest.fn(),
 	};
 
 	beforeAll(() => {

packages/auth/__tests__/providers/cognito/utils/oauth/completeOAuthFlow.test.ts

@@ -13,6 +13,7 @@ import { AuthErrorTypes } from '../../../../../src/types/Auth';
 import { OAuthStore } from '../../../../../src/providers/cognito/utils/types';
 import { completeOAuthFlow } from '../../../../../src/providers/cognito/utils/oauth/completeOAuthFlow';
 
+jest.mock('../../../../../src/providers/cognito/tokenProvider');
 jest.mock('@aws-amplify/core', () => ({
 	Hub: {
 		dispatch: jest.fn(),

packages/auth/__tests__/providers/cognito/utils/oauth/handleOAuthSignOut.test.ts

@@ -1,6 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+import { TokenOrchestrator } from '../../../../../src/providers/cognito';
 import { completeOAuthSignOut } from '../../../../../src/providers/cognito/utils/oauth/completeOAuthSignOut';
 import { handleOAuthSignOut } from '../../../../../src/providers/cognito/utils/oauth/handleOAuthSignOut';
 import { oAuthSignOutRedirect } from '../../../../../src/providers/cognito/utils/oauth/oAuthSignOutRedirect';
@@ -12,6 +13,7 @@ jest.mock(
 jest.mock(
 	'../../../../../src/providers/cognito/utils/oauth/oAuthSignOutRedirect',
 );
+jest.mock('../../../../../src/providers/cognito/tokenProvider');
 
 describe('handleOAuthSignOut', () => {
 	const region = 'us-west-2';
@@ -27,9 +29,13 @@ describe('handleOAuthSignOut', () => {
 	const mockStore = {
 		loadOAuthSignIn: jest.fn(),
 	} as unknown as jest.Mocked<DefaultOAuthStore>;
+	const mockTokenOrchestrator = {
+		getOAuthMetadata: jest.fn(),
+	} as unknown as jest.Mocked<TokenOrchestrator>;
 
 	afterEach(() => {
 		mockStore.loadOAuthSignIn.mockReset();
+		mockTokenOrchestrator.getOAuthMetadata.mockReset();
 		mockCompleteOAuthSignOut.mockClear();
 		mockOAuthSignOutRedirect.mockClear();
 	});
@@ -39,7 +45,21 @@ describe('handleOAuthSignOut', () => {
 			isOAuthSignIn: true,
 			preferPrivateSession: false,
 		});
-		await handleOAuthSignOut(cognitoConfig, mockStore);
+		await handleOAuthSignOut(cognitoConfig, mockStore, mockTokenOrchestrator);
+
+		expect(mockCompleteOAuthSignOut).toHaveBeenCalledWith(mockStore);
+		expect(mockOAuthSignOutRedirect).toHaveBeenCalledWith(cognitoConfig);
+	});
+
+	it('should complete OAuth sign out and redirect when there oauth metadata in tokenOrchestrator', async () => {
+		mockTokenOrchestrator.getOAuthMetadata.mockResolvedValue({
+			oauthSignIn: true,
+		});
+		mockStore.loadOAuthSignIn.mockResolvedValue({
+			isOAuthSignIn: false,
+			preferPrivateSession: false,
+		});
+		await handleOAuthSignOut(cognitoConfig, mockStore, mockTokenOrchestrator);
 
 		expect(mockCompleteOAuthSignOut).toHaveBeenCalledWith(mockStore);
 		expect(mockOAuthSignOutRedirect).toHaveBeenCalledWith(cognitoConfig);
@@ -50,7 +70,7 @@ describe('handleOAuthSignOut', () => {
 			isOAuthSignIn: false,
 			preferPrivateSession: false,
 		});
-		await handleOAuthSignOut(cognitoConfig, mockStore);
+		await handleOAuthSignOut(cognitoConfig, mockStore, mockTokenOrchestrator);
 
 		expect(mockCompleteOAuthSignOut).toHaveBeenCalledWith(mockStore);
 		expect(mockOAuthSignOutRedirect).not.toHaveBeenCalled();

packages/auth/src/providers/cognito/apis/signOut.ts

@@ -65,7 +65,11 @@ export async function signOut(input?: SignOutInput): Promise<void> {
 		const oAuthStore = new DefaultOAuthStore(defaultStorage);
 		oAuthStore.setAuthConfig(cognitoConfig);
 		const { type } =
-			(await handleOAuthSignOut(cognitoConfig, oAuthStore)) ?? {};
+			(await handleOAuthSignOut(
+				cognitoConfig,
+				oAuthStore,
+				tokenOrchestrator,
+			)) ?? {};
 		if (type === 'error') {
 			throw new AuthError({
 				name: OAUTH_SIGNOUT_EXCEPTION,

packages/auth/src/providers/cognito/tokenProvider/TokenOrchestrator.ts

@@ -25,6 +25,7 @@ import {
 	AuthTokenStore,
 	CognitoAuthTokens,
 	DeviceMetadata,
+	OAuthMetadata,
 	TokenRefresher,
 } from './types';
 
@@ -203,4 +204,12 @@ export class TokenOrchestrator implements AuthTokenOrchestrator {
 	clearDeviceMetadata(username?: string): Promise<void> {
 		return this.getTokenStore().clearDeviceMetadata(username);
 	}
+
+	setOAuthMetadata(metadata: OAuthMetadata): Promise<void> {
+		return this.getTokenStore().setOAuthMetadata(metadata);
+	}
+
+	getOAuthMetadata(): Promise<OAuthMetadata | null> {
+		return this.getTokenStore().getOAuthMetadata();
+	}
 }

packages/auth/src/providers/cognito/tokenProvider/TokenStore.ts

@@ -14,6 +14,7 @@ import {
 	AuthTokenStore,
 	CognitoAuthTokens,
 	DeviceMetadata,
+	OAuthMetadata,
 } from './types';
 import { TokenProviderErrorCode, assert } from './errorHelpers';
 
@@ -163,6 +164,7 @@ export class DefaultTokenStore implements AuthTokenStore {
 			this.getKeyValueStorage().removeItem(authKeys.refreshToken),
 			this.getKeyValueStorage().removeItem(authKeys.signInDetails),
 			this.getKeyValueStorage().removeItem(this.getLastAuthUserKey()),
+			this.getKeyValueStorage().removeItem(authKeys.oauthMetadata),
 		]);
 	}
 
@@ -222,6 +224,22 @@ export class DefaultTokenStore implements AuthTokenStore {
 
 		return lastAuthUser;
 	}
+
+	async setOAuthMetadata(metadata: OAuthMetadata): Promise<void> {
+		const { oauthMetadata: oauthMetadataKey } = await this.getAuthKeys();
+		await this.getKeyValueStorage().setItem(
+			oauthMetadataKey,
+			JSON.stringify(metadata),
+		);
+	}
+
+	async getOAuthMetadata(): Promise<OAuthMetadata | null> {
+		const { oauthMetadata: oauthMetadataKey } = await this.getAuthKeys();
+		const oauthMetadata =
+			await this.getKeyValueStorage().getItem(oauthMetadataKey);
+
+		return oauthMetadata && JSON.parse(oauthMetadata);
+	}
 }
 
 export const createKeysForAuthStorage = (

packages/auth/src/providers/cognito/tokenProvider/types.ts

@@ -34,6 +34,7 @@ export const AuthTokenStorageKeys = {
 	randomPasswordKey: 'randomPasswordKey',
 	deviceGroupKey: 'deviceGroupKey',
 	signInDetails: 'signInDetails',
+	oauthMetadata: 'oauthMetadata',
 };
 
 export interface AuthTokenStore {
@@ -44,6 +45,8 @@ export interface AuthTokenStore {
 	setKeyValueStorage(keyValueStorage: KeyValueStorageInterface): void;
 	getDeviceMetadata(username?: string): Promise<DeviceMetadata | null>;
 	clearDeviceMetadata(username?: string): Promise<void>;
+	setOAuthMetadata(metadata: OAuthMetadata): Promise<void>;
+	getOAuthMetadata(): Promise<OAuthMetadata | null>;
 }
 
 export interface AuthTokenOrchestrator {
@@ -58,6 +61,8 @@ export interface AuthTokenOrchestrator {
 	clearTokens(): Promise<void>;
 	getDeviceMetadata(username?: string): Promise<DeviceMetadata | null>;
 	clearDeviceMetadata(username?: string): Promise<void>;
+	setOAuthMetadata(metadata: OAuthMetadata): Promise<void>;
+	getOAuthMetadata(): Promise<OAuthMetadata | null>;
 }
 
 export interface CognitoUserPoolTokenProviderType extends TokenProvider {
@@ -78,3 +83,7 @@ export interface DeviceMetadata {
 	deviceGroupKey?: string;
 	randomPassword: string;
 }
+
+export interface OAuthMetadata {
+	oauthSignIn: boolean;
+}

packages/auth/src/providers/cognito/utils/oauth/completeOAuthFlow.ts

@@ -11,6 +11,7 @@ import { Hub, decodeJWT } from '@aws-amplify/core';
 
 import { cacheCognitoTokens } from '../../tokenProvider/cacheTokens';
 import { dispatchSignedInHubEvent } from '../dispatchSignedInHubEvent';
+import { tokenOrchestrator } from '../../tokenProvider';
 
 import { createOAuthError } from './createOAuthError';
 import { resolveAndClearInflightPromises } from './inflightPromise';
@@ -227,6 +228,9 @@ const completeFlow = async ({
 	redirectUri: string;
 	state: string;
 }) => {
+	await tokenOrchestrator.setOAuthMetadata({
+		oauthSignIn: true,
+	});
 	await oAuthStore.clearOAuthData();
 	await oAuthStore.storeOAuthSignIn(true, preferPrivateSession);
 

packages/auth/src/providers/cognito/utils/oauth/handleOAuthSignOut.ts

@@ -5,21 +5,30 @@ import { CognitoUserPoolConfig } from '@aws-amplify/core';
 
 import { OpenAuthSessionResult } from '../../../../utils/types';
 import { DefaultOAuthStore } from '../../utils/signInWithRedirectStore';
+import { TokenOrchestrator } from '../../tokenProvider';
 
 import { completeOAuthSignOut } from './completeOAuthSignOut';
 import { oAuthSignOutRedirect } from './oAuthSignOutRedirect';
 
 export const handleOAuthSignOut = async (
 	cognitoConfig: CognitoUserPoolConfig,
 	store: DefaultOAuthStore,
+	tokenOrchestrator: TokenOrchestrator,
 ): Promise<void | OpenAuthSessionResult> => {
 	const { isOAuthSignIn } = await store.loadOAuthSignIn();
+	const oauthMetadata = await tokenOrchestrator.getOAuthMetadata();
 
 	// Clear everything before attempting to visted logout endpoint since the current application
 	// state could be wiped away on redirect
 	await completeOAuthSignOut(store);
 
-	if (isOAuthSignIn) {
+	// The isOAuthSignIn flag is propagated by the oAuthToken store which manages oauth keys in local storage only.
+	// These keys are used to determine if a user is in an inflight or signedIn oauth states.
+	// However, this behavior represents an issue when 2 apps share the same set of tokens in Cookie storage because the app that didn't
+	// start the OAuth will not have access to the oauth keys.
+	// A heuristic solution is to add oauth metadata to the tokenOrchestrator which will have access to the underlying
+	// storage mechanism that is used by Amplify.
+	if (isOAuthSignIn || oauthMetadata?.oauthSignIn) {
 		// On web, this will always end up being a void action
 		return oAuthSignOutRedirect(cognitoConfig);
 	}