fix(adapter-nextjs): create jwt verifier once (#13825)

* fix: create validator once * fix: move token validator initialization to createRunWithAmplifyServerContext * Fix logic issue * Fixed outdated comments --------- Co-authored-by: Chris Fang <cshfang@gmail.com> Co-authored-by: Chris F <5827964+cshfang@users.noreply.github.com>
aws-amplify · Dec 17, 2024 · 88f9eef · 88f9eef
1 parent ec2ff53
commit 88f9eef
Show file tree

Hide file tree

Showing 9 changed files with 199 additions and 162 deletions.
diff --git a/packages/adapter-nextjs/__tests__/createServerRunner.test.ts b/packages/adapter-nextjs/__tests__/createServerRunner.test.ts
@@ -20,7 +20,6 @@ const mockAmplifyConfig: ResourcesConfig = {
 		},
 	},
 };
-
 jest.mock(
 	'../src/utils/createCookieStorageAdapterFromNextServerContext',
 	() => ({
@@ -30,6 +29,7 @@ jest.mock(
 
 describe('createServerRunner', () => {
 	let createServerRunner: any;
+	let createRunWithAmplifyServerContextSpy: any;
 
 	const mockParseAmplifyConfig = jest.fn();
 	const mockCreateAWSCredentialsAndIdentityIdProvider = jest.fn();
@@ -50,11 +50,16 @@ describe('createServerRunner', () => {
 		jest.doMock('@aws-amplify/core/internals/utils', () => ({
 			parseAmplifyConfig: mockParseAmplifyConfig,
 		}));
+		createRunWithAmplifyServerContextSpy = jest.spyOn(
+			require('../src/utils/createRunWithAmplifyServerContext'),
+			'createRunWithAmplifyServerContext',
+		);
 
 		({ createServerRunner } = require('../src'));
 	});
 
 	afterEach(() => {
+		createRunWithAmplifyServerContextSpy.mockClear();
 		mockParseAmplifyConfig.mockClear();
 		mockCreateAWSCredentialsAndIdentityIdProvider.mockClear();
 		mockCreateKeyValueStorageFromCookieStorageAdapter.mockClear();
@@ -98,6 +103,10 @@ describe('createServerRunner', () => {
 					{},
 					operation,
 				);
+				expect(createRunWithAmplifyServerContextSpy).toHaveBeenCalledWith({
+					config: mockAmplifyConfigWithoutAuth,
+					tokenValidator: undefined,
+				});
 			});
 		});
 
@@ -120,6 +129,12 @@ describe('createServerRunner', () => {
 						mockAmplifyConfig.Auth,
 						sharedInMemoryStorage,
 					);
+					expect(createRunWithAmplifyServerContextSpy).toHaveBeenCalledWith({
+						config: mockAmplifyConfig,
+						tokenValidator: expect.objectContaining({
+							getItem: expect.any(Function),
+						}),
+					});
 				});
 			});
 
@@ -162,6 +177,12 @@ describe('createServerRunner', () => {
 						mockAmplifyConfig.Auth,
 						mockCookieStorageAdapter,
 					);
+					expect(createRunWithAmplifyServerContextSpy).toHaveBeenCalledWith({
+						config: mockAmplifyConfig,
+						tokenValidator: expect.objectContaining({
+							getItem: expect.any(Function),
+						}),
+					});
 				});
 			});
 		});

diff --git a/packages/adapter-nextjs/__tests__/utils/createTokenValidator.test.ts b/packages/adapter-nextjs/__tests__/utils/createTokenValidator.test.ts
@@ -1,85 +1,108 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+import { CognitoJwtVerifier } from 'aws-jwt-verify';
+
 import { isValidCognitoToken } from '../../src/utils/isValidCognitoToken';
 import { createTokenValidator } from '../../src/utils/createTokenValidator';
+import { JwtVerifier } from '../../src/types';
 
+jest.mock('aws-jwt-verify');
 jest.mock('../../src/utils/isValidCognitoToken');
 
-const mockIsValidCognitoToken = isValidCognitoToken as jest.Mock;
-
-const userPoolId = 'userPoolId';
-const userPoolClientId = 'clientId';
-const tokenValidatorInput = {
-	userPoolId,
-	userPoolClientId,
-};
-const accessToken = {
-	key: 'CognitoIdentityServiceProvider.clientId.usersub.accessToken',
-	value:
-		'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMTEiLCJpc3MiOiJodHRwc',
-};
-const idToken = {
-	key: 'CognitoIdentityServiceProvider.clientId.usersub.idToken',
-	value: 'eyJzdWIiOiIxMTEiLCJpc3MiOiJodHRwc.XAiOiJKV1QiLCJhbGciOiJIUzI1NiJ',
-};
-
-const tokenValidator = createTokenValidator({
-	userPoolId,
-	userPoolClientId,
-});
+describe('createTokenValidator', () => {
+	const userPoolId = 'userPoolId';
+	const userPoolClientId = 'clientId';
+	const accessToken = {
+		key: 'CognitoIdentityServiceProvider.clientId.usersub.accessToken',
+		value: 'access-token-value',
+	};
+	const idToken = {
+		key: 'CognitoIdentityServiceProvider.clientId.usersub.idToken',
+		value: 'id-token-value',
+	};
+
+	const mockIsValidCognitoToken = jest.mocked(isValidCognitoToken);
+	const mockCognitoJwtVerifier = {
+		create: jest.mocked(CognitoJwtVerifier.create),
+	};
 
-describe('Validator', () => {
 	afterEach(() => {
-		jest.resetAllMocks();
-	});
-	it('should return a validator', () => {
-		expect(createTokenValidator(tokenValidatorInput)).toBeDefined();
+		mockIsValidCognitoToken.mockClear();
 	});
 
-	it('should return true for non-token keys', async () => {
-		const result = await tokenValidator.getItem?.('mockKey', 'mockValue');
-		expect(result).toBe(true);
-		expect(mockIsValidCognitoToken).toHaveBeenCalledTimes(0);
+	it('should return a token validator', () => {
+		expect(
+			createTokenValidator({
+				userPoolId,
+				userPoolClientId,
+			}),
+		).toStrictEqual({
+			getItem: expect.any(Function),
+		});
 	});
 
-	it('should return true for valid accessToken', async () => {
-		mockIsValidCognitoToken.mockImplementation(() => Promise.resolve(true));
-
-		const result = await tokenValidator.getItem?.(
-			accessToken.key,
-			accessToken.value,
-		);
-
-		expect(result).toBe(true);
-		expect(mockIsValidCognitoToken).toHaveBeenCalledTimes(1);
-		expect(mockIsValidCognitoToken).toHaveBeenCalledWith({
-			userPoolId,
-			clientId: userPoolClientId,
-			token: accessToken.value,
-			tokenType: 'access',
+	describe('created token validator', () => {
+		afterEach(() => {
+			mockCognitoJwtVerifier.create.mockReset();
 		});
-	});
 
-	it('should return true for valid idToken', async () => {
-		mockIsValidCognitoToken.mockImplementation(() => Promise.resolve(true));
-
-		const result = await tokenValidator.getItem?.(idToken.key, idToken.value);
-		expect(result).toBe(true);
-		expect(mockIsValidCognitoToken).toHaveBeenCalledTimes(1);
-		expect(mockIsValidCognitoToken).toHaveBeenCalledWith({
-			userPoolId,
-			clientId: userPoolClientId,
-			token: idToken.value,
-			tokenType: 'id',
+		it('should return true if key is not for access or id tokens', async () => {
+			const tokenValidator = createTokenValidator({
+				userPoolId,
+				userPoolClientId,
+			});
+
+			expect(await tokenValidator.getItem?.('key', 'value')).toBe(true);
+			expect(mockIsValidCognitoToken).not.toHaveBeenCalled();
 		});
-	});
 
-	it('should return false if invalid tokenType is access', async () => {
-		mockIsValidCognitoToken.mockImplementation(() => Promise.resolve(false));
+		it('should return false if validator created without user pool or client ids', async () => {
+			const tokenValidator = createTokenValidator({});
 
-		const result = await tokenValidator.getItem?.(idToken.key, idToken.value);
-		expect(result).toBe(false);
-		expect(mockIsValidCognitoToken).toHaveBeenCalledTimes(1);
+			expect(
+				await tokenValidator.getItem?.(accessToken.key, accessToken.value),
+			).toBe(false);
+			expect(await tokenValidator.getItem?.(idToken.key, idToken.value)).toBe(
+				false,
+			);
+			expect(mockIsValidCognitoToken).not.toHaveBeenCalled();
+		});
+
+		describe.each([
+			{ tokenUse: 'access', token: accessToken },
+			{ tokenUse: 'id', token: idToken },
+		])('$tokenUse token verifier', ({ tokenUse, token }) => {
+			const mockTokenVerifier = {} as JwtVerifier;
+			const tokenValidator = createTokenValidator({
+				userPoolId,
+				userPoolClientId,
+			});
+
+			beforeAll(() => {
+				mockCognitoJwtVerifier.create.mockReturnValue(mockTokenVerifier);
+			});
+
+			it('should create a jwt verifier and use it to validate', async () => {
+				await tokenValidator.getItem?.(token.key, token.value);
+
+				expect(mockCognitoJwtVerifier.create).toHaveBeenCalledWith({
+					userPoolId,
+					clientId: userPoolClientId,
+					tokenUse,
+				});
+				expect(mockIsValidCognitoToken).toHaveBeenCalledWith({
+					token: token.value,
+					verifier: mockTokenVerifier,
+				});
+			});
+
+			it('should not re-create the jwt verifier', async () => {
+				await tokenValidator.getItem?.(token.key, token.value);
+
+				expect(mockCognitoJwtVerifier.create).not.toHaveBeenCalled();
+				expect(mockIsValidCognitoToken).toHaveBeenCalled();
+			});
+		});
 	});
 });
diff --git a/packages/adapter-nextjs/__tests__/utils/isValidCognitoToken.test.ts b/packages/adapter-nextjs/__tests__/utils/isValidCognitoToken.test.ts
@@ -1,94 +1,57 @@
-import { CognitoJwtVerifier } from 'aws-jwt-verify';
 import { JwtExpiredError } from 'aws-jwt-verify/error';
 
 import { isValidCognitoToken } from '../../src/utils/isValidCognitoToken';
-
-jest.mock('aws-jwt-verify', () => {
-	return {
-		CognitoJwtVerifier: {
-			create: jest.fn(),
-		},
-	};
-});
-
-const mockedCreate = CognitoJwtVerifier.create as jest.MockedFunction<
-	typeof CognitoJwtVerifier.create
->;
+import { JwtVerifier } from '../../src/types';
 
 describe('isValidCognitoToken', () => {
 	const token = 'mocked-token';
-	const userPoolId = 'us-east-1_test';
-	const clientId = 'client-id-test';
-	const tokenType = 'id';
 
 	beforeEach(() => {
 		jest.clearAllMocks();
 	});
 
 	it('should return true for a valid token', async () => {
-		const mockVerifier: any = {
-			verify: jest.fn().mockResolvedValue({}),
+		// @ts-expect-error - partial mock
+		const mockVerifier: JwtVerifier = {
+			verify: jest.fn().mockResolvedValue(null),
 		};
-		mockedCreate.mockReturnValue(mockVerifier);
 
-		const isValid = await isValidCognitoToken({
-			token,
-			userPoolId,
-			clientId,
-			tokenType,
-		});
-		expect(isValid).toBe(true);
-		expect(CognitoJwtVerifier.create).toHaveBeenCalledWith({
-			userPoolId,
-			clientId,
-			tokenUse: tokenType,
-		});
+		expect(
+			await isValidCognitoToken({
+				token,
+				verifier: mockVerifier,
+			}),
+		).toBe(true);
 		expect(mockVerifier.verify).toHaveBeenCalledWith(token);
 	});
 
-	it('should return true for a token that has valid signature and expired', async () => {
-		const mockVerifier: any = {
+	it('should return true for a token that has valid signature but is expired', async () => {
+		// @ts-expect-error - partial mock
+		const mockVerifier: JwtVerifier = {
 			verify: jest
 				.fn()
-				.mockRejectedValue(
-					new JwtExpiredError('Token expired', 'mocked-token'),
-				),
+				.mockRejectedValue(new JwtExpiredError('Token expired', token)),
 		};
-		mockedCreate.mockReturnValue(mockVerifier);
 
-		const isValid = await isValidCognitoToken({
-			token,
-			userPoolId,
-			clientId,
-			tokenType,
-		});
-		expect(isValid).toBe(true);
-		expect(CognitoJwtVerifier.create).toHaveBeenCalledWith({
-			userPoolId,
-			clientId,
-			tokenUse: tokenType,
-		});
-		expect(mockVerifier.verify).toHaveBeenCalledWith(token);
+		expect(
+			await isValidCognitoToken({
+				token,
+				verifier: mockVerifier,
+			}),
+		).toBe(true);
 	});
 
 	it('should return false for an invalid token', async () => {
-		const mockVerifier: any = {
-			verify: jest.fn().mockRejectedValue(new Error('Invalid token')),
+		// @ts-expect-error - partial mock
+		const mockVerifier: JwtVerifier = {
+			verify: jest.fn().mockRejectedValue(null),
 		};
-		mockedCreate.mockReturnValue(mockVerifier);
 
-		const isValid = await isValidCognitoToken({
-			token,
-			userPoolId,
-			clientId,
-			tokenType,
-		});
-		expect(isValid).toBe(false);
-		expect(CognitoJwtVerifier.create).toHaveBeenCalledWith({
-			userPoolId,
-			clientId,
-			tokenUse: tokenType,
-		});
-		expect(mockVerifier.verify).toHaveBeenCalledWith(token);
+		expect(
+			await isValidCognitoToken({
+				token,
+				verifier: mockVerifier,
+			}),
+		).toBe(false);
 	});
 });
diff --git a/packages/adapter-nextjs/src/createServerRunner.ts b/packages/adapter-nextjs/src/createServerRunner.ts
@@ -2,10 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { ResourcesConfig } from 'aws-amplify';
+import { KeyValueStorageMethodValidator } from '@aws-amplify/core/internals/adapter-core';
 import { parseAmplifyConfig } from '@aws-amplify/core/internals/utils';
 
 import { createRunWithAmplifyServerContext } from './utils';
 import { NextServer } from './types';
+import { createTokenValidator } from './utils/createTokenValidator';
 
 /**
  * Creates the `runWithAmplifyServerContext` function to run Amplify server side APIs in an isolated request context.
@@ -30,9 +32,19 @@ export const createServerRunner: NextServer.CreateServerRunner = ({
 }) => {
 	const amplifyConfig = parseAmplifyConfig(config);
 
+	let tokenValidator: KeyValueStorageMethodValidator | undefined;
+	if (amplifyConfig?.Auth) {
+		const { Cognito } = amplifyConfig.Auth;
+		tokenValidator = createTokenValidator({
+			userPoolId: Cognito?.userPoolId,
+			userPoolClientId: Cognito?.userPoolClientId,
+		});
+	}
+
 	return {
 		runWithAmplifyServerContext: createRunWithAmplifyServerContext({
 			config: amplifyConfig,
+			tokenValidator,
 		}),
 	};
 };