Skip to content

Commit 086697a

Browse files
authored
feat(Auth): Add TOTP Support (#2537)
1 parent 0882075 commit 086697a

File tree

45 files changed

+2355
-81
lines changed

Some content is hidden

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

45 files changed

+2355
-81
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,179 @@
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.test.R
24+
import com.amplifyframework.auth.options.AuthSignUpOptions
25+
import com.amplifyframework.auth.result.step.AuthSignInStep
26+
import com.amplifyframework.core.Amplify
27+
import com.amplifyframework.core.AmplifyConfiguration
28+
import com.amplifyframework.core.category.CategoryConfiguration
29+
import com.amplifyframework.core.category.CategoryType
30+
import com.amplifyframework.logging.AndroidLoggingPlugin
31+
import com.amplifyframework.logging.LogLevel
32+
import com.amplifyframework.testutils.sync.SynchronousAuth
33+
import dev.robinohs.totpkt.otp.totp.TotpGenerator
34+
import dev.robinohs.totpkt.otp.totp.timesupport.generateCode
35+
import java.util.Random
36+
import java.util.UUID
37+
import java.util.concurrent.CountDownLatch
38+
import java.util.concurrent.TimeUnit
39+
import org.junit.After
40+
import org.junit.Assert
41+
import org.junit.Before
42+
import org.junit.Test
43+
import org.junit.runner.RunWith
44+
45+
@RunWith(AndroidJUnit4::class)
46+
class AWSCognitoAuthPluginTOTPTests {
47+
48+
private lateinit var authPlugin: AWSCognitoAuthPlugin
49+
private lateinit var synchronousAuth: SynchronousAuth
50+
private val password = UUID.randomUUID().toString()
51+
private val userName = "testUser${Random().nextInt()}"
52+
private val email = "$userName@testdomain.com"
53+
54+
@Before
55+
fun setup() {
56+
val context = ApplicationProvider.getApplicationContext<Context>()
57+
Amplify.addPlugin(AndroidLoggingPlugin(LogLevel.VERBOSE))
58+
val config = AmplifyConfiguration.fromConfigFile(context, R.raw.amplifyconfiguration_totp)
59+
val authConfig: CategoryConfiguration = config.forCategoryType(CategoryType.AUTH)
60+
val authConfigJson = authConfig.getPluginConfig("awsCognitoAuthPlugin")
61+
authPlugin = AWSCognitoAuthPlugin()
62+
authPlugin.configure(authConfigJson, context)
63+
synchronousAuth = SynchronousAuth.delegatingTo(authPlugin)
64+
signUpNewUser(userName, password, email)
65+
synchronousAuth.signOut()
66+
}
67+
68+
@After
69+
fun tearDown() {
70+
synchronousAuth.deleteUser()
71+
}
72+
73+
/*
74+
* This test signs up a new user and goes thru successful MFA Setup process.
75+
* */
76+
@Test
77+
fun mfa_setup() {
78+
val result = synchronousAuth.signIn(userName, password)
79+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
80+
val otp = TotpGenerator().generateCode(
81+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
82+
System.currentTimeMillis()
83+
)
84+
synchronousAuth.confirmSignIn(otp)
85+
val currentUser = synchronousAuth.currentUser
86+
Assert.assertEquals(userName.lowercase(), currentUser.username)
87+
}
88+
89+
/*
90+
* This test signs up a new user, enter incorrect MFA code during verification and
91+
* then enter correct OTP code to successfully set TOTP MFA.
92+
* */
93+
@Test
94+
fun mfasetup_with_incorrect_otp() {
95+
val result = synchronousAuth.signIn(userName, password)
96+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
97+
try {
98+
synchronousAuth.confirmSignIn("123456")
99+
} catch (e: Exception) {
100+
Assert.assertEquals("Code mismatch", e.cause?.message)
101+
val otp = TotpGenerator().generateCode(
102+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
103+
System.currentTimeMillis()
104+
)
105+
synchronousAuth.confirmSignIn(otp)
106+
val currentUser = synchronousAuth.currentUser
107+
Assert.assertEquals(userName.lowercase(), currentUser.username)
108+
}
109+
}
110+
111+
/*
112+
* This test signs up a new user, successfully setup MFA, sign-out and then goes thru sign-in with TOTP.
113+
* */
114+
@Test
115+
fun signIn_with_totp_after_mfa_setup() {
116+
val result = synchronousAuth.signIn(userName, password)
117+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
118+
val otp = TotpGenerator().generateCode(
119+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray()
120+
)
121+
synchronousAuth.confirmSignIn(otp)
122+
synchronousAuth.signOut()
123+
124+
val signInResult = synchronousAuth.signIn(userName, password)
125+
Assert.assertEquals(AuthSignInStep.CONFIRM_SIGN_IN_WITH_TOTP_CODE, signInResult.nextStep.signInStep)
126+
val otpCode = TotpGenerator().generateCode(
127+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
128+
System.currentTimeMillis() + 30 * 1000 // 30 sec is added to generate new OTP code
129+
)
130+
synchronousAuth.confirmSignIn(otpCode)
131+
val currentUser = synchronousAuth.currentUser
132+
Assert.assertEquals(userName.lowercase(), currentUser.username)
133+
}
134+
135+
/*
136+
* This test signs up a new user, successfully setup MFA, update user attribute to add phone number,
137+
* sign-out the user, goes thru MFA selection flow during sign-in, select TOTP MFA type,
138+
* successfully sign-in using TOTP
139+
* */
140+
@Test
141+
fun select_mfa_type() {
142+
val result = synchronousAuth.signIn(userName, password)
143+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_TOTP_SETUP, result.nextStep.signInStep)
144+
val otp = TotpGenerator().generateCode(
145+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray()
146+
)
147+
synchronousAuth.confirmSignIn(otp)
148+
synchronousAuth.updateUserAttribute(AuthUserAttribute(AuthUserAttributeKey.phoneNumber(), "+19876543210"))
149+
updateMFAPreference(MFAPreference.Enabled, MFAPreference.Enabled)
150+
synchronousAuth.signOut()
151+
val signInResult = synchronousAuth.signIn(userName, password)
152+
Assert.assertEquals(AuthSignInStep.CONTINUE_SIGN_IN_WITH_MFA_SELECTION, signInResult.nextStep.signInStep)
153+
val totpSignInResult = synchronousAuth.confirmSignIn(MFAType.TOTP.value)
154+
Assert.assertEquals(AuthSignInStep.CONFIRM_SIGN_IN_WITH_TOTP_CODE, totpSignInResult.nextStep.signInStep)
155+
val otpCode = TotpGenerator().generateCode(
156+
result.nextStep.totpSetupDetails!!.sharedSecret.toByteArray(),
157+
System.currentTimeMillis() + 30 * 1000 // 30 sec is added to generate new OTP code
158+
)
159+
synchronousAuth.confirmSignIn(otpCode)
160+
val currentUser = synchronousAuth.currentUser
161+
Assert.assertEquals(userName.lowercase(), currentUser.username)
162+
}
163+
164+
private fun signUpNewUser(userName: String, password: String, email: String) {
165+
val options = AuthSignUpOptions.builder()
166+
.userAttributes(
167+
listOf(
168+
AuthUserAttribute(AuthUserAttributeKey.email(), email)
169+
)
170+
).build()
171+
synchronousAuth.signUp(userName, password, options)
172+
}
173+
174+
private fun updateMFAPreference(sms: MFAPreference, totp: MFAPreference) {
175+
val latch = CountDownLatch(1)
176+
authPlugin.updateMFAPreference(sms, totp, { latch.countDown() }, { latch.countDown() })
177+
latch.await(5, TimeUnit.SECONDS)
178+
}
179+
}

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

+71
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import android.content.Context
2020
import android.content.Intent
2121
import androidx.annotation.VisibleForTesting
2222
import com.amplifyframework.AmplifyException
23+
import com.amplifyframework.TOTPSetupDetails
2324
import com.amplifyframework.annotations.InternalAmplifyApi
2425
import com.amplifyframework.auth.AWSCognitoAuthMetadataType
2526
import com.amplifyframework.auth.AuthCodeDeliveryDetails
@@ -32,6 +33,7 @@ import com.amplifyframework.auth.AuthUser
3233
import com.amplifyframework.auth.AuthUserAttribute
3334
import com.amplifyframework.auth.AuthUserAttributeKey
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.setupMFA()
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
@@ -88,6 +89,8 @@ internal class CognitoAuthExceptionConverter {
8889
com.amplifyframework.auth.cognito.exceptions.service.TooManyRequestsException(error)
8990
is PasswordResetRequiredException ->
9091
com.amplifyframework.auth.cognito.exceptions.service.PasswordResetRequiredException(error)
92+
is EnableSoftwareTokenMfaException ->
93+
com.amplifyframework.auth.cognito.exceptions.service.EnableSoftwareTokenMfaException(error)
9194
else -> UnknownException(fallbackMessage, error)
9295
}
9396
}

0 commit comments

Comments
 (0)