fix: device tracking for MFA (#14626)

Co-authored-by: Philipp Andreas Paul <[email protected]>

Author: soberm

packages/auth/__tests__/providers/cognito/confirmSignInHappyCases.test.ts

@@ -9,6 +9,7 @@ import {
 	signIn,
 } from '../../../src/providers/cognito/';
 import * as signInHelpers from '../../../src/providers/cognito/utils/signInHelpers';
+import { handleDeviceSRPAuth } from '../../../src/providers/cognito/utils/handleDeviceSRPAuth';
 import {
 	cognitoUserPoolsTokenProvider,
 	tokenOrchestrator,
@@ -27,6 +28,7 @@ jest.mock('../../../src/providers/cognito/apis/getCurrentUser');
 jest.mock(
 	'../../../src/foundation/factories/serviceClients/cognitoIdentityProvider',
 );
+jest.mock('../../../src/providers/cognito/utils/handleDeviceSRPAuth');
 
 const authConfig = {
 	Cognito: {
@@ -447,6 +449,69 @@ describe('confirmSignIn API happy path cases', () => {
 			options,
 		);
 	});
+
+	test('device metadata is retrieved and handled during MFA challenge', async () => {
+		const mockGetDeviceMetadata = jest
+			.spyOn(tokenOrchestrator, 'getDeviceMetadata')
+			.mockResolvedValue({
+				deviceKey: 'device-123',
+				deviceGroupKey: 'group-456',
+				randomPassword: 'random-password',
+			});
+
+		const mockHandleDeviceSRPAuth = jest.mocked(handleDeviceSRPAuth);
+		mockHandleDeviceSRPAuth.mockResolvedValue({
+			ChallengeName: undefined,
+			$metadata: {},
+			AuthenticationResult: {
+				AccessToken: 'access-token',
+				IdToken: 'id-token',
+				RefreshToken: 'refresh-token',
+			},
+		});
+
+		const mockRespondToAuthChallenge = jest.fn().mockResolvedValue({
+			ChallengeName: 'DEVICE_SRP_AUTH',
+			Session: 'device-session',
+			$metadata: {},
+		});
+
+		const mockCreateRespondToAuthChallengeClient = jest.mocked(
+			createRespondToAuthChallengeClient,
+		);
+		mockCreateRespondToAuthChallengeClient.mockReturnValueOnce(
+			mockRespondToAuthChallenge,
+		);
+
+		await signInHelpers.handleMFAChallenge({
+			challengeName: 'SOFTWARE_TOKEN_MFA',
+			challengeResponse: '123456',
+			clientMetadata: undefined,
+			session: 'test-session',
+			username,
+			config: authConfig.Cognito,
+			tokenOrchestrator,
+		});
+
+		expect(mockGetDeviceMetadata).toHaveBeenCalledWith(username);
+		expect(mockRespondToAuthChallenge).toHaveBeenCalledWith(
+			expect.anything(),
+			expect.objectContaining({
+				ChallengeResponses: expect.objectContaining({
+					DEVICE_KEY: 'device-123',
+					SOFTWARE_TOKEN_MFA_CODE: '123456',
+					USERNAME: username,
+				}),
+			}),
+		);
+		expect(mockHandleDeviceSRPAuth).toHaveBeenCalledWith({
+			username,
+			config: authConfig.Cognito,
+			clientMetadata: undefined,
+			session: 'device-session',
+			tokenOrchestrator,
+		});
+	});
 });
 
 describe('Cognito ASF', () => {

packages/auth/src/providers/cognito/utils/signInHelpers.ts

@@ -71,6 +71,7 @@ interface HandleAuthChallengeRequest {
 	deviceName?: string;
 	requiredAttributes?: AuthUserAttributes;
 	config: CognitoUserPoolConfig;
+	tokenOrchestrator: AuthTokenOrchestrator;
 }
 
 function isWebAuthnResultAuthSignInOutput(
@@ -834,6 +835,7 @@ export async function handleChallengeName(
 				session,
 				username,
 				config,
+				tokenOrchestrator,
 			});
 		case 'MFA_SETUP':
 			return handleMFASetupChallenge({
@@ -843,6 +845,7 @@ export async function handleChallengeName(
 				username,
 				deviceName,
 				config,
+				tokenOrchestrator,
 			});
 		case 'NEW_PASSWORD_REQUIRED':
 			return handleCompleteNewPasswordChallenge({
@@ -852,6 +855,7 @@ export async function handleChallengeName(
 				username,
 				requiredAttributes: userAttributes,
 				config,
+				tokenOrchestrator,
 			});
 		case 'CUSTOM_CHALLENGE':
 			return retryOnResourceNotFoundException(
@@ -880,6 +884,7 @@ export async function handleChallengeName(
 				session,
 				username,
 				config,
+				tokenOrchestrator,
 			});
 		case 'PASSWORD':
 			return handleSelectChallengeWithPassword(
@@ -967,6 +972,7 @@ export async function handleMFAChallenge({
 	session,
 	username,
 	config,
+	tokenOrchestrator,
 }: HandleAuthChallengeRequest & {
 	challengeName: Extract<
 		ChallengeName,
@@ -995,6 +1001,11 @@ export async function handleMFAChallenge({
 		challengeResponses.SOFTWARE_TOKEN_MFA_CODE = challengeResponse;
 	}
 
+	const deviceMetadata = await tokenOrchestrator?.getDeviceMetadata(username);
+	if (deviceMetadata && deviceMetadata.deviceKey) {
+		challengeResponses.DEVICE_KEY = deviceMetadata.deviceKey;
+	}
+
 	const userContextData = getUserContextData({
 		username,
 		userPoolId,
@@ -1016,11 +1027,22 @@ export async function handleMFAChallenge({
 		}),
 	});
 
-	return respondToAuthChallenge(
+	const response = await respondToAuthChallenge(
 		{
 			region: getRegionFromUserPoolId(userPoolId),
 			userAgentValue: getAuthUserAgentValue(AuthAction.ConfirmSignIn),
 		},
 		jsonReq,
 	);
+	if (response.ChallengeName === 'DEVICE_SRP_AUTH') {
+		return handleDeviceSRPAuth({
+			username,
+			config,
+			clientMetadata,
+			session: response.Session,
+			tokenOrchestrator,
+		});
+	}
+
+	return response;
 }