fix(auth): Allow retries with verifyTotpSetup() (#4532)

Author: Equartey

packages/auth/amplify_auth_cognito/example/integration_test/mfa_totp_optional_test.dart

@@ -93,6 +93,55 @@ void main() {
           'w/ device name',
           (_) => runTest(friendlyDeviceName: friendlyDeviceName),
         );
+
+        asyncTest('verifyTotpSetup allows retries', (_) async {
+          final username = generateUsername();
+          final password = generatePassword();
+
+          await adminCreateUser(
+            username,
+            password,
+            verifyAttributes: true,
+            autoConfirm: true,
+          );
+
+          final signInRes = await Amplify.Auth.signIn(
+            username: username,
+            password: password,
+          );
+
+          check(
+            because: 'TOTP is optional',
+            signInRes.nextStep.signInStep,
+          ).equals(AuthSignInStep.done);
+
+          final totpSetupResult = await Amplify.Auth.setUpTotp();
+
+          await Amplify.Auth.verifyTotpSetup(
+            '555555',
+          );
+
+          check(
+            await cognitoPlugin.fetchMfaPreference(),
+            because: 'TOTP should not be enabled',
+          ).equals(
+            const UserMfaPreference(
+              enabled: {},
+              preferred: null,
+            ),
+          );
+
+          await Amplify.Auth.verifyTotpSetup(
+            await generateTotpCode(totpSetupResult.sharedSecret),
+          );
+
+          check(await cognitoPlugin.fetchMfaPreference()).equals(
+            const UserMfaPreference(
+              enabled: {MfaType.totp},
+              preferred: MfaType.totp,
+            ),
+          );
+        });
       });
     });
   });

packages/auth/amplify_auth_cognito_dart/lib/src/auth_plugin_impl.dart

@@ -932,7 +932,7 @@ class AmplifyAuthCognitoDart extends AuthPluginInterface
       defaultPluginOptions: const CognitoVerifyTotpSetupPluginOptions(),
     );
     final machine = _stateMachine.getOrCreate(TotpSetupStateMachine.type);
-    await machine.dispatchAndComplete<TotpSetupSuccess>(
+    await machine.dispatchAndComplete<TotpSetupState>(
       TotpSetupEvent.verify(
         code: totpCode,
         friendlyDeviceName: pluginOptions.friendlyDeviceName,

packages/auth/amplify_auth_cognito_dart/lib/src/state/machines/totp_setup_state_machine.dart

@@ -2,8 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import 'package:amplify_auth_cognito_dart/src/jwt/src/cognito.dart';
-import 'package:amplify_auth_cognito_dart/src/sdk/cognito_identity_provider.dart';
+import 'package:amplify_auth_cognito_dart/src/sdk/cognito_identity_provider.dart'
+    hide EnableSoftwareTokenMfaException;
 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';
 import 'package:amplify_auth_cognito_dart/src/state/state.dart';
 import 'package:amplify_core/amplify_core.dart';
@@ -48,6 +50,7 @@ final class TotpSetupStateMachine
   CognitoIdentityProviderClient get _cognitoIdp => expect();
 
   String? _session;
+  TotpSetupDetails? _details;
 
   Future<void> _onInitiate(TotpSetupInitiate event) async {
     final tokens = await manager.getUserPoolTokens();
@@ -59,29 +62,58 @@ final class TotpSetupStateMachine
         )
         .result;
     _session = response.session;
+    _details = TotpSetupDetails(
+      username: CognitoIdToken(tokens.idToken).username,
+      sharedSecret: response.secretCode!,
+    );
     emit(
       TotpSetupState.requiresVerification(
-        TotpSetupDetails(
-          username: CognitoIdToken(tokens.idToken).username,
-          sharedSecret: response.secretCode!,
-        ),
+        _details!,
       ),
     );
   }
 
   Future<void> _onVerify(TotpSetupVerify event) async {
     final tokens = await manager.getUserPoolTokens();
     final accessToken = tokens.accessToken.raw;
-    await _cognitoIdp
-        .verifySoftwareToken(
-          VerifySoftwareTokenRequest(
-            accessToken: accessToken,
-            session: _session,
-            userCode: event.code,
-            friendlyDeviceName: event.friendlyDeviceName,
+    try {
+      await _cognitoIdp
+          .verifySoftwareToken(
+            VerifySoftwareTokenRequest(
+              accessToken: accessToken,
+              session: _session,
+              userCode: event.code,
+              friendlyDeviceName: event.friendlyDeviceName,
+            ),
+          )
+          .result;
+    } on Exception catch (e, st) {
+      // 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) {
+        assert(
+          _details != null,
+          'TotpSetupDetails should not be null. Please report this issue.',
+        );
+        logger.verbose(
+          'Failed to verify TOTP code. Retrying...',
+          e,
+        );
+        emit(
+          TotpSetupState.requiresVerification(
+            _details!,
           ),
-        )
-        .result;
+        );
+        return;
+      }
+      logger.error(
+        'Failed to verify TOTP code. Please try again.',
+        e,
+        st,
+      );
+      emit(TotpSetupState.failure(e, st));
+    }
+
     try {
       await _cognitoIdp.setMfaSettings(
         accessToken: accessToken,