fix(auth): Add MFAType.challengeResponse extension (#2680)

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/MFATypeUtil.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+@file:JvmName("MFATypeUtil")
+
+package com.amplifyframework.auth.cognito
+
+import com.amplifyframework.auth.MFAType
+
+/**
+ * Returns the cognito-specific string to pass to Amplify.Auth.confirmSignIn for a specific [MFAType] when making
+ * an MFA selection during the sign-in process.
+ */
+val MFAType.challengeResponse: String
+    get() = when (this) {
+        MFAType.SMS -> "SMS_MFA"
+        MFAType.TOTP -> "SOFTWARE_TOKEN_MFA"
+    }

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPlugin.kt

@@ -71,7 +71,9 @@ import com.amplifyframework.auth.cognito.helpers.HostedUIHelper
 import com.amplifyframework.auth.cognito.helpers.SessionHelper
 import com.amplifyframework.auth.cognito.helpers.SignInChallengeHelper
 import com.amplifyframework.auth.cognito.helpers.getMFAType
+import com.amplifyframework.auth.cognito.helpers.getMFATypeOrNull
 import com.amplifyframework.auth.cognito.helpers.identityProviderName
+import com.amplifyframework.auth.cognito.helpers.value
 import com.amplifyframework.auth.cognito.options.AWSCognitoAuthConfirmResetPasswordOptions
 import com.amplifyframework.auth.cognito.options.AWSCognitoAuthConfirmSignInOptions
 import com.amplifyframework.auth.cognito.options.AWSCognitoAuthConfirmSignUpOptions
@@ -789,14 +791,25 @@ internal class RealAWSCognitoAuthPlugin(
                 val userAttributes = awsCognitoConfirmSignInOptions?.userAttributes ?: emptyList()
                 when (signInState) {
                     is SignInState.ResolvingChallenge -> {
-                        val event = SignInChallengeEvent(
-                            SignInChallengeEvent.EventType.VerifyChallengeAnswer(
-                                challengeResponse,
-                                metadata,
-                                userAttributes
+                        val challengeState = signInState.challengeState
+                        if (challengeState is SignInChallengeState.WaitingForAnswer &&
+                            challengeState.challenge.challengeName == "SELECT_MFA_TYPE" &&
+                            getMFATypeOrNull(challengeResponse) == null
+                        ) {
+                            val error = InvalidParameterException(
+                                message = "Value for challengeResponse must be one of SMS_MFA or SOFTWARE_TOKEN_MFA"
                             )
-                        )
-                        authStateMachine.send(event)
+                            onError.accept(error)
+                        } else {
+                            val event = SignInChallengeEvent(
+                                SignInChallengeEvent.EventType.VerifyChallengeAnswer(
+                                    challengeResponse,
+                                    metadata,
+                                    userAttributes
+                                )
+                            )
+                            authStateMachine.send(event)
+                        }
                     }
 
                     is SignInState.ResolvingTOTPSetup -> {

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/helpers/MFAHelper.kt

@@ -16,15 +16,18 @@
 package com.amplifyframework.auth.cognito.helpers
 
 import com.amplifyframework.auth.MFAType
-import kotlin.jvm.Throws
 
 @Throws(IllegalArgumentException::class)
-internal fun getMFAType(value: String): MFAType {
-    return when (value) {
-        "SMS_MFA" -> MFAType.SMS
-        "SOFTWARE_TOKEN_MFA" -> MFAType.TOTP
-        else -> throw IllegalArgumentException("Unsupported MFA type")
-    }
+internal fun getMFAType(value: String) = when (value) {
+    "SMS_MFA" -> MFAType.SMS
+    "SOFTWARE_TOKEN_MFA" -> MFAType.TOTP
+    else -> throw IllegalArgumentException("Unsupported MFA type")
+}
+
+internal fun getMFATypeOrNull(value: String) = when (value) {
+    "SMS_MFA" -> MFAType.SMS
+    "SOFTWARE_TOKEN_MFA" -> MFAType.TOTP
+    else -> null
 }
 
 internal val MFAType.value: String

aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/MFATypeUtilTest.kt

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amplifyframework.auth.cognito
+
+import com.amplifyframework.auth.MFAType
+import kotlin.test.assertEquals
+import org.junit.Test
+
+class MFATypeUtilTest {
+    @Test
+    fun challengeResponse_returns_correct_sms_string() {
+        assertEquals("SMS_MFA", MFAType.SMS.challengeResponse)
+    }
+
+    @Test
+    fun challengeResponse_returns_correct_totp_string() {
+        assertEquals("SOFTWARE_TOKEN_MFA", MFAType.TOTP.challengeResponse)
+    }
+}