feat(core): validate if access and id tokens are valid cognito tokens (…

…#13385) * feat(adapter-nextjs): validate if access and id tokens are valid cognito tokens * add unit tests and cleanups * increase bundle size * increase bundle size * Update packages/core/src/adapterCore/serverContext/types/KeyValueStorageValidator.ts Co-authored-by: Hui Zhao <[email protected]> * Update packages/adapter-nextjs/__tests__/utils/createTokenValidator.test.ts Co-authored-by: Hui Zhao <[email protected]> * Update packages/core/src/adapterCore/serverContext/types/KeyValueStorageValidator.ts Co-authored-by: Hui Zhao <[email protected]> * address feedback * address feedback * Update packages/adapter-nextjs/src/utils/createTokenValidator.ts Co-authored-by: israx <[email protected]> * address feedback --------- Co-authored-by: Ashwin Kumar <[email protected]> Co-authored-by: Hui Zhao <[email protected]> Co-authored-by: israx <[email protected]>
aws-amplify · May 21, 2024 · 0b72b32 · 0b72b32
1 parent d7d33da
commit 0b72b32
Show file tree

Hide file tree

Showing 15 changed files with 339 additions and 35 deletions.
diff --git a/packages/adapter-nextjs/__tests__/utils/createTokenValidator.test.ts b/packages/adapter-nextjs/__tests__/utils/createTokenValidator.test.ts
@@ -0,0 +1,88 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { isValidCognitoToken } from '@aws-amplify/core/internals/utils';
+
+import { createTokenValidator } from '../../src/utils/createTokenValidator';
+
+jest.mock('@aws-amplify/core/internals/utils', () => ({
+	...jest.requireActual('@aws-amplify/core/internals/utils'),
+	isValidCognitoToken: jest.fn(),
+}));
+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('Validator', () => {
+	afterEach(() => {
+		jest.resetAllMocks();
+	});
+	it('should return a validator', () => {
+		expect(createTokenValidator(tokenValidatorInput)).toBeDefined();
+	});
+
+	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 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',
+		});
+	});
+
+	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 false if invalid tokenType is access', async () => {
+		mockIsValidCognitoToken.mockImplementation(() => Promise.resolve(false));
+
+		const result = await tokenValidator.getItem?.(idToken.key, idToken.value);
+		expect(result).toBe(false);
+		expect(mockIsValidCognitoToken).toHaveBeenCalledTimes(1);
+	});
+});
diff --git a/packages/adapter-nextjs/src/utils/createRunWithAmplifyServerContext.ts b/packages/adapter-nextjs/src/utils/createRunWithAmplifyServerContext.ts
@@ -11,6 +11,7 @@ import {
 
 import { NextServer } from '../types';
 
+import { createTokenValidator } from './createTokenValidator';
 import { createCookieStorageAdapterFromNextServerContext } from './createCookieStorageAdapterFromNextServerContext';
 
 export const createRunWithAmplifyServerContext = ({
@@ -34,6 +35,11 @@ export const createRunWithAmplifyServerContext = ({
 								createCookieStorageAdapterFromNextServerContext(
 									nextServerContext,
 								),
+								createTokenValidator({
+									userPoolId: resourcesConfig?.Auth.Cognito?.userPoolId,
+									userPoolClientId:
+										resourcesConfig?.Auth.Cognito?.userPoolClientId,
+								}),
 							);
 				const credentialsProvider = createAWSCredentialsAndIdentityIdProvider(
 					resourcesConfig.Auth,

diff --git a/packages/adapter-nextjs/src/utils/createTokenValidator.ts b/packages/adapter-nextjs/src/utils/createTokenValidator.ts
@@ -0,0 +1,38 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { isValidCognitoToken } from '@aws-amplify/core/internals/utils';
+import { KeyValueStorageMethodValidator } from '@aws-amplify/core/internals/adapter-core';
+
+interface CreateTokenValidatorInput {
+	userPoolId?: string;
+	userPoolClientId?: string;
+}
+/**
+ * Creates a validator object for validating methods in a KeyValueStorage.
+ */
+export const createTokenValidator = ({
+	userPoolId,
+	userPoolClientId: clientId,
+}: CreateTokenValidatorInput): KeyValueStorageMethodValidator => {
+	return {
+		// validate access, id tokens
+		getItem: async (key: string, value: string): Promise<boolean> => {
+			const tokenType = key.includes('.accessToken')
+				? 'access'
+				: key.includes('.idToken')
+					? 'id'
+					: null;
+			if (!tokenType) return true;
+
+			if (!userPoolId || !clientId) return false;
+
+			return isValidCognitoToken({
+				clientId,
+				userPoolId,
+				tokenType,
+				token: value,
+			});
+		},
+	};
+};
diff --git a/...ests__/adapterCore/storageFactories/createKeyValueStorageFromCookieStorageAdapter.test.ts b/...ests__/adapterCore/storageFactories/createKeyValueStorageFromCookieStorageAdapter.test.ts
@@ -84,5 +84,44 @@ describe('keyValueStorage', () => {
 				}).toThrow('This method has not implemented.');
 			});
 		});
+
+		describe('in conjunction with token validator', () => {
+			const testKey = 'testKey';
+			const testValue = 'testValue';
+
+			beforeEach(() => {
+				mockCookiesStorageAdapter.get.mockReturnValueOnce({
+					name: testKey,
+					value: testValue,
+				});
+			});
+			afterEach(() => {
+				jest.clearAllMocks();
+			});
+
+			it('should return item successfully if validation passes when getting item', async () => {
+				const getItemValidator = jest.fn().mockImplementation(() => true);
+				const keyValueStorage = createKeyValueStorageFromCookieStorageAdapter(
+					mockCookiesStorageAdapter,
+					{ getItem: getItemValidator },
+				);
+
+				const value = await keyValueStorage.getItem(testKey);
+				expect(value).toBe(testValue);
+				expect(getItemValidator).toHaveBeenCalledTimes(1);
+			});
+
+			it('should return null if validation fails when getting item', async () => {
+				const getItemValidator = jest.fn().mockImplementation(() => false);
+				const keyValueStorage = createKeyValueStorageFromCookieStorageAdapter(
+					mockCookiesStorageAdapter,
+					{ getItem: getItemValidator },
+				);
+
+				const value = await keyValueStorage.getItem(testKey);
+				expect(value).toBe(null);
+				expect(getItemValidator).toHaveBeenCalledTimes(1);
+			});
+		});
 	});
 });