fix(auth): verifyTotp throw EnableSoftwareTokenMfaException (#4558)

aws-amplify · Mar 15, 2024 · 396d0ec · 396d0ec
1 parent e964950
commit 396d0ec
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 17 deletions.
diff --git a/packages/auth/amplify_auth_cognito/example/integration_test/mfa_totp_optional_test.dart b/packages/auth/amplify_auth_cognito/example/integration_test/mfa_totp_optional_test.dart
@@ -1,6 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+import 'package:amplify_auth_cognito_dart/amplify_auth_cognito_dart.dart';
 import 'package:amplify_auth_integration_test/amplify_auth_integration_test.dart';
 import 'package:amplify_flutter/amplify_flutter.dart';
 import 'package:amplify_integration_test/amplify_integration_test.dart';
@@ -117,9 +118,15 @@ void main() {
 
           final totpSetupResult = await Amplify.Auth.setUpTotp();
 
-          await Amplify.Auth.verifyTotpSetup(
-            '555555',
-          );
+          try {
+            await Amplify.Auth.verifyTotpSetup('555555');
+            fail('Expected to fail');
+          } on AuthException catch (e) {
+            check(
+              e,
+              because: 'Invalid TOTP code should fail verification',
+            ).isA<EnableSoftwareTokenMfaException>();
+          }
 
           check(
             await cognitoPlugin.fetchMfaPreference(),
@@ -131,9 +138,13 @@ void main() {
             ),
           );
 
-          await Amplify.Auth.verifyTotpSetup(
-            await generateTotpCode(totpSetupResult.sharedSecret),
-          );
+          try {
+            await Amplify.Auth.verifyTotpSetup(
+              await generateTotpCode(totpSetupResult.sharedSecret),
+            );
+          } on Exception catch (e) {
+            fail('Expected to succeed, but got $e');
+          }
 
           check(await cognitoPlugin.fetchMfaPreference()).equals(
             const UserMfaPreference(

diff --git a/packages/auth/amplify_auth_cognito_dart/lib/src/auth_plugin_impl.dart b/packages/auth/amplify_auth_cognito_dart/lib/src/auth_plugin_impl.dart
@@ -932,12 +932,24 @@ class AmplifyAuthCognitoDart extends AuthPluginInterface
       defaultPluginOptions: const CognitoVerifyTotpSetupPluginOptions(),
     );
     final machine = _stateMachine.getOrCreate(TotpSetupStateMachine.type);
-    await machine.dispatchAndComplete<TotpSetupState>(
+    final state = await machine.dispatchAndComplete<TotpSetupState>(
       TotpSetupEvent.verify(
         code: totpCode,
         friendlyDeviceName: pluginOptions.friendlyDeviceName,
       ),
     );
+
+    switch (state) {
+      case TotpSetupRequiresVerification _:
+        // TODO(equartey): Change to `CodeMismatchException` in next major version as breaking change
+        throw const EnableSoftwareTokenMfaException(
+          'The code provided was incorrect, try again',
+        );
+      case TotpSetupFailure(:final exception, :final stackTrace):
+        Error.throwWithStackTrace(exception, stackTrace);
+      default:
+        return;
+    }
   }
 
   @override

diff --git a/packages/auth/amplify_auth_cognito_dart/lib/src/state/machines/totp_setup_state_machine.dart b/packages/auth/amplify_auth_cognito_dart/lib/src/state/machines/totp_setup_state_machine.dart
@@ -3,7 +3,7 @@
 
 import 'package:amplify_auth_cognito_dart/src/jwt/src/cognito.dart';
 import 'package:amplify_auth_cognito_dart/src/sdk/cognito_identity_provider.dart'
-    hide EnableSoftwareTokenMfaException;
+    hide EnableSoftwareTokenMfaException, CodeMismatchException;
 import 'package:amplify_auth_cognito_dart/src/sdk/sdk_bridge.dart';
 import 'package:amplify_auth_cognito_dart/src/sdk/sdk_exception.dart';
 import 'package:amplify_auth_cognito_dart/src/state/cognito_state_machine.dart';
@@ -76,6 +76,7 @@ final class TotpSetupStateMachine
   Future<void> _onVerify(TotpSetupVerify event) async {
     final tokens = await manager.getUserPoolTokens();
     final accessToken = tokens.accessToken.raw;
+
     try {
       await _cognitoIdp
           .verifySoftwareToken(
@@ -87,16 +88,16 @@ final class TotpSetupStateMachine
             ),
           )
           .result;
-    } on Exception catch (e, st) {
+    } on Exception catch (e) {
       // Handle mismatch code exception that may occur during TOTP verification.
       // See: https://docs.aws.amazon.com/cognito-user-identity-pools/latest/APIReference/API_VerifySoftwareToken.html#API_VerifySoftwareToken_Errors
-      if (e is EnableSoftwareTokenMfaException) {
+      if (e is EnableSoftwareTokenMfaException || e is CodeMismatchException) {
         assert(
           _details != null,
           'TotpSetupDetails should not be null. Please report this issue.',
         );
         logger.verbose(
-          'Failed to verify TOTP code. Retrying...',
+          'Failed to verify TOTP code. Allowing retry...',
           e,
         );
         emit(
@@ -106,12 +107,7 @@ final class TotpSetupStateMachine
         );
         return;
       }
-      logger.error(
-        'Failed to verify TOTP code. Please try again.',
-        e,
-        st,
-      );
-      emit(TotpSetupState.failure(e, st));
+      rethrow;
     }
 
     try {