fix(aws-amplify): cookie store is not used for identity store in SSR (#14326)

* fix(aws-amplify): cookie store is not used for identity store in SSR

* chore: resolve comments

* chore: resolve comments

Author: HuiSF

packages/aws-amplify/__tests__/initSingleton.test.ts

@@ -11,6 +11,8 @@ import {
 import { AmplifyOutputs } from '@aws-amplify/core/internals/utils';
 
 import {
+	CognitoAWSCredentialsAndIdentityIdProvider,
+	DefaultIdentityIdStore,
 	cognitoCredentialsProvider,
 	cognitoUserPoolsTokenProvider,
 } from '../src/auth/cognito';
@@ -23,6 +25,8 @@ jest.mock('../src/auth/cognito', () => ({
 		setKeyValueStorage: jest.fn(),
 	},
 	cognitoCredentialsProvider: jest.fn(),
+	DefaultIdentityIdStore: jest.fn(),
+	CognitoAWSCredentialsAndIdentityIdProvider: jest.fn(),
 }));
 
 const mockCognitoUserPoolsTokenProviderSetAuthConfig =
@@ -32,6 +36,10 @@ const mockCognitoUserPoolsTokenProviderSetKeyValueStorage =
 const mockAmplifySingletonConfigure = AmplifySingleton.configure as jest.Mock;
 const mockAmplifySingletonGetConfig = AmplifySingleton.getConfig as jest.Mock;
 const MockCookieStorage = CookieStorage as jest.Mock;
+const MockDefaultIdentityIdStore = jest.mocked(DefaultIdentityIdStore);
+const MockCognitoAWSCredentialsAndIdentityIdProvider = jest.mocked(
+	CognitoAWSCredentialsAndIdentityIdProvider,
+);
 
 const mockResourceConfig: ResourcesConfig = {
 	Auth: {
@@ -50,8 +58,16 @@ const mockResourceConfig: ResourcesConfig = {
 
 describe('initSingleton (DefaultAmplify)', () => {
 	const mockCookieStorageInstance = {};
+	const mockCognitoAWSCredentialsAndIdentityIdProviderInstance = {} as any;
+	const mockDefaultIdentityIdStoreInstance = {} as any;
 	beforeAll(() => {
 		MockCookieStorage.mockImplementation(() => mockCookieStorageInstance);
+		MockDefaultIdentityIdStore.mockImplementation(
+			() => mockDefaultIdentityIdStoreInstance,
+		);
+		MockCognitoAWSCredentialsAndIdentityIdProvider.mockImplementation(
+			() => mockCognitoAWSCredentialsAndIdentityIdProviderInstance,
+		);
 	});
 	beforeEach(() => {
 		mockAmplifySingletonConfigure.mockImplementation((_, libraryOptions) => {
@@ -64,6 +80,8 @@ describe('initSingleton (DefaultAmplify)', () => {
 
 	afterEach(() => {
 		MockCookieStorage.mockClear();
+		MockCognitoAWSCredentialsAndIdentityIdProvider.mockClear();
+		MockDefaultIdentityIdStore.mockClear();
 		mockCognitoUserPoolsTokenProviderSetAuthConfig.mockReset();
 		mockCognitoUserPoolsTokenProviderSetKeyValueStorage.mockReset();
 		mockAmplifySingletonConfigure.mockReset();
@@ -252,13 +270,20 @@ describe('initSingleton (DefaultAmplify)', () => {
 					expect(
 						mockCognitoUserPoolsTokenProviderSetKeyValueStorage,
 					).toHaveBeenCalledWith(mockCookieStorageInstance);
+					expect(MockDefaultIdentityIdStore).toHaveBeenCalledWith(
+						mockCookieStorageInstance,
+					);
+					expect(
+						MockCognitoAWSCredentialsAndIdentityIdProvider,
+					).toHaveBeenCalledWith(mockDefaultIdentityIdStoreInstance);
 					expect(mockAmplifySingletonConfigure).toHaveBeenCalledWith(
 						mockResourceConfig,
 						{
 							...libraryOptions,
 							Auth: {
 								tokenProvider: cognitoUserPoolsTokenProvider,
-								credentialsProvider: cognitoCredentialsProvider,
+								credentialsProvider:
+									mockCognitoAWSCredentialsAndIdentityIdProviderInstance,
 							},
 						},
 					);
@@ -345,7 +370,6 @@ describe('initSingleton (DefaultAmplify)', () => {
 					expect(
 						mockCognitoUserPoolsTokenProviderSetAuthConfig,
 					).not.toHaveBeenCalled();
-					expect(MockCookieStorage).not.toHaveBeenCalled();
 					expect(
 						mockCognitoUserPoolsTokenProviderSetKeyValueStorage,
 					).not.toHaveBeenCalled();

packages/aws-amplify/src/initSingleton.ts

@@ -14,6 +14,8 @@ import {
 } from '@aws-amplify/core/internals/utils';
 
 import {
+	CognitoAWSCredentialsAndIdentityIdProvider,
+	DefaultIdentityIdStore,
 	cognitoCredentialsProvider,
 	cognitoUserPoolsTokenProvider,
 } from './auth/cognito';
@@ -36,6 +38,15 @@ export const DefaultAmplify = {
 		libraryOptions?: LibraryOptions,
 	): void {
 		const resolvedResourceConfig = parseAmplifyConfig(resourceConfig);
+		const cookieBasedKeyValueStorage = new CookieStorage({ sameSite: 'lax' });
+		const resolvedKeyValueStorage = libraryOptions?.ssr
+			? cookieBasedKeyValueStorage
+			: defaultStorage;
+		const resolvedCredentialsProvider = libraryOptions?.ssr
+			? new CognitoAWSCredentialsAndIdentityIdProvider(
+					new DefaultIdentityIdStore(cookieBasedKeyValueStorage),
+				)
+			: cognitoCredentialsProvider;
 
 		// If no Auth config is provided, no special handling will be required, configure as is.
 		// Otherwise, we can assume an Auth config is provided from here on.
@@ -58,16 +69,14 @@ export const DefaultAmplify = {
 			cognitoUserPoolsTokenProvider.setAuthConfig(resolvedResourceConfig.Auth);
 			cognitoUserPoolsTokenProvider.setKeyValueStorage(
 				// TODO: allow configure with a public interface
-				libraryOptions?.ssr
-					? new CookieStorage({ sameSite: 'lax' })
-					: defaultStorage,
+				resolvedKeyValueStorage,
 			);
 
 			Amplify.configure(resolvedResourceConfig, {
 				...libraryOptions,
 				Auth: {
 					tokenProvider: cognitoUserPoolsTokenProvider,
-					credentialsProvider: cognitoCredentialsProvider,
+					credentialsProvider: resolvedCredentialsProvider,
 				},
 			});
 
@@ -77,18 +86,19 @@ export const DefaultAmplify = {
 		// At this point, Auth libraryOptions would have been previously configured and no overriding
 		// Auth options were given, so we should preserve the currently configured Auth libraryOptions.
 		if (libraryOptions) {
+			const authLibraryOptions = Amplify.libraryOptions.Auth;
 			// If ssr is provided through libraryOptions, we should respect the intentional reconfiguration.
 			if (libraryOptions.ssr !== undefined) {
 				cognitoUserPoolsTokenProvider.setKeyValueStorage(
 					// TODO: allow configure with a public interface
-					libraryOptions.ssr
-						? new CookieStorage({ sameSite: 'lax' })
-						: defaultStorage,
+					resolvedKeyValueStorage,
 				);
+
+				authLibraryOptions.credentialsProvider = resolvedCredentialsProvider;
 			}
 
 			Amplify.configure(resolvedResourceConfig, {
-				Auth: Amplify.libraryOptions.Auth,
+				Auth: authLibraryOptions,
 				...libraryOptions,
 			});