Skip to content

Commit e31c581

Browse files
authored
chore(auth): Migrate TOTP setup to use case classes (#2992)
1 parent f51182e commit e31c581

File tree

10 files changed

+308
-228
lines changed

10 files changed

+308
-228
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -414,7 +414,7 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
414414
}
415415

416416
override fun setUpTOTP(onSuccess: Consumer<TOTPSetupDetails>, onError: Consumer<AuthException>) =
417-
enqueue(onSuccess, onError) { queueFacade.setUpTOTP() }
417+
enqueue(onSuccess, onError) { useCaseFactory.setupTotp().execute() }
418418

419419
override fun verifyTOTPSetup(code: String, onSuccess: Action, onError: Consumer<AuthException>) {
420420
verifyTOTPSetup(code, AWSCognitoAuthVerifyTOTPSetupOptions.CognitoBuilder().build(), onSuccess, onError)
@@ -425,7 +425,7 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
425425
options: AuthVerifyTOTPSetupOptions,
426426
onSuccess: Action,
427427
onError: Consumer<AuthException>
428-
) = enqueue(onSuccess, onError) { queueFacade.verifyTOTPSetup(code, options) }
428+
) = enqueue(onSuccess, onError) { useCaseFactory.verifyTotpSetup().execute(code, options) }
429429

430430
override fun associateWebAuthnCredential(
431431
callingActivity: Activity,

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

Lines changed: 0 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ import android.content.Intent
2020
import com.amplifyframework.auth.AuthCodeDeliveryDetails
2121
import com.amplifyframework.auth.AuthProvider
2222
import com.amplifyframework.auth.AuthSession
23-
import com.amplifyframework.auth.TOTPSetupDetails
2423
import com.amplifyframework.auth.cognito.options.FederateToIdentityPoolOptions
2524
import com.amplifyframework.auth.cognito.result.FederateToIdentityPoolResult
2625
import com.amplifyframework.auth.options.AuthConfirmResetPasswordOptions
@@ -32,7 +31,6 @@ import com.amplifyframework.auth.options.AuthResetPasswordOptions
3231
import com.amplifyframework.auth.options.AuthSignInOptions
3332
import com.amplifyframework.auth.options.AuthSignOutOptions
3433
import com.amplifyframework.auth.options.AuthSignUpOptions
35-
import com.amplifyframework.auth.options.AuthVerifyTOTPSetupOptions
3634
import com.amplifyframework.auth.options.AuthWebUISignInOptions
3735
import com.amplifyframework.auth.result.AuthResetPasswordResult
3836
import com.amplifyframework.auth.result.AuthSignInResult
@@ -286,21 +284,6 @@ internal class KotlinAuthFacadeInternal(private val delegate: RealAWSCognitoAuth
286284
)
287285
}
288286

289-
suspend fun setUpTOTP(): TOTPSetupDetails = suspendCoroutine { continuation ->
290-
delegate.setUpTOTP(
291-
{ continuation.resume(it) },
292-
{ continuation.resumeWithException(it) }
293-
)
294-
}
295-
suspend fun verifyTOTPSetup(code: String, options: AuthVerifyTOTPSetupOptions) = suspendCoroutine { continuation ->
296-
delegate.verifyTOTPSetup(
297-
code,
298-
options,
299-
{ continuation.resume(Unit) },
300-
{ continuation.resumeWithException(it) }
301-
)
302-
}
303-
304287
suspend fun fetchMFAPreference(): UserMFAPreference = suspendCoroutine { continuation ->
305288
delegate.fetchMFAPreference(
306289
{ continuation.resume(it) },

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

Lines changed: 0 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ package com.amplifyframework.auth.cognito
1818
import android.app.Activity
1919
import android.content.Intent
2020
import androidx.annotation.WorkerThread
21-
import aws.sdk.kotlin.services.cognitoidentityprovider.associateSoftwareToken
2221
import aws.sdk.kotlin.services.cognitoidentityprovider.confirmForgotPassword
2322
import aws.sdk.kotlin.services.cognitoidentityprovider.getUser
2423
import aws.sdk.kotlin.services.cognitoidentityprovider.model.AnalyticsMetadataType
@@ -27,10 +26,8 @@ import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChangePasswordReque
2726
import aws.sdk.kotlin.services.cognitoidentityprovider.model.EmailMfaSettingsType
2827
import aws.sdk.kotlin.services.cognitoidentityprovider.model.SmsMfaSettingsType
2928
import aws.sdk.kotlin.services.cognitoidentityprovider.model.SoftwareTokenMfaSettingsType
30-
import aws.sdk.kotlin.services.cognitoidentityprovider.model.VerifySoftwareTokenResponseType
3129
import aws.sdk.kotlin.services.cognitoidentityprovider.resendConfirmationCode
3230
import aws.sdk.kotlin.services.cognitoidentityprovider.setUserMfaPreference
33-
import aws.sdk.kotlin.services.cognitoidentityprovider.verifySoftwareToken
3431
import com.amplifyframework.AmplifyException
3532
import com.amplifyframework.annotations.InternalAmplifyApi
3633
import com.amplifyframework.auth.AWSCognitoAuthMetadataType
@@ -43,7 +40,6 @@ import com.amplifyframework.auth.AuthFactorType
4340
import com.amplifyframework.auth.AuthProvider
4441
import com.amplifyframework.auth.AuthSession
4542
import com.amplifyframework.auth.MFAType
46-
import com.amplifyframework.auth.TOTPSetupDetails
4743
import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidOauthConfigurationException
4844
import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidUserPoolConfigurationException
4945
import com.amplifyframework.auth.cognito.exceptions.invalidstate.SignedInException
@@ -53,7 +49,6 @@ import com.amplifyframework.auth.cognito.exceptions.service.InvalidParameterExce
5349
import com.amplifyframework.auth.cognito.exceptions.service.UserCancelledException
5450
import com.amplifyframework.auth.cognito.helpers.AuthHelper
5551
import com.amplifyframework.auth.cognito.helpers.HostedUIHelper
56-
import com.amplifyframework.auth.cognito.helpers.SessionHelper
5752
import com.amplifyframework.auth.cognito.helpers.SignInChallengeHelper
5853
import com.amplifyframework.auth.cognito.helpers.collectWhile
5954
import com.amplifyframework.auth.cognito.helpers.getAllowedMFATypesFromChallengeParameters
@@ -70,7 +65,6 @@ import com.amplifyframework.auth.cognito.options.AWSCognitoAuthResendSignUpCodeO
7065
import com.amplifyframework.auth.cognito.options.AWSCognitoAuthSignInOptions
7166
import com.amplifyframework.auth.cognito.options.AWSCognitoAuthSignOutOptions
7267
import com.amplifyframework.auth.cognito.options.AWSCognitoAuthSignUpOptions
73-
import com.amplifyframework.auth.cognito.options.AWSCognitoAuthVerifyTOTPSetupOptions
7468
import com.amplifyframework.auth.cognito.options.AWSCognitoAuthWebUISignInOptions
7569
import com.amplifyframework.auth.cognito.options.AuthFlowType
7670
import com.amplifyframework.auth.cognito.options.FederateToIdentityPoolOptions
@@ -96,7 +90,6 @@ import com.amplifyframework.auth.options.AuthResetPasswordOptions
9690
import com.amplifyframework.auth.options.AuthSignInOptions
9791
import com.amplifyframework.auth.options.AuthSignOutOptions
9892
import com.amplifyframework.auth.options.AuthSignUpOptions
99-
import com.amplifyframework.auth.options.AuthVerifyTOTPSetupOptions
10093
import com.amplifyframework.auth.options.AuthWebUISignInOptions
10194
import com.amplifyframework.auth.result.AuthResetPasswordResult
10295
import com.amplifyframework.auth.result.AuthSignInResult
@@ -1948,56 +1941,6 @@ internal class RealAWSCognitoAuthPlugin(
19481941
}
19491942
}
19501943

1951-
fun setUpTOTP(onSuccess: Consumer<TOTPSetupDetails>, onError: Consumer<AuthException>) {
1952-
authStateMachine.getCurrentState { authState ->
1953-
when (authState.authNState) {
1954-
is AuthenticationState.SignedIn -> {
1955-
GlobalScope.launch {
1956-
try {
1957-
val accessToken = getSession().userPoolTokensResult.value?.accessToken
1958-
accessToken?.let { token ->
1959-
SessionHelper.getUsername(token)?.let { username ->
1960-
authEnvironment.cognitoAuthService
1961-
.cognitoIdentityProviderClient?.associateSoftwareToken {
1962-
this.accessToken = token
1963-
}?.also { response ->
1964-
response.secretCode?.let { secret ->
1965-
onSuccess.accept(
1966-
TOTPSetupDetails(
1967-
secret,
1968-
username
1969-
)
1970-
)
1971-
}
1972-
}
1973-
}
1974-
} ?: onError.accept(SignedOutException())
1975-
} catch (error: Exception) {
1976-
onError.accept(
1977-
CognitoAuthExceptionConverter.lookup(
1978-
error,
1979-
"Cannot find a multi-factor authentication (MFA) method."
1980-
)
1981-
)
1982-
}
1983-
}
1984-
}
1985-
1986-
else -> onError.accept(InvalidStateException())
1987-
}
1988-
}
1989-
}
1990-
1991-
fun verifyTOTPSetup(
1992-
code: String,
1993-
options: AuthVerifyTOTPSetupOptions,
1994-
onSuccess: Action,
1995-
onError: Consumer<AuthException>
1996-
) {
1997-
val cognitoOptions = options as? AWSCognitoAuthVerifyTOTPSetupOptions
1998-
verifyTotp(code, cognitoOptions?.friendlyDeviceName, onSuccess, onError)
1999-
}
2000-
20011944
fun fetchMFAPreference(onSuccess: Consumer<UserMFAPreference>, onError: Consumer<AuthException>) {
20021945
authStateMachine.getCurrentState { authState ->
20031946
when (authState.authNState) {
@@ -2131,50 +2074,6 @@ internal class RealAWSCognitoAuthPlugin(
21312074
})
21322075
}
21332076

2134-
private fun verifyTotp(
2135-
code: String,
2136-
friendlyDeviceName: String?,
2137-
onSuccess: Action,
2138-
onError: Consumer<AuthException>
2139-
) {
2140-
authStateMachine.getCurrentState { authState ->
2141-
when (authState.authNState) {
2142-
is AuthenticationState.SignedIn -> {
2143-
GlobalScope.launch {
2144-
try {
2145-
val accessToken = getSession().userPoolTokensResult.value?.accessToken
2146-
accessToken?.let { token ->
2147-
authEnvironment.cognitoAuthService
2148-
.cognitoIdentityProviderClient?.verifySoftwareToken {
2149-
this.userCode = code
2150-
this.friendlyDeviceName = friendlyDeviceName
2151-
this.accessToken = token
2152-
}?.also {
2153-
when (it.status) {
2154-
is VerifySoftwareTokenResponseType.Success -> onSuccess.call()
2155-
else -> throw ServiceException(
2156-
message = "An unknown service error has occurred",
2157-
recoverySuggestion = AmplifyException.TODO_RECOVERY_SUGGESTION
2158-
)
2159-
}
2160-
}
2161-
} ?: onError.accept(SignedOutException())
2162-
} catch (error: Exception) {
2163-
onError.accept(
2164-
CognitoAuthExceptionConverter.lookup(
2165-
error,
2166-
"Amazon Cognito cannot find a multi-factor authentication (MFA) method."
2167-
)
2168-
)
2169-
}
2170-
}
2171-
}
2172-
2173-
else -> onError.accept(InvalidStateException())
2174-
}
2175-
}
2176-
}
2177-
21782077
private fun _clearFederationToIdentityPool(onSuccess: Action, onError: Consumer<AuthException>) {
21792078
_signOut(sendHubEvent = false) {
21802079
when (it) {

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,4 +96,16 @@ internal class AuthUseCaseFactory(
9696
fetchAuthSession = fetchAuthSession(),
9797
stateMachine = stateMachine
9898
)
99+
100+
fun setupTotp() = SetupTotpUseCase(
101+
client = authEnvironment.requireIdentityProviderClient(),
102+
fetchAuthSession = fetchAuthSession(),
103+
stateMachine = stateMachine
104+
)
105+
106+
fun verifyTotpSetup() = VerifyTotpSetupUseCase(
107+
client = authEnvironment.requireIdentityProviderClient(),
108+
fetchAuthSession = fetchAuthSession(),
109+
stateMachine = stateMachine
110+
)
99111
}
Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/*
2+
* Copyright 2025 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+
16+
package com.amplifyframework.auth.cognito.usecases
17+
18+
import aws.sdk.kotlin.services.cognitoidentityprovider.CognitoIdentityProviderClient
19+
import aws.sdk.kotlin.services.cognitoidentityprovider.associateSoftwareToken
20+
import com.amplifyframework.AmplifyException
21+
import com.amplifyframework.auth.AuthException
22+
import com.amplifyframework.auth.TOTPSetupDetails
23+
import com.amplifyframework.auth.cognito.AuthStateMachine
24+
import com.amplifyframework.auth.cognito.requireAccessToken
25+
import com.amplifyframework.auth.cognito.requireSignedInState
26+
27+
internal class SetupTotpUseCase(
28+
private val fetchAuthSession: FetchAuthSessionUseCase,
29+
private val client: CognitoIdentityProviderClient,
30+
private val stateMachine: AuthStateMachine
31+
) {
32+
33+
suspend fun execute(): TOTPSetupDetails {
34+
val state = stateMachine.requireSignedInState()
35+
36+
val token = fetchAuthSession.execute().requireAccessToken()
37+
38+
val response = client.associateSoftwareToken {
39+
accessToken = token
40+
}
41+
42+
val sharedSecret = response.secretCode ?: throw AuthException(
43+
"Shared secret missing from response",
44+
AmplifyException.REPORT_BUG_TO_AWS_SUGGESTION
45+
)
46+
47+
return TOTPSetupDetails(
48+
sharedSecret = sharedSecret,
49+
username = state.signedInData.username
50+
)
51+
}
52+
}
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
* Copyright 2025 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+
16+
package com.amplifyframework.auth.cognito.usecases
17+
18+
import aws.sdk.kotlin.services.cognitoidentityprovider.CognitoIdentityProviderClient
19+
import aws.sdk.kotlin.services.cognitoidentityprovider.model.VerifySoftwareTokenResponseType
20+
import aws.sdk.kotlin.services.cognitoidentityprovider.verifySoftwareToken
21+
import com.amplifyframework.AmplifyException
22+
import com.amplifyframework.auth.cognito.AuthStateMachine
23+
import com.amplifyframework.auth.cognito.options.AWSCognitoAuthVerifyTOTPSetupOptions
24+
import com.amplifyframework.auth.cognito.requireAccessToken
25+
import com.amplifyframework.auth.cognito.requireSignedInState
26+
import com.amplifyframework.auth.exceptions.ServiceException
27+
import com.amplifyframework.auth.options.AuthVerifyTOTPSetupOptions
28+
29+
internal class VerifyTotpSetupUseCase(
30+
private val fetchAuthSession: FetchAuthSessionUseCase,
31+
private val client: CognitoIdentityProviderClient,
32+
private val stateMachine: AuthStateMachine
33+
) {
34+
35+
suspend fun execute(code: String, options: AuthVerifyTOTPSetupOptions) {
36+
val cognitoOptions = options as? AWSCognitoAuthVerifyTOTPSetupOptions
37+
val deviceName = cognitoOptions?.friendlyDeviceName
38+
39+
stateMachine.requireSignedInState()
40+
41+
val token = fetchAuthSession.execute().requireAccessToken()
42+
43+
val response = client.verifySoftwareToken {
44+
userCode = code
45+
friendlyDeviceName = deviceName
46+
accessToken = token
47+
}
48+
49+
if (response.status != VerifySoftwareTokenResponseType.Success) {
50+
throw ServiceException(
51+
message = "An unknown service error has occurred",
52+
recoverySuggestion = AmplifyException.TODO_RECOVERY_SUGGESTION
53+
)
54+
}
55+
}
56+
}

aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/AWSCognitoAuthPluginTest.kt

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -692,17 +692,23 @@ class AWSCognitoAuthPluginTest {
692692
fun setUpTOTP() {
693693
val expectedOnSuccess = Consumer<TOTPSetupDetails> { }
694694
val expectedOnError = Consumer<AuthException> { }
695+
696+
val useCase = authPlugin.useCaseFactory.setupTotp()
697+
695698
authPlugin.setUpTOTP(expectedOnSuccess, expectedOnError)
696-
verify(timeout = CHANNEL_TIMEOUT) { realPlugin.setUpTOTP(any(), any()) }
699+
coVerify(timeout = CHANNEL_TIMEOUT) { useCase.execute() }
697700
}
698701

699702
@Test
700703
fun verifyTOTPSetup() {
701704
val code = "123456"
702705
val expectedOnSuccess = Action { }
703706
val expectedOnError = Consumer<AuthException> { }
707+
708+
val useCase = authPlugin.useCaseFactory.verifyTotpSetup()
709+
704710
authPlugin.verifyTOTPSetup(code, expectedOnSuccess, expectedOnError)
705-
verify(timeout = CHANNEL_TIMEOUT) { realPlugin.verifyTOTPSetup(code, any(), any(), any()) }
711+
coVerify(timeout = CHANNEL_TIMEOUT) { useCase.execute(code, any()) }
706712
}
707713

708714
@Test
@@ -711,8 +717,11 @@ class AWSCognitoAuthPluginTest {
711717
val options = AWSCognitoAuthVerifyTOTPSetupOptions.CognitoBuilder().friendlyDeviceName("DEVICE_NAME").build()
712718
val expectedOnSuccess = Action { }
713719
val expectedOnError = Consumer<AuthException> { }
720+
721+
val useCase = authPlugin.useCaseFactory.verifyTotpSetup()
722+
714723
authPlugin.verifyTOTPSetup(code, options, expectedOnSuccess, expectedOnError)
715-
verify(timeout = CHANNEL_TIMEOUT) { realPlugin.verifyTOTPSetup(code, options, any(), any()) }
724+
coVerify(timeout = CHANNEL_TIMEOUT) { useCase.execute(code, options) }
716725
}
717726

718727
@Test

0 commit comments

Comments
 (0)