fix(auth): associate unAuth identityId to newly authenticated user's identityId

packages/auth/__tests__/providers/cognito/credentialsProvider.test.ts → packages/auth/__tests__/providers/cognito/credentialsProvider/credentialsProvider.test.ts

@@ -11,18 +11,17 @@ import {
 import {
 	CognitoAWSCredentialsAndIdentityIdProvider,
 	DefaultIdentityIdStore,
-} from '../../../src/providers/cognito';
-import { AuthError } from '../../../src/errors/AuthError';
-
-import { authAPITestParams } from './testUtils/authApiTestParams';
+} from '../../../../src/providers/cognito';
+import { AuthError } from '../../../../src/errors/AuthError';
+import { authAPITestParams } from '../testUtils/authApiTestParams';
 
 jest.mock('@aws-amplify/core', () => ({
 	...jest.requireActual('@aws-amplify/core'),
 	getCredentialsForIdentity: jest.fn(),
 }));
 
 jest.mock(
-	'./../../../src/providers/cognito/credentialsProvider/IdentityIdProvider',
+	'./../../../../src/providers/cognito/credentialsProvider/IdentityIdProvider',
 	() => ({
 		cognitoIdentityIdProvider: jest
 			.fn()

packages/auth/__tests__/providers/cognito/identityIdProvider.test.ts → packages/auth/__tests__/providers/cognito/credentialsProvider/identityIdProvider.test.ts

@@ -8,17 +8,18 @@ import {
 } from '@aws-amplify/core/internals/aws-clients/cognitoIdentity';
 import { CognitoIdentityPoolConfig } from '@aws-amplify/core/internals/utils';
 
-import { DefaultIdentityIdStore } from '../../../src/providers/cognito/credentialsProvider/IdentityIdStore';
-import { cognitoIdentityIdProvider } from '../../../src/providers/cognito/credentialsProvider/IdentityIdProvider';
-
-import { authAPITestParams } from './testUtils/authApiTestParams';
+import { DefaultIdentityIdStore } from '../../../../src/providers/cognito/credentialsProvider/IdentityIdStore';
+import { cognitoIdentityIdProvider } from '../../../../src/providers/cognito/credentialsProvider/IdentityIdProvider';
+import { authAPITestParams } from '../testUtils/authApiTestParams';
 
 jest.mock('@aws-amplify/core', () => ({
 	...jest.requireActual('@aws-amplify/core'),
 	getId: jest.fn(),
 }));
 jest.mock('@aws-amplify/core/internals/aws-clients/cognitoIdentity');
-jest.mock('../../../src/providers/cognito/credentialsProvider/IdentityIdStore');
+jest.mock(
+	'../../../../src/providers/cognito/credentialsProvider/IdentityIdStore',
+);
 
 const ampConfig: ResourcesConfig = {
 	Auth: {
@@ -140,4 +141,59 @@ describe('Cognito IdentityId Provider Happy Path Cases:', () => {
 		).toBe(authAPITestParams.PrimaryIdentityId.id);
 		expect(mockGetId).toHaveBeenCalledTimes(1);
 	});
+	test('Should return the identityId irresspective of the type if present', async () => {
+		mockDefaultIdentityIdStoreInstance.loadIdentityId.mockImplementationOnce(
+			async () => {
+				return authAPITestParams.PrimaryIdentityId as Identity;
+			},
+		);
+		expect(
+			await cognitoIdentityIdProvider({
+				tokens: authAPITestParams.ValidAuthTokens,
+				authConfig: {
+					identityPoolId: 'XXXXXXXXXXXXXXXXX',
+				},
+				identityIdStore: mockDefaultIdentityIdStoreInstance,
+			}),
+		).toBe(authAPITestParams.PrimaryIdentityId.id);
+
+		mockDefaultIdentityIdStoreInstance.loadIdentityId.mockImplementationOnce(
+			async () => {
+				return authAPITestParams.GuestIdentityId as Identity;
+			},
+		);
+		expect(
+			await cognitoIdentityIdProvider({
+				tokens: authAPITestParams.ValidAuthTokens,
+				authConfig: {
+					identityPoolId: 'XXXXXXXXXXXXXXXXX',
+				},
+				identityIdStore: mockDefaultIdentityIdStoreInstance,
+			}),
+		).toBe(authAPITestParams.GuestIdentityId.id);
+		expect(mockGetId).toHaveBeenCalledTimes(0);
+	});
+	test('Should fetch from Cognito when there is no identityId cached', async () => {
+		mockDefaultIdentityIdStoreInstance.loadIdentityId.mockImplementationOnce(
+			async () => {
+				return undefined;
+			},
+		);
+		mockDefaultIdentityIdStoreInstance.storeIdentityId.mockImplementationOnce(
+			async (identity: Identity) => {
+				expect(identity.id).toBe(authAPITestParams.PrimaryIdentityId.id);
+				expect(identity.type).toBe(authAPITestParams.PrimaryIdentityId.type);
+			},
+		);
+		expect(
+			await cognitoIdentityIdProvider({
+				tokens: authAPITestParams.ValidAuthTokens,
+				authConfig: {
+					identityPoolId: 'us-east-1:test-id',
+				},
+				identityIdStore: mockDefaultIdentityIdStoreInstance,
+			}),
+		).toBe(authAPITestParams.PrimaryIdentityId.id);
+		expect(mockGetId).toHaveBeenCalledTimes(1);
+	});
 });

packages/auth/src/providers/cognito/credentialsProvider/IdentityIdProvider.ts

@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { AuthTokens, ConsoleLogger, Identity, getId } from '@aws-amplify/core';
+import { AuthTokens, Identity, getId } from '@aws-amplify/core';
 import { CognitoIdentityPoolConfig } from '@aws-amplify/core/internals/utils';
 
 import { AuthError } from '../../../errors/AuthError';
@@ -11,7 +11,6 @@ import { GetIdException } from '../types/errors';
 import { IdentityIdStore } from './types';
 import { formLoginsMap } from './utils';
 
-const logger = new ConsoleLogger('CognitoIdentityIdProvider');
 /**
  * Provides a Cognito identityId
  *
@@ -33,46 +32,22 @@ export async function cognitoIdentityIdProvider({
 	identityIdStore.setAuthConfig({ Cognito: authConfig });
 
 	// will return null only if there is no identityId cached or if there is an error retrieving it
-	let identityId: Identity | null = await identityIdStore.loadIdentityId();
+	const identityId: Identity | null = await identityIdStore.loadIdentityId();
 
-	// Tokens are available so return primary identityId
-	if (tokens) {
-		// If there is existing primary identityId in-memory return that
-		if (identityId && identityId.type === 'primary') {
-			return identityId.id;
-		} else {
-			const logins = tokens.idToken
-				? formLoginsMap(tokens.idToken.toString())
-				: {};
-
-			const generatedIdentityId = await generateIdentityId(logins, authConfig);
-
-			if (identityId && identityId.id === generatedIdentityId) {
-				logger.debug(
-					`The guest identity ${identityId.id} has become the primary identity.`,
-				);
-			}
-			identityId = {
-				id: generatedIdentityId,
-				type: 'primary',
-			};
-		}
-	} else {
-		// If there is existing guest identityId cached return that
-		if (identityId && identityId.type === 'guest') {
-			return identityId.id;
-		} else {
-			identityId = {
-				id: await generateIdentityId({}, authConfig),
-				type: 'guest',
-			};
-		}
+	if (identityId) {
+		return identityId.id;
 	}
+	const logins = tokens?.idToken
+		? formLoginsMap(tokens.idToken.toString())
+		: {};
+	const generatedIdentityId = await generateIdentityId(logins, authConfig);
+	// Store generated identityId
+	identityIdStore.storeIdentityId({
+		id: generatedIdentityId,
+		type: tokens ? 'primary' : 'guest',
+	});
 
-	// Store in-memory or local storage depending on guest or primary identityId
-	identityIdStore.storeIdentityId(identityId);
-
-	return identityId.id;
+	return generatedIdentityId;
 }
 
 async function generateIdentityId(

packages/auth/src/providers/cognito/credentialsProvider/credentialsProvider.ts

@@ -138,11 +138,10 @@ export class CognitoAWSCredentialsAndIdentityIdProvider
 				},
 				identityId,
 			};
-			const identityIdRes = clientResult.IdentityId;
-			if (identityIdRes) {
-				res.identityId = identityIdRes;
+			if (clientResult.IdentityId) {
+				res.identityId = clientResult.IdentityId;
 				this._identityIdStore.storeIdentityId({
-					id: identityIdRes,
+					id: clientResult.IdentityId,
 					type: 'guest',
 				});
 			}
@@ -216,11 +215,10 @@ export class CognitoAWSCredentialsAndIdentityIdProvider
 			};
 			this._nextCredentialsRefresh = new Date().getTime() + CREDENTIALS_TTL;
 
-			const identityIdRes = clientResult.IdentityId;
-			if (identityIdRes) {
-				res.identityId = identityIdRes;
+			if (clientResult.IdentityId) {
+				res.identityId = clientResult.IdentityId;
 				this._identityIdStore.storeIdentityId({
-					id: identityIdRes,
+					id: clientResult.IdentityId,
 					type: 'primary',
 				});
 			}