Skip to content

Commit 19cc6e4

Browse files
gpanshusdhuka
andauthored
feat(Auth): Add TOTP Support (#2537) (#2568)
Co-authored-by: Saijad Dhuka <[email protected]>
1 parent aac074a commit 19cc6e4

File tree

50 files changed

+2980
-70
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

50 files changed

+2980
-70
lines changed

aws-auth-cognito/build.gradle.kts

+4
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ dependencies {
3939
implementation(libs.aws.cognitoidentityprovider)
4040

4141
testImplementation(project(":testutils"))
42+
testImplementation(project(":core"))
43+
testImplementation(project(":aws-core"))
4244
//noinspection GradleDependency
4345
testImplementation(libs.test.json)
4446

@@ -60,6 +62,8 @@ dependencies {
6062
androidTestImplementation(libs.test.androidx.runner)
6163
androidTestImplementation(libs.test.androidx.junit)
6264
androidTestImplementation(libs.test.kotlin.coroutines)
65+
androidTestImplementation(libs.test.totp)
66+
6367
androidTestImplementation(project(":aws-api"))
6468
androidTestImplementation(project(":testutils"))
6569
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
/*
2+
* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
package com.amplifyframework.auth.cognito
16+
17+
import android.content.Context
18+
import androidx.test.core.app.ApplicationProvider
19+
import androidx.test.ext.junit.runners.AndroidJUnit4
20+
import com.amplifyframework.auth.AuthUserAttribute
21+
import com.amplifyframework.auth.AuthUserAttributeKey
22+
import com.amplifyframework.auth.MFAType
23+
import com.amplifyframework.auth.cognito.helpers.value
24+
import com.amplifyframework.auth.cognito.test.R
25+
import com.amplifyframework.auth.options.AuthSignUpOptions
26+
import com.amplifyframework.auth.result.step.AuthSignInStep
27+
import com.amplifyframework.core.AmplifyConfiguration
28+
import com.amplifyframework.core.category.CategoryConfiguration
29+
import com.amplifyframework.core.category.CategoryType
30+
import com.amplifyframework.testutils.sync.SynchronousAuth
31+
import dev.robinohs.totpkt.otp.totp.TotpGenerator
32+
import dev.robinohs.totpkt.otp.totp.timesupport.generateCode
33+
import java.util.Random
34+
import java.util.UUID
35+
import java.util.concurrent.CountDownLatch
36+
import java.util.concurrent.TimeUnit
37+
import org.junit.After
38+
import org.junit.Assert
39+
import org.junit.Before
40+
import org.junit.Test
41+
import org.junit.runner.RunWith
42+
43+
@RunWith(AndroidJUnit4::class)
44+
class AWSCognitoAuthPluginTOTPTests {
45+
46+
private lateinit var authPlugin: AWSCognitoAuthPlugin
47+
private lateinit var synchronousAuth: SynchronousAuth
48+
private val password = UUID.randomUUID().toString()
49+
private val userName = "testUser${Random().nextInt()}"
50+
private val email = "$userName@testdomain.com"
51+
52+
@Before
53+
fun setup() {
54+
val context = ApplicationProvider.getApplicationContext<Context>()
55+
val config = AmplifyConfiguration.fromConfigFile(context, R.raw.amplifyconfiguration_totp)
56+
val authConfig: CategoryConfiguration = config.forCategoryType(CategoryType.AUTH)
57+
val authConfigJson = authConfig.getPluginConfig("awsCognitoAuthPlugin")
58+
authPlugin = AWSCognitoAuthPlugin()
59+
authPlugin.configure(authConfigJson, context)
60+
synchronousAuth = SynchronousAuth.delegatingTo(authPlugin)
61+
signUpNewUser(userName, password, email)
62+
synchronousAuth.signOut()
63+
}
64+
65+
@After
66+
fun tearDown() {
67+
synchronousAuth.deleteUser()
68+
}
69+
70+
/*
71+
* This test signs up a new user and goes thru successful MFA Setup process.
72+
* */
73+
@Test
74+
fun mfa_setup() {
75+
val result = synchronousAuth.signIn(userName, password)
76+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
77+
val otp = TotpGenerator().generateCode(
78+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
79+
System.currentTimeMillis()
80+
)
81+
synchronousAuth.confirmSignIn(otp)
82+
val currentUser = synchronousAuth.currentUser
83+
Assert.assertEquals(userName.lowercase(), currentUser.username)
84+
}
85+
86+
/*
87+
* This test signs up a new user, enter incorrect MFA code during verification and
88+
* then enter correct OTP code to successfully set TOTP MFA.
89+
* */
90+
@Test
91+
fun mfasetup_with_incorrect_otp() {
92+
val result = synchronousAuth.signIn(userName, password)
93+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
94+
try {
95+
synchronousAuth.confirmSignIn("123456")
96+
} catch (e: Exception) {
97+
Assert.assertEquals("Code mismatch", e.cause?.message)
98+
val otp = TotpGenerator().generateCode(
99+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
100+
System.currentTimeMillis()
101+
)
102+
synchronousAuth.confirmSignIn(otp)
103+
val currentUser = synchronousAuth.currentUser
104+
Assert.assertEquals(userName.lowercase(), currentUser.username)
105+
}
106+
}
107+
108+
/*
109+
* This test signs up a new user, successfully setup MFA, sign-out and then goes thru sign-in with TOTP.
110+
* */
111+
@Test
112+
fun signIn_with_totp_after_mfa_setup() {
113+
val result = synchronousAuth.signIn(userName, password)
114+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
115+
val otp = TotpGenerator().generateCode(
116+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray()
117+
)
118+
synchronousAuth.confirmSignIn(otp)
119+
synchronousAuth.signOut()
120+
121+
val signInResult = synchronousAuth.signIn(userName, password)
122+
Assert.assertEquals(AuthSignInStep.CONFIRM_SIGN_IN_WITH_TOTP_CODE, signInResult.nextStep.signInStep)
123+
val otpCode = TotpGenerator().generateCode(
124+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
125+
System.currentTimeMillis() + 30 * 1000 // 30 sec is added to generate new OTP code
126+
)
127+
synchronousAuth.confirmSignIn(otpCode)
128+
val currentUser = synchronousAuth.currentUser
129+
Assert.assertEquals(userName.lowercase(), currentUser.username)
130+
}
131+
132+
/*
133+
* This test signs up a new user, successfully setup MFA, update user attribute to add phone number,
134+
* sign-out the user, goes thru MFA selection flow during sign-in, select TOTP MFA type,
135+
* successfully sign-in using TOTP
136+
* */
137+
@Test
138+
fun select_mfa_type() {
139+
val result = synchronousAuth.signIn(userName, password)
140+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
141+
val otp = TotpGenerator().generateCode(
142+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray()
143+
)
144+
synchronousAuth.confirmSignIn(otp)
145+
synchronousAuth.updateUserAttribute(AuthUserAttribute(AuthUserAttributeKey.phoneNumber(), "+19876543210"))
146+
updateMFAPreference(MFAPreference.ENABLED, MFAPreference.ENABLED)
147+
synchronousAuth.signOut()
148+
val signInResult = synchronousAuth.signIn(userName, password)
149+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_MFA_SELECTION, signInResult.nextStep.signInStep)
150+
val totpSignInResult = synchronousAuth.confirmSignIn(MFAType.TOTP.value)
151+
Assert.assertEquals(AuthSignInStep.CONFIRM_SIGN_IN_WITH_TOTP_CODE, totpSignInResult.nextStep.signInStep)
152+
val otpCode = TotpGenerator().generateCode(
153+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
154+
System.currentTimeMillis() + 30 * 1000 // 30 sec is added to generate new OTP code
155+
)
156+
synchronousAuth.confirmSignIn(otpCode)
157+
val currentUser = synchronousAuth.currentUser
158+
Assert.assertEquals(userName.lowercase(), currentUser.username)
159+
}
160+
161+
private fun signUpNewUser(userName: String, password: String, email: String) {
162+
val options = AuthSignUpOptions.builder()
163+
.userAttributes(
164+
listOf(
165+
AuthUserAttribute(AuthUserAttributeKey.email(), email)
166+
)
167+
).build()
168+
synchronousAuth.signUp(userName, password, options)
169+
}
170+
171+
private fun updateMFAPreference(sms: MFAPreference, totp: MFAPreference) {
172+
val latch = CountDownLatch(1)
173+
authPlugin.updateMFAPreference(sms, totp, { latch.countDown() }, { latch.countDown() })
174+
latch.await(5, TimeUnit.SECONDS)
175+
}
176+
}

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

+71
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,9 @@ import com.amplifyframework.auth.AuthSession
3131
import com.amplifyframework.auth.AuthUser
3232
import com.amplifyframework.auth.AuthUserAttribute
3333
import com.amplifyframework.auth.AuthUserAttributeKey
34+
import com.amplifyframework.auth.TOTPSetupDetails
3435
import com.amplifyframework.auth.cognito.asf.UserContextDataProvider
36+
import com.amplifyframework.auth.cognito.options.AWSCognitoAuthVerifyTOTPSetupOptions
3537
import com.amplifyframework.auth.cognito.options.FederateToIdentityPoolOptions
3638
import com.amplifyframework.auth.cognito.result.FederateToIdentityPoolResult
3739
import com.amplifyframework.auth.exceptions.ConfigurationException
@@ -48,6 +50,7 @@ import com.amplifyframework.auth.options.AuthSignOutOptions
4850
import com.amplifyframework.auth.options.AuthSignUpOptions
4951
import com.amplifyframework.auth.options.AuthUpdateUserAttributeOptions
5052
import com.amplifyframework.auth.options.AuthUpdateUserAttributesOptions
53+
import com.amplifyframework.auth.options.AuthVerifyTOTPSetupOptions
5154
import com.amplifyframework.auth.options.AuthWebUISignInOptions
5255
import com.amplifyframework.auth.result.AuthResetPasswordResult
5356
import com.amplifyframework.auth.result.AuthSignInResult
@@ -768,6 +771,40 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
768771
)
769772
}
770773

774+
override fun setUpTOTP(onSuccess: Consumer<TOTPSetupDetails>, onError: Consumer<AuthException>) {
775+
queueChannel.trySend(
776+
pluginScope.launch(start = CoroutineStart.LAZY) {
777+
try {
778+
val result = queueFacade.setUpTOTP()
779+
onSuccess.accept(result)
780+
} catch (e: Exception) {
781+
onError.accept(e.toAuthException())
782+
}
783+
}
784+
)
785+
}
786+
787+
override fun verifyTOTPSetup(code: String, onSuccess: Action, onError: Consumer<AuthException>) {
788+
verifyTOTPSetup(code, AWSCognitoAuthVerifyTOTPSetupOptions.CognitoBuilder().build(), onSuccess, onError)
789+
}
790+
791+
override fun verifyTOTPSetup(
792+
code: String,
793+
options: AuthVerifyTOTPSetupOptions,
794+
onSuccess: Action,
795+
onError: Consumer<AuthException>
796+
) {
797+
queueChannel.trySend(
798+
pluginScope.launch(start = CoroutineStart.LAZY) {
799+
try {
800+
queueFacade.verifyTOTPSetup(code, options)
801+
onSuccess.call()
802+
} catch (e: Exception) {
803+
onError.accept(e.toAuthException())
804+
}
805+
}
806+
)
807+
}
771808
override fun getEscapeHatch() = realPlugin.escapeHatch()
772809

773810
override fun getPluginKey() = AWS_COGNITO_AUTH_PLUGIN_KEY
@@ -852,4 +889,38 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
852889
}
853890
)
854891
}
892+
893+
fun fetchMFAPreference(
894+
onSuccess: Consumer<UserMFAPreference>,
895+
onError: Consumer<AuthException>
896+
) {
897+
queueChannel.trySend(
898+
pluginScope.launch(start = CoroutineStart.LAZY) {
899+
try {
900+
val result = queueFacade.fetchMFAPreference()
901+
onSuccess.accept(result)
902+
} catch (e: Exception) {
903+
onError.accept(e.toAuthException())
904+
}
905+
}
906+
)
907+
}
908+
909+
fun updateMFAPreference(
910+
sms: MFAPreference?,
911+
totp: MFAPreference?,
912+
onSuccess: Action,
913+
onError: Consumer<AuthException>
914+
) {
915+
queueChannel.trySend(
916+
pluginScope.launch(start = CoroutineStart.LAZY) {
917+
try {
918+
queueFacade.updateMFAPreference(sms, totp)
919+
onSuccess.call()
920+
} catch (e: Exception) {
921+
onError.accept(e.toAuthException())
922+
}
923+
}
924+
)
925+
}
855926
}

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

+6-2
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import com.amplifyframework.auth.cognito.actions.FetchAuthSessionCognitoActions
2424
import com.amplifyframework.auth.cognito.actions.HostedUICognitoActions
2525
import com.amplifyframework.auth.cognito.actions.MigrateAuthCognitoActions
2626
import com.amplifyframework.auth.cognito.actions.SRPCognitoActions
27+
import com.amplifyframework.auth.cognito.actions.SetupTOTPCognitoActions
2728
import com.amplifyframework.auth.cognito.actions.SignInChallengeCognitoActions
2829
import com.amplifyframework.auth.cognito.actions.SignInCognitoActions
2930
import com.amplifyframework.auth.cognito.actions.SignInCustomCognitoActions
@@ -42,6 +43,7 @@ import com.amplifyframework.statemachine.codegen.states.HostedUISignInState
4243
import com.amplifyframework.statemachine.codegen.states.MigrateSignInState
4344
import com.amplifyframework.statemachine.codegen.states.RefreshSessionState
4445
import com.amplifyframework.statemachine.codegen.states.SRPSignInState
46+
import com.amplifyframework.statemachine.codegen.states.SetupTOTPState
4547
import com.amplifyframework.statemachine.codegen.states.SignInChallengeState
4648
import com.amplifyframework.statemachine.codegen.states.SignInState
4749
import com.amplifyframework.statemachine.codegen.states.SignOutState
@@ -62,10 +64,11 @@ internal class AuthStateMachine(
6264
SignInChallengeState.Resolver(SignInChallengeCognitoActions),
6365
HostedUISignInState.Resolver(HostedUICognitoActions),
6466
DeviceSRPSignInState.Resolver(DeviceSRPCognitoSignInActions),
67+
SetupTOTPState.Resolver(SetupTOTPCognitoActions),
6568
SignInCognitoActions
6669
),
6770
SignOutState.Resolver(SignOutCognitoActions),
68-
AuthenticationCognitoActions,
71+
AuthenticationCognitoActions
6972
),
7073
AuthorizationState.Resolver(
7174
FetchAuthSessionState.Resolver(FetchAuthSessionCognitoActions),
@@ -93,10 +96,11 @@ internal class AuthStateMachine(
9396
SignInChallengeState.Resolver(SignInChallengeCognitoActions).logging(),
9497
HostedUISignInState.Resolver(HostedUICognitoActions).logging(),
9598
DeviceSRPSignInState.Resolver(DeviceSRPCognitoSignInActions).logging(),
99+
SetupTOTPState.Resolver(SetupTOTPCognitoActions).logging(),
96100
SignInCognitoActions
97101
).logging(),
98102
SignOutState.Resolver(SignOutCognitoActions).logging(),
99-
AuthenticationCognitoActions,
103+
AuthenticationCognitoActions
100104
).logging(),
101105
AuthorizationState.Resolver(
102106
FetchAuthSessionState.Resolver(FetchAuthSessionCognitoActions).logging(),

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

+3
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ package com.amplifyframework.auth.cognito
1818
import aws.sdk.kotlin.services.cognitoidentityprovider.model.AliasExistsException
1919
import aws.sdk.kotlin.services.cognitoidentityprovider.model.CodeDeliveryFailureException
2020
import aws.sdk.kotlin.services.cognitoidentityprovider.model.CodeMismatchException
21+
import aws.sdk.kotlin.services.cognitoidentityprovider.model.EnableSoftwareTokenMfaException
2122
import aws.sdk.kotlin.services.cognitoidentityprovider.model.ExpiredCodeException
2223
import aws.sdk.kotlin.services.cognitoidentityprovider.model.InvalidParameterException
2324
import aws.sdk.kotlin.services.cognitoidentityprovider.model.InvalidPasswordException
@@ -89,6 +90,8 @@ internal class CognitoAuthExceptionConverter {
8990
com.amplifyframework.auth.cognito.exceptions.service.TooManyRequestsException(error)
9091
is PasswordResetRequiredException ->
9192
com.amplifyframework.auth.cognito.exceptions.service.PasswordResetRequiredException(error)
93+
is EnableSoftwareTokenMfaException ->
94+
com.amplifyframework.auth.cognito.exceptions.service.EnableSoftwareTokenMFAException(error)
9295
is UserLambdaValidationException ->
9396
com.amplifyframework.auth.cognito.exceptions.service.UserLambdaValidationException(
9497
error.message,

0 commit comments

Comments
 (0)