Skip to content

Commit

Permalink
feat(Auth): Add TOTP Support (#2537) (#2568)
Browse files Browse the repository at this point in the history
Co-authored-by: Saijad Dhuka <[email protected]>
  • Loading branch information
gpanshu and sdhuka authored Aug 29, 2023
1 parent aac074a commit 19cc6e4
Show file tree
Hide file tree
Showing 50 changed files with 2,980 additions and 70 deletions.
4 changes: 4 additions & 0 deletions aws-auth-cognito/build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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"))
}
Original file line number Diff line number Diff line change
@@ -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 = "$userName@testdomain.com"

@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)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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())
}
}
)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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),
Expand Down Expand Up @@ -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(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Loading

0 comments on commit 19cc6e4

Please sign in to comment.