feat(aws-amplify|adapter-nextjs): add runtimeOptions.cookies to creat…

…eServerRunner
aws-amplify · Sep 11, 2024 · 6cbdad1 · 6cbdad1
1 parent a76b594
commit 6cbdad1
Show file tree

Hide file tree

Showing 6 changed files with 139 additions and 19 deletions.
diff --git a/packages/adapter-nextjs/__tests__/createServerRunner.test.ts b/packages/adapter-nextjs/__tests__/createServerRunner.test.ts
@@ -24,12 +24,22 @@ const mockAmplifyConfig: ResourcesConfig = {
 jest.mock(
 	'../src/utils/createCookieStorageAdapterFromNextServerContext',
 	() => ({
-		createCookieStorageAdapterFromNextServerContext: jest.fn(),
+		createCookieStorageAdapterFromNextServerContext: jest.fn(() => ({
+			get: jest.fn(),
+			set: jest.fn(),
+			delete: jest.fn(),
+			getAll: jest.fn(),
+		})),
 	}),
 );
 
+jest.mock('../src/utils/createTokenValidator', () => ({
+	createTokenValidator: jest.fn(() => ({
+		getItem: jest.fn(),
+	})),
+}));
 describe('createServerRunner', () => {
-	let createServerRunner: any;
+	let createServerRunner: NextServer.CreateServerRunner;
 
 	const mockParseAmplifyConfig = jest.fn();
 	const mockCreateAWSCredentialsAndIdentityIdProvider = jest.fn();
@@ -124,26 +134,28 @@ describe('createServerRunner', () => {
 			});
 
 			describe('when nextServerContext is not null', () => {
+				const mockNextServerContext = {
+					req: {
+						headers: {
+							cookie: 'cookie',
+						},
+					},
+					res: {
+						setHeader: jest.fn(),
+					},
+				};
+				const mockCookieStorageAdapter = {
+					get: jest.fn(),
+					set: jest.fn(),
+					remove: jest.fn(),
+				};
+
 				it('should create auth providers with cookie storage adapter', () => {
 					const operation = jest.fn();
-					const mockCookieStorageAdapter = {
-						get: jest.fn(),
-						set: jest.fn(),
-						remove: jest.fn(),
-					};
+
 					mockCreateKeyValueStorageFromCookieStorageAdapter.mockReturnValueOnce(
 						mockCookieStorageAdapter,
 					);
-					const mockNextServerContext = {
-						req: {
-							headers: {
-								cookie: 'cookie',
-							},
-						},
-						res: {
-							setHeader: jest.fn(),
-						},
-					};
 					const { runWithAmplifyServerContext } = createServerRunner({
 						config: mockAmplifyConfig,
 					});
@@ -163,6 +175,54 @@ describe('createServerRunner', () => {
 						mockCookieStorageAdapter,
 					);
 				});
+
+				it('should call createKeyValueStorageFromCookieStorageAdapter with specified runtimeOptions.cookies', () => {
+					const testCookiesOptions: NextServer.CreateServerRunnerRuntimeOptions['cookies'] =
+						{
+							domain: '.example.com',
+							sameSite: 'lax',
+							expires: new Date('2024-09-05'),
+						};
+					mockCreateKeyValueStorageFromCookieStorageAdapter.mockReturnValueOnce(
+						mockCookieStorageAdapter,
+					);
+
+					const { runWithAmplifyServerContext } = createServerRunner({
+						config: mockAmplifyConfig,
+						runtimeOptions: {
+							cookies: testCookiesOptions,
+						},
+					});
+
+					runWithAmplifyServerContext({
+						nextServerContext:
+							mockNextServerContext as unknown as NextServer.Context,
+						operation: jest.fn(),
+					});
+
+					expect(
+						mockCreateKeyValueStorageFromCookieStorageAdapter,
+					).toHaveBeenCalledWith(
+						expect.any(Object),
+						expect.any(Object),
+						testCookiesOptions,
+					);
+
+					// modify by reference should not affect the original configuration
+					testCookiesOptions.sameSite = 'strict';
+					runWithAmplifyServerContext({
+						nextServerContext:
+							mockNextServerContext as unknown as NextServer.Context,
+						operation: jest.fn(),
+					});
+
+					expect(
+						mockCreateKeyValueStorageFromCookieStorageAdapter,
+					).toHaveBeenCalledWith(expect.any(Object), expect.any(Object), {
+						...testCookiesOptions,
+						sameSite: 'lax',
+					});
+				});
 			});
 		});
 	});

diff --git a/packages/adapter-nextjs/src/createServerRunner.ts b/packages/adapter-nextjs/src/createServerRunner.ts
@@ -27,12 +27,14 @@ import { NextServer } from './types';
  */
 export const createServerRunner: NextServer.CreateServerRunner = ({
 	config,
+	runtimeOptions,
 }) => {
 	const amplifyConfig = parseAmplifyConfig(config);
 
 	return {
 		runWithAmplifyServerContext: createRunWithAmplifyServerContext({
 			config: amplifyConfig,
+			runtimeOptions,
 		}),
 	};
 };
diff --git a/packages/adapter-nextjs/src/types/NextServer.ts b/packages/adapter-nextjs/src/types/NextServer.ts
@@ -5,7 +5,10 @@ import { GetServerSidePropsContext as NextGetServerSidePropsContext } from 'next
 import { NextRequest, NextResponse } from 'next/server.js';
 import { cookies } from 'next/headers.js';
 import { AmplifyOutputs, LegacyConfig } from 'aws-amplify/adapter-core';
-import { AmplifyServer } from '@aws-amplify/core/internals/adapter-core';
+import {
+	AmplifyServer,
+	CookieStorage,
+} from '@aws-amplify/core/internals/adapter-core';
 import { ResourcesConfig } from '@aws-amplify/core';
 
 export declare namespace NextServer {
@@ -73,8 +76,16 @@ export declare namespace NextServer {
 		input: RunWithContextInput<OperationResult>,
 	) => Promise<OperationResult>;
 
+	export interface CreateServerRunnerRuntimeOptions {
+		cookies?: Pick<
+			CookieStorage.SetCookieOptions,
+			'domain' | 'expires' | 'sameSite'
+		>;
+	}
+
 	export interface CreateServerRunnerInput {
 		config: ResourcesConfig | LegacyConfig | AmplifyOutputs;
+		runtimeOptions?: CreateServerRunnerRuntimeOptions;
 	}
 
 	export interface CreateServerRunnerOutput {

diff --git a/packages/adapter-nextjs/src/utils/createRunWithAmplifyServerContext.ts b/packages/adapter-nextjs/src/utils/createRunWithAmplifyServerContext.ts
@@ -16,9 +16,14 @@ import { createCookieStorageAdapterFromNextServerContext } from './createCookieS
 
 export const createRunWithAmplifyServerContext = ({
 	config: resourcesConfig,
+	runtimeOptions = { cookies: {} },
 }: {
 	config: ResourcesConfig;
+	runtimeOptions?: NextServer.CreateServerRunnerRuntimeOptions;
 }) => {
+	const setCookieOptions = {
+		...runtimeOptions.cookies,
+	};
 	const runWithAmplifyServerContext: NextServer.RunOperationWithContext =
 		async ({ nextServerContext, operation }) => {
 			// When the Auth config is presented, attempt to create a Amplify server
@@ -40,6 +45,7 @@ export const createRunWithAmplifyServerContext = ({
 									userPoolClientId:
 										resourcesConfig?.Auth.Cognito?.userPoolClientId,
 								}),
+								setCookieOptions,
 							);
 				const credentialsProvider = createAWSCredentialsAndIdentityIdProvider(
 					resourcesConfig.Auth,

diff --git a/...ests__/adapterCore/storageFactories/createKeyValueStorageFromCookieStorageAdapter.test.ts b/...ests__/adapterCore/storageFactories/createKeyValueStorageFromCookieStorageAdapter.test.ts
@@ -1,7 +1,10 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { createKeyValueStorageFromCookieStorageAdapter } from '../../../src/adapter-core';
+import {
+	CookieStorage,
+	createKeyValueStorageFromCookieStorageAdapter,
+} from '../../../src/adapter-core';
 import { defaultSetCookieOptions } from '../../../src/adapter-core/storageFactories/createKeyValueStorageFromCookieStorageAdapter';
 
 const mockCookiesStorageAdapter = {
@@ -12,6 +15,13 @@ const mockCookiesStorageAdapter = {
 };
 
 describe('keyValueStorage', () => {
+	afterEach(() => {
+		mockCookiesStorageAdapter.delete.mockClear();
+		mockCookiesStorageAdapter.get.mockClear();
+		mockCookiesStorageAdapter.set.mockClear();
+		mockCookiesStorageAdapter.getAll.mockClear();
+	});
+
 	describe('createKeyValueStorageFromCookiesStorageAdapter', () => {
 		it('should return a key value storage', () => {
 			const keyValueStorage = createKeyValueStorageFromCookieStorageAdapter(
@@ -100,6 +110,31 @@ describe('keyValueStorage', () => {
 			});
 		});
 
+		describe('passing setCookieOptions parameter', () => {
+			const testSetCookieOptions: CookieStorage.SetCookieOptions = {
+				httpOnly: true,
+				sameSite: 'strict',
+				expires: new Date('2024-09-05'),
+			};
+			const keyValueStorage = createKeyValueStorageFromCookieStorageAdapter(
+				mockCookiesStorageAdapter,
+				undefined,
+				testSetCookieOptions,
+			);
+
+			it('sets item with specified setCookieOptions', async () => {
+				keyValueStorage.setItem('testKey', 'testValue');
+				expect(mockCookiesStorageAdapter.set).toHaveBeenCalledWith(
+					'testKey',
+					'testValue',
+					{
+						...defaultSetCookieOptions,
+						...testSetCookieOptions,
+					},
+				);
+			});
+		});
+
 		describe('in conjunction with token validator', () => {
 			const testKey = 'testKey';
 			const testValue = 'testValue';

diff --git a/...mplify/src/adapter-core/storageFactories/createKeyValueStorageFromCookieStorageAdapter.ts b/...mplify/src/adapter-core/storageFactories/createKeyValueStorageFromCookieStorageAdapter.ts
@@ -24,7 +24,12 @@ const ONE_YEAR_IN_MS = 365 * 24 * 60 * 60 * 1000;
 export const createKeyValueStorageFromCookieStorageAdapter = (
 	cookieStorageAdapter: CookieStorage.Adapter,
 	validatorMap?: KeyValueStorageMethodValidator,
+	setCookieOptions: CookieStorage.SetCookieOptions = {},
 ): KeyValueStorageInterface => {
+	const overrideCookieAttributes = {
+		...setCookieOptions,
+	};
+
 	return {
 		setItem(key, value) {
 			// Delete the cookie item first then set it. This results:
@@ -36,6 +41,7 @@ export const createKeyValueStorageFromCookieStorageAdapter = (
 			cookieStorageAdapter.set(key, value, {
 				...defaultSetCookieOptions,
 				expires: new Date(Date.now() + ONE_YEAR_IN_MS),
+				...overrideCookieAttributes,
 			});
 
 			return Promise.resolve();