chore(adapter-nextjs): validate if access and id tokens are valid cog…

…nito tokens

packages/adapter-nextjs/__tests__/utils/createTokenValidator.test.ts

@@ -0,0 +1,85 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { isValidCognitoToken } from '../../src/utils/isValidCognitoToken';
+import { createTokenValidator } from '../../src/utils/createTokenValidator';
+
+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('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);
+	});
+});

packages/adapter-nextjs/__tests__/utils/isValidCognitoToken.test.ts

@@ -0,0 +1,68 @@
+import { CognitoJwtVerifier } from 'aws-jwt-verify';
+
+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
+>;
+
+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({}),
+		};
+		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);
+	});
+
+	it('should return false for an invalid token', async () => {
+		const mockVerifier: any = {
+			verify: jest.fn().mockRejectedValue(new Error('Invalid token')),
+		};
+		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);
+	});
+});

packages/adapter-nextjs/package.json

@@ -8,6 +8,7 @@
 		"next": ">=13.5.0 <15.0.0"
 	},
 	"dependencies": {
+		"aws-jwt-verify": "^4.0.1",
 		"cookie": "0.5.0"
 	},
 	"devDependencies": {

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,

packages/adapter-nextjs/src/utils/createTokenValidator.ts

@@ -0,0 +1,39 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { KeyValueStorageMethodValidator } from '@aws-amplify/core/internals/adapter-core';
+
+import { isValidCognitoToken } from './isValidCognitoToken';
+
+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,
+			});
+		},
+	};
+};

packages/adapter-nextjs/src/utils/index.ts

@@ -2,3 +2,4 @@
 // SPDX-License-Identifier: Apache-2.0
 
 export { createRunWithAmplifyServerContext } from './createRunWithAmplifyServerContext';
+export { isValidCognitoToken } from './isValidCognitoToken';

packages/adapter-nextjs/src/utils/isValidCognitoToken.ts

@@ -0,0 +1,37 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { CognitoJwtVerifier } from 'aws-jwt-verify';
+
+/**
+ * Verifies a Cognito JWT token for its validity.
+ *
+ * @param input - An object containing:
+ *                - token: The JWT token as a string that needs to be verified.
+ *                - userPoolId: The ID of the AWS Cognito User Pool to which the token belongs.
+ *                - clientId: The Client ID associated with the Cognito User Pool.
+ * @internal
+ */
+export const isValidCognitoToken = async (input: {
+	token: string;
+	userPoolId: string;
+	clientId: string;
+	tokenType: 'id' | 'access';
+}): Promise<boolean> => {
+	const { userPoolId, clientId, tokenType, token } = input;
+
+	try {
+		const verifier = CognitoJwtVerifier.create({
+			userPoolId,
+			tokenUse: tokenType,
+			clientId,
+		});
+		await verifier.verify(token);
+
+		return true;
+	} catch (error) {
+		// TODO (ashwinkumar6): surface invalid cognito token error to customer
+		// TODO: clear invalid tokens from Storage
+		return false;
+	}
+};

packages/aws-amplify/__tests__/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);
+			});
+		});
 	});
 });

packages/aws-amplify/src/adapter-core/storageFactories/createKeyValueStorageFromCookieStorageAdapter.ts

@@ -2,7 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { KeyValueStorageInterface } from '@aws-amplify/core';
-import { CookieStorage } from '@aws-amplify/core/internals/adapter-core';
+import {
+	CookieStorage,
+	KeyValueStorageMethodValidator,
+} from '@aws-amplify/core/internals/adapter-core';
 
 export const defaultSetCookieOptions: CookieStorage.SetCookieOptions = {
 	// TODO: allow configure with a public interface
@@ -18,6 +21,7 @@ const ONE_YEAR_IN_MS = 365 * 24 * 60 * 60 * 1000;
  */
 export const createKeyValueStorageFromCookieStorageAdapter = (
 	cookieStorageAdapter: CookieStorage.Adapter,
+	validatorMap?: KeyValueStorageMethodValidator,
 ): KeyValueStorageInterface => {
 	return {
 		setItem(key, value) {
@@ -29,10 +33,16 @@ export const createKeyValueStorageFromCookieStorageAdapter = (
 
 			return Promise.resolve();
 		},
-		getItem(key) {
+		async getItem(key) {
 			const cookie = cookieStorageAdapter.get(key);
+			const value = cookie?.value ?? null;
 
-			return Promise.resolve(cookie?.value ?? null);
+			if (value && validatorMap?.getItem) {
+				const isValid = await validatorMap.getItem(key, value);
+				if (!isValid) return null;
+			}
+
+			return value;
 		},
 		removeItem(key) {
 			cookieStorageAdapter.delete(key);

packages/core/src/adapterCore/index.ts

@@ -7,5 +7,6 @@ export {
 	destroyAmplifyServerContext,
 	AmplifyServer,
 	CookieStorage,
+	KeyValueStorageMethodValidator,
 } from './serverContext';
 export { AmplifyServerContextError } from './error';

packages/core/src/adapterCore/serverContext/index.ts

@@ -7,4 +7,8 @@ export {
 	getAmplifyServerContext,
 } from './serverContext';
 
-export { AmplifyServer, CookieStorage } from './types';
+export {
+	AmplifyServer,
+	CookieStorage,
+	KeyValueStorageMethodValidator,
+} from './types';

packages/core/src/adapterCore/serverContext/types/KeyValueStorageMethodValidator.ts

@@ -0,0 +1,10 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { KeyValueStorageInterface } from '../../../types/storage';
+
+export type KeyValueStorageMethodValidator = Partial<
+	Record<keyof KeyValueStorageInterface, ValidatorFunction>
+>;
+
+type ValidatorFunction = (...args: any[]) => Promise<boolean>;

packages/core/src/adapterCore/serverContext/types/index.ts

@@ -7,3 +7,4 @@ type AmplifyServerContextSpec = AmplifyServer.ContextSpec;
 
 export { AmplifyServerContextSpec, AmplifyServer };
 export { CookieStorage } from './cookieStorage';
+export { KeyValueStorageMethodValidator } from './KeyValueStorageMethodValidator';