feat(Auth): Add TOTP Support

aws-auth-cognito/build.gradle.kts

@@ -39,6 +39,8 @@ dependencies {
     implementation(dependency.aws.cognitoidentityprovider)
 
     testImplementation(project(":testutils"))
+    testImplementation(project(":core"))
+    testImplementation(project(":aws-core"))
     //noinspection GradleDependency
     testImplementation(testDependency.json)
 

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

@@ -20,6 +20,7 @@ import android.content.Context
 import android.content.Intent
 import androidx.annotation.VisibleForTesting
 import com.amplifyframework.AmplifyException
+import com.amplifyframework.TOTPSetupDetails
 import com.amplifyframework.annotations.InternalAmplifyApi
 import com.amplifyframework.auth.AWSCognitoAuthMetadataType
 import com.amplifyframework.auth.AuthCodeDeliveryDetails
@@ -32,6 +33,7 @@ import com.amplifyframework.auth.AuthUser
 import com.amplifyframework.auth.AuthUserAttribute
 import com.amplifyframework.auth.AuthUserAttributeKey
 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.setupMFA()
+                    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
@@ -88,6 +89,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)
                 else -> UnknownException(fallbackMessage, error)
             }
         }

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

@@ -17,6 +17,7 @@ package com.amplifyframework.auth.cognito
 
 import android.app.Activity
 import android.content.Intent
+import com.amplifyframework.TOTPSetupDetails
 import com.amplifyframework.auth.AuthCodeDeliveryDetails
 import com.amplifyframework.auth.AuthDevice
 import com.amplifyframework.auth.AuthProvider
@@ -38,6 +39,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
@@ -520,4 +522,46 @@ internal class KotlinAuthFacadeInternal(private val delegate: RealAWSCognitoAuth
             )
         }
     }
+
+    suspend fun setupMFA(): TOTPSetupDetails {
+        return suspendCoroutine { continuation ->
+            delegate.setUpTOTP(
+                { continuation.resume(it) },
+                { continuation.resumeWithException(it) }
+            )
+        }
+    }
+    suspend fun verifyTOTPSetup(code: String, options: AuthVerifyTOTPSetupOptions) {
+        return suspendCoroutine { continuation ->
+            delegate.verifyTOTPSetup(
+                code,
+                options,
+                { continuation.resume(Unit) },
+                { continuation.resumeWithException(it) }
+            )
+        }
+    }
+
+    suspend fun fetchMFAPreference(): UserMFAPreference {
+        return suspendCoroutine { continuation ->
+            delegate.fetchMFAPreference(
+                { continuation.resume(it) },
+                { continuation.resumeWithException(it) }
+            )
+        }
+    }
+
+    suspend fun updateMFAPreference(
+        sms: MFAPreference?,
+        totp: MFAPreference?
+    ) {
+        return suspendCoroutine { continuation ->
+            delegate.updateMFAPreference(
+                sms,
+                totp,
+                { continuation.resume(Unit) },
+                { continuation.resumeWithException(it) }
+            )
+        }
+    }
 }