fix(auth): update service returned identity id to the in-memory store (#14268)

Author: HuiSF

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

@@ -69,18 +69,20 @@ describe('Guest Credentials', () => {
 
 	describe('Happy Path Cases:', () => {
 		beforeEach(() => {
+			const identityIdStore = new DefaultIdentityIdStore(sharedInMemoryStorage);
+			identityIdStore.setAuthConfig(validAuthConfig.Auth!);
 			cognitoCredentialsProvider =
-				new CognitoAWSCredentialsAndIdentityIdProvider(
-					new DefaultIdentityIdStore(sharedInMemoryStorage),
-				);
+				new CognitoAWSCredentialsAndIdentityIdProvider(identityIdStore);
 			credentialsForIdentityIdSpy.mockImplementationOnce(async () => {
 				return authAPITestParams.CredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;
 			});
 		});
+
 		afterEach(() => {
 			cognitoCredentialsProvider.clearCredentials();
 			credentialsForIdentityIdSpy?.mockReset();
 		});
+
 		test('Should call identityIdClient with no logins to obtain guest creds', async () => {
 			const res = await cognitoCredentialsProvider.getCredentialsAndIdentityId({
 				authenticated: false,
@@ -137,6 +139,7 @@ describe('Guest Credentials', () => {
 		afterAll(() => {
 			credentialsForIdentityIdSpy?.mockReset();
 		});
+
 		test('Should not throw AuthError when allowGuestAccess is false in the config', async () => {
 			expect(
 				await cognitoCredentialsProvider.getCredentialsAndIdentityId({
@@ -145,6 +148,7 @@ describe('Guest Credentials', () => {
 				}),
 			).toBe(undefined);
 		});
+
 		test('Should not throw AuthError when there is no Cognito object in the config', async () => {
 			expect(
 				await cognitoCredentialsProvider.getCredentialsAndIdentityId({
@@ -160,31 +164,47 @@ describe('Primary Credentials', () => {
 	let cognitoCredentialsProvider: CognitoAWSCredentialsAndIdentityIdProvider;
 	describe('Happy Path Cases:', () => {
 		beforeEach(() => {
+			const identityIdStore = new DefaultIdentityIdStore(sharedInMemoryStorage);
+			identityIdStore.setAuthConfig(validAuthConfig.Auth!);
 			cognitoCredentialsProvider =
-				new CognitoAWSCredentialsAndIdentityIdProvider(
-					new DefaultIdentityIdStore(sharedInMemoryStorage),
-				);
+				new CognitoAWSCredentialsAndIdentityIdProvider(identityIdStore);
 			credentialsForIdentityIdSpy.mockImplementation(async () => {
 				return authAPITestParams.CredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;
 			});
 		});
+
 		afterEach(() => {
 			cognitoCredentialsProvider.clearCredentials();
 			credentialsForIdentityIdSpy?.mockReset();
 		});
+
 		test('Should call identityIdClient with the logins map to obtain primary creds', async () => {
 			const res = await cognitoCredentialsProvider.getCredentialsAndIdentityId({
 				authenticated: true,
 				authConfig: validAuthConfig.Auth!,
 				tokens: authAPITestParams.ValidAuthTokens,
 			});
-			expect(res?.credentials.accessKeyId).toEqual(
-				authAPITestParams.CredentialsForIdentityIdResult.Credentials
-					.AccessKeyId,
-			);
+			expect(res).toMatchObject({
+				credentials: {
+					accessKeyId:
+						authAPITestParams.CredentialsForIdentityIdResult.Credentials
+							.AccessKeyId,
+					expiration:
+						authAPITestParams.CredentialsForIdentityIdResult.Credentials
+							.Expiration,
+					secretAccessKey:
+						authAPITestParams.CredentialsForIdentityIdResult.Credentials
+							.SecretKey,
+					sessionToken:
+						authAPITestParams.CredentialsForIdentityIdResult.Credentials
+							.SessionToken,
+				},
+				identityId: authAPITestParams.CredentialsForIdentityIdResult.IdentityId,
+			});
 
 			expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
 		});
+
 		test('in-memory primary creds are returned if not expired and not past TTL', async () => {
 			await cognitoCredentialsProvider.getCredentialsAndIdentityId({
 				authenticated: true,
@@ -211,6 +231,7 @@ describe('Primary Credentials', () => {
 			// expecting to be called only once becasue in-memory creds should be returned
 			expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
 		});
+
 		test('Should get new credentials when tokens have changed', async () => {
 			await cognitoCredentialsProvider.getCredentialsAndIdentityId({
 				authenticated: true,
@@ -239,19 +260,23 @@ describe('Primary Credentials', () => {
 			expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(2);
 		});
 	});
+
 	describe('Error Path Cases:', () => {
 		beforeEach(() => {
 			cognitoCredentialsProvider =
 				new CognitoAWSCredentialsAndIdentityIdProvider(
 					new DefaultIdentityIdStore(sharedInMemoryStorage),
 				);
 		});
+
 		afterEach(() => {
 			cognitoCredentialsProvider.clearCredentials();
 		});
+
 		afterAll(() => {
 			credentialsForIdentityIdSpy?.mockReset();
 		});
+
 		test('Should throw AuthError if either Credentials, accessKeyId or secretKey is absent in the response', async () => {
 			credentialsForIdentityIdSpy.mockImplementationOnce(async () => {
 				return authAPITestParams.NoAccessKeyCredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;

packages/auth/__tests__/providers/cognito/testUtils/authApiTestParams.ts

@@ -149,6 +149,8 @@ export const authAPITestParams = {
 			SessionToken: 'SessionToken',
 			Expiration: new Date('2023-07-29'),
 		},
+		IdentityId:
+			'I am expected to be returned as a part of the result as I am the source of truth',
 		$metadata: {},
 	},
 	NoAccessKeyCredentialsForIdentityIdResult: {

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

@@ -129,6 +129,7 @@ export class CognitoAWSCredentialsAndIdentityIdProvider
 			clientResult.Credentials.SecretKey
 		) {
 			this._nextCredentialsRefresh = new Date().getTime() + CREDENTIALS_TTL;
+
 			const res: CredentialsAndIdentityId = {
 				credentials: {
 					accessKeyId: clientResult.Credentials.AccessKeyId,
@@ -198,6 +199,8 @@ export class CognitoAWSCredentialsAndIdentityIdProvider
 			clientResult.Credentials.AccessKeyId &&
 			clientResult.Credentials.SecretKey
 		) {
+			this._nextCredentialsRefresh = new Date().getTime() + CREDENTIALS_TTL;
+
 			const res: CredentialsAndIdentityId = {
 				credentials: {
 					accessKeyId: clientResult.Credentials.AccessKeyId,
@@ -207,22 +210,23 @@ export class CognitoAWSCredentialsAndIdentityIdProvider
 				},
 				identityId,
 			};
-			// Store the credentials in-memory along with the expiration
-			this._credentialsAndIdentityId = {
-				...res,
-				isAuthenticatedCreds: true,
-				associatedIdToken: authTokens.idToken?.toString(),
-			};
-			this._nextCredentialsRefresh = new Date().getTime() + CREDENTIALS_TTL;
 
 			if (clientResult.IdentityId) {
 				res.identityId = clientResult.IdentityId;
+				// note: the following call removes guest identityId from the persistent store (localStorage)
 				this._identityIdStore.storeIdentityId({
 					id: clientResult.IdentityId,
 					type: 'primary',
 				});
 			}
 
+			// Store the credentials in-memory along with the expiration
+			this._credentialsAndIdentityId = {
+				...res,
+				isAuthenticatedCreds: true,
+				associatedIdToken: authTokens.idToken?.toString(),
+			};
+
 			return res;
 		} else {
 			throw new AuthError({