feat(Auth): Add TOTP Support (#2537)

aws-auth-cognito/build.gradle.kts

@@ -39,6 +39,8 @@ dependencies {
     implementation(libs.aws.cognitoidentityprovider)
 
     testImplementation(project(":testutils"))
+    testImplementation(project(":core"))
+    testImplementation(project(":aws-core"))
     //noinspection GradleDependency
     testImplementation(libs.test.json)
 
@@ -60,6 +62,8 @@ dependencies {
     androidTestImplementation(libs.test.androidx.runner)
     androidTestImplementation(libs.test.androidx.junit)
     androidTestImplementation(libs.test.kotlin.coroutines)
+    androidTestImplementation(libs.test.totp)
+
     androidTestImplementation(project(":aws-api"))
     androidTestImplementation(project(":testutils"))
 }

aws-auth-cognito/src/androidTest/java/com/amplifyframework/auth/cognito/AWSCognitoAuthPluginTOTPTests.kt

@@ -0,0 +1,176 @@
+/*
+ * Copyright 2023 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 android.content.Context
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.amplifyframework.auth.AuthUserAttribute
+import com.amplifyframework.auth.AuthUserAttributeKey
+import com.amplifyframework.auth.MFAType
+import com.amplifyframework.auth.cognito.helpers.value
+import com.amplifyframework.auth.cognito.test.R
+import com.amplifyframework.auth.options.AuthSignUpOptions
+import com.amplifyframework.auth.result.step.AuthSignInStep
+import com.amplifyframework.core.AmplifyConfiguration
+import com.amplifyframework.core.category.CategoryConfiguration
+import com.amplifyframework.core.category.CategoryType
+import com.amplifyframework.testutils.sync.SynchronousAuth
+import dev.robinohs.totpkt.otp.totp.TotpGenerator
+import dev.robinohs.totpkt.otp.totp.timesupport.generateCode
+import java.util.Random
+import java.util.UUID
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import org.junit.After
+import org.junit.Assert
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class AWSCognitoAuthPluginTOTPTests {
+
+    private lateinit var authPlugin: AWSCognitoAuthPlugin
+    private lateinit var synchronousAuth: SynchronousAuth
+    private val password = UUID.randomUUID().toString()
+    private val userName = "testUser${Random().nextInt()}"
+    private val email = "[email protected]"
+
+    @Before
+    fun setup() {
+        val context = ApplicationProvider.getApplicationContext<Context>()
+        val config = AmplifyConfiguration.fromConfigFile(context, R.raw.amplifyconfiguration_totp)
+        val authConfig: CategoryConfiguration = config.forCategoryType(CategoryType.AUTH)
+        val authConfigJson = authConfig.getPluginConfig("awsCognitoAuthPlugin")
+        authPlugin = AWSCognitoAuthPlugin()
+        authPlugin.configure(authConfigJson, context)
+        synchronousAuth = SynchronousAuth.delegatingTo(authPlugin)
+        signUpNewUser(userName, password, email)
+        synchronousAuth.signOut()
+    }
+
+    @After
+    fun tearDown() {
+        synchronousAuth.deleteUser()
+    }
+
+    /*
+    * This test signs up a new user and goes thru successful MFA Setup process.
+    * */
+    @Test
+    fun mfa_setup() {
+        val result = synchronousAuth.signIn(userName, password)
+        Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
+        val otp = TotpGenerator().generateCode(
+            result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
+            System.currentTimeMillis()
+        )
+        synchronousAuth.confirmSignIn(otp)
+        val currentUser = synchronousAuth.currentUser
+        Assert.assertEquals(userName.lowercase(), currentUser.username)
+    }
+
+    /*
+    * This test signs up a new user, enter incorrect MFA code during verification and
+    * then enter correct OTP code to successfully set TOTP MFA.
+    * */
+    @Test
+    fun mfasetup_with_incorrect_otp() {
+        val result = synchronousAuth.signIn(userName, password)
+        Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
+        try {
+            synchronousAuth.confirmSignIn("123456")
+        } catch (e: Exception) {
+            Assert.assertEquals("Code mismatch", e.cause?.message)
+            val otp = TotpGenerator().generateCode(
+                result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
+                System.currentTimeMillis()
+            )
+            synchronousAuth.confirmSignIn(otp)
+            val currentUser = synchronousAuth.currentUser
+            Assert.assertEquals(userName.lowercase(), currentUser.username)
+        }
+    }
+
+    /*
+    * This test signs up a new user, successfully setup MFA, sign-out and then goes thru sign-in with TOTP.
+    * */
+    @Test
+    fun signIn_with_totp_after_mfa_setup() {
+        val result = synchronousAuth.signIn(userName, password)
+        Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
+        val otp = TotpGenerator().generateCode(
+            result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray()
+        )
+        synchronousAuth.confirmSignIn(otp)
+        synchronousAuth.signOut()
+
+        val signInResult = synchronousAuth.signIn(userName, password)
+        Assert.assertEquals(AuthSignInStep.CONFIRM_SIGN_IN_WITH_TOTP_CODE, signInResult.nextStep.signInStep)
+        val otpCode = TotpGenerator().generateCode(
+            result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
+            System.currentTimeMillis() + 30 * 1000 // 30 sec is added to generate new OTP code
+        )
+        synchronousAuth.confirmSignIn(otpCode)
+        val currentUser = synchronousAuth.currentUser
+        Assert.assertEquals(userName.lowercase(), currentUser.username)
+    }
+
+    /*
+    * This test signs up a new user, successfully setup MFA, update user attribute to add phone number,
+    * sign-out the user, goes thru MFA selection flow during sign-in, select TOTP MFA type,
+    * successfully sign-in using TOTP
+    * */
+    @Test
+    fun select_mfa_type() {
+        val result = synchronousAuth.signIn(userName, password)
+        Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
+        val otp = TotpGenerator().generateCode(
+            result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray()
+        )
+        synchronousAuth.confirmSignIn(otp)
+        synchronousAuth.updateUserAttribute(AuthUserAttribute(AuthUserAttributeKey.phoneNumber(), "+19876543210"))
+        updateMFAPreference(MFAPreference.ENABLED, MFAPreference.ENABLED)
+        synchronousAuth.signOut()
+        val signInResult = synchronousAuth.signIn(userName, password)
+        Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_MFA_SELECTION, signInResult.nextStep.signInStep)
+        val totpSignInResult = synchronousAuth.confirmSignIn(MFAType.TOTP.value)
+        Assert.assertEquals(AuthSignInStep.CONFIRM_SIGN_IN_WITH_TOTP_CODE, totpSignInResult.nextStep.signInStep)
+        val otpCode = TotpGenerator().generateCode(
+            result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
+            System.currentTimeMillis() + 30 * 1000 // 30 sec is added to generate new OTP code
+        )
+        synchronousAuth.confirmSignIn(otpCode)
+        val currentUser = synchronousAuth.currentUser
+        Assert.assertEquals(userName.lowercase(), currentUser.username)
+    }
+
+    private fun signUpNewUser(userName: String, password: String, email: String) {
+        val options = AuthSignUpOptions.builder()
+            .userAttributes(
+                listOf(
+                    AuthUserAttribute(AuthUserAttributeKey.email(), email)
+                )
+            ).build()
+        synchronousAuth.signUp(userName, password, options)
+    }
+
+    private fun updateMFAPreference(sms: MFAPreference, totp: MFAPreference) {
+        val latch = CountDownLatch(1)
+        authPlugin.updateMFAPreference(sms, totp, { latch.countDown() }, { latch.countDown() })
+        latch.await(5, TimeUnit.SECONDS)
+    }
+}

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

@@ -31,7 +31,9 @@ import com.amplifyframework.auth.AuthSession
 import com.amplifyframework.auth.AuthUser
 import com.amplifyframework.auth.AuthUserAttribute
 import com.amplifyframework.auth.AuthUserAttributeKey
+import com.amplifyframework.auth.TOTPSetupDetails
 import com.amplifyframework.auth.cognito.asf.UserContextDataProvider
+import com.amplifyframework.auth.cognito.options.AWSCognitoAuthVerifyTOTPSetupOptions
 import com.amplifyframework.auth.cognito.options.FederateToIdentityPoolOptions
 import com.amplifyframework.auth.cognito.result.FederateToIdentityPoolResult
 import com.amplifyframework.auth.exceptions.ConfigurationException
@@ -48,6 +50,7 @@ import com.amplifyframework.auth.options.AuthSignOutOptions
 import com.amplifyframework.auth.options.AuthSignUpOptions
 import com.amplifyframework.auth.options.AuthUpdateUserAttributeOptions
 import com.amplifyframework.auth.options.AuthUpdateUserAttributesOptions
+import com.amplifyframework.auth.options.AuthVerifyTOTPSetupOptions
 import com.amplifyframework.auth.options.AuthWebUISignInOptions
 import com.amplifyframework.auth.result.AuthResetPasswordResult
 import com.amplifyframework.auth.result.AuthSignInResult
@@ -768,6 +771,40 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
         )
     }
 
+    override fun setUpTOTP(onSuccess: Consumer<TOTPSetupDetails>, onError: Consumer<AuthException>) {
+        queueChannel.trySend(
+            pluginScope.launch(start = CoroutineStart.LAZY) {
+                try {
+                    val result = queueFacade.setUpTOTP()
+                    onSuccess.accept(result)
+                } catch (e: Exception) {
+                    onError.accept(e.toAuthException())
+                }
+            }
+        )
+    }
+
+    override fun verifyTOTPSetup(code: String, onSuccess: Action, onError: Consumer<AuthException>) {
+        verifyTOTPSetup(code, AWSCognitoAuthVerifyTOTPSetupOptions.CognitoBuilder().build(), onSuccess, onError)
+    }
+
+    override fun verifyTOTPSetup(
+        code: String,
+        options: AuthVerifyTOTPSetupOptions,
+        onSuccess: Action,
+        onError: Consumer<AuthException>
+    ) {
+        queueChannel.trySend(
+            pluginScope.launch(start = CoroutineStart.LAZY) {
+                try {
+                    queueFacade.verifyTOTPSetup(code, options)
+                    onSuccess.call()
+                } catch (e: Exception) {
+                    onError.accept(e.toAuthException())
+                }
+            }
+        )
+    }
     override fun getEscapeHatch() = realPlugin.escapeHatch()
 
     override fun getPluginKey() = AWS_COGNITO_AUTH_PLUGIN_KEY
@@ -852,4 +889,38 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
             }
         )
     }
+
+    fun fetchMFAPreference(
+        onSuccess: Consumer<UserMFAPreference>,
+        onError: Consumer<AuthException>
+    ) {
+        queueChannel.trySend(
+            pluginScope.launch(start = CoroutineStart.LAZY) {
+                try {
+                    val result = queueFacade.fetchMFAPreference()
+                    onSuccess.accept(result)
+                } catch (e: Exception) {
+                    onError.accept(e.toAuthException())
+                }
+            }
+        )
+    }
+
+    fun updateMFAPreference(
+        sms: MFAPreference?,
+        totp: MFAPreference?,
+        onSuccess: Action,
+        onError: Consumer<AuthException>
+    ) {
+        queueChannel.trySend(
+            pluginScope.launch(start = CoroutineStart.LAZY) {
+                try {
+                    queueFacade.updateMFAPreference(sms, totp)
+                    onSuccess.call()
+                } catch (e: Exception) {
+                    onError.accept(e.toAuthException())
+                }
+            }
+        )
+    }
 }

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

@@ -24,6 +24,7 @@ import com.amplifyframework.auth.cognito.actions.FetchAuthSessionCognitoActions
 import com.amplifyframework.auth.cognito.actions.HostedUICognitoActions
 import com.amplifyframework.auth.cognito.actions.MigrateAuthCognitoActions
 import com.amplifyframework.auth.cognito.actions.SRPCognitoActions
+import com.amplifyframework.auth.cognito.actions.SetupTOTPCognitoActions
 import com.amplifyframework.auth.cognito.actions.SignInChallengeCognitoActions
 import com.amplifyframework.auth.cognito.actions.SignInCognitoActions
 import com.amplifyframework.auth.cognito.actions.SignInCustomCognitoActions
@@ -42,6 +43,7 @@ import com.amplifyframework.statemachine.codegen.states.HostedUISignInState
 import com.amplifyframework.statemachine.codegen.states.MigrateSignInState
 import com.amplifyframework.statemachine.codegen.states.RefreshSessionState
 import com.amplifyframework.statemachine.codegen.states.SRPSignInState
+import com.amplifyframework.statemachine.codegen.states.SetupTOTPState
 import com.amplifyframework.statemachine.codegen.states.SignInChallengeState
 import com.amplifyframework.statemachine.codegen.states.SignInState
 import com.amplifyframework.statemachine.codegen.states.SignOutState
@@ -62,10 +64,11 @@ internal class AuthStateMachine(
                     SignInChallengeState.Resolver(SignInChallengeCognitoActions),
                     HostedUISignInState.Resolver(HostedUICognitoActions),
                     DeviceSRPSignInState.Resolver(DeviceSRPCognitoSignInActions),
+                    SetupTOTPState.Resolver(SetupTOTPCognitoActions),
                     SignInCognitoActions
                 ),
                 SignOutState.Resolver(SignOutCognitoActions),
-                AuthenticationCognitoActions,
+                AuthenticationCognitoActions
             ),
             AuthorizationState.Resolver(
                 FetchAuthSessionState.Resolver(FetchAuthSessionCognitoActions),
@@ -93,10 +96,11 @@ internal class AuthStateMachine(
                         SignInChallengeState.Resolver(SignInChallengeCognitoActions).logging(),
                         HostedUISignInState.Resolver(HostedUICognitoActions).logging(),
                         DeviceSRPSignInState.Resolver(DeviceSRPCognitoSignInActions).logging(),
+                        SetupTOTPState.Resolver(SetupTOTPCognitoActions).logging(),
                         SignInCognitoActions
                     ).logging(),
                     SignOutState.Resolver(SignOutCognitoActions).logging(),
-                    AuthenticationCognitoActions,
+                    AuthenticationCognitoActions
                 ).logging(),
                 AuthorizationState.Resolver(
                     FetchAuthSessionState.Resolver(FetchAuthSessionCognitoActions).logging(),

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

@@ -18,6 +18,7 @@ package com.amplifyframework.auth.cognito
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.AliasExistsException
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.CodeDeliveryFailureException
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.CodeMismatchException
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.EnableSoftwareTokenMfaException
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ExpiredCodeException
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.InvalidParameterException
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.InvalidPasswordException
@@ -89,6 +90,8 @@ internal class CognitoAuthExceptionConverter {
                     com.amplifyframework.auth.cognito.exceptions.service.TooManyRequestsException(error)
                 is PasswordResetRequiredException ->
                     com.amplifyframework.auth.cognito.exceptions.service.PasswordResetRequiredException(error)
+                is EnableSoftwareTokenMfaException ->
+                    com.amplifyframework.auth.cognito.exceptions.service.EnableSoftwareTokenMFAException(error)
                 is UserLambdaValidationException ->
                     com.amplifyframework.auth.cognito.exceptions.service.UserLambdaValidationException(
                         error.message,