fix(auth): Add userAttributes to confirmSignIn call (#2640)

tylerjroach · web-flow · commit 3f2ea63725f0 · 2023-11-21T16:25:28.000-05:00
diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPlugin.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPlugin.kt
@@ -785,12 +785,15 @@ internal class RealAWSCognitoAuthPlugin(
             },
             {
                 val awsCognitoConfirmSignInOptions = options as? AWSCognitoAuthConfirmSignInOptions
+                val metadata = awsCognitoConfirmSignInOptions?.metadata ?: emptyMap()
+                val userAttributes = awsCognitoConfirmSignInOptions?.userAttributes ?: emptyList()
                 when (signInState) {
                     is SignInState.ResolvingChallenge -> {
                         val event = SignInChallengeEvent(
                             SignInChallengeEvent.EventType.VerifyChallengeAnswer(
                                 challengeResponse,
-                                awsCognitoConfirmSignInOptions?.metadata ?: mapOf()
+                                metadata,
+                                userAttributes
                             )
                         )
                         authStateMachine.send(event)
diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/actions/SignInChallengeCognitoActions.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/actions/SignInChallengeCognitoActions.kt
@@ -18,6 +18,7 @@ package com.amplifyframework.auth.cognito.actions
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChallengeNameType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ResourceNotFoundException
 import aws.sdk.kotlin.services.cognitoidentityprovider.respondToAuthChallenge
+import com.amplifyframework.auth.AuthUserAttribute
 import com.amplifyframework.auth.cognito.AuthEnvironment
 import com.amplifyframework.auth.cognito.helpers.AuthHelper
 import com.amplifyframework.auth.cognito.helpers.SignInChallengeHelper
@@ -32,9 +33,11 @@ import com.amplifyframework.statemachine.codegen.events.SignInChallengeEvent
 internal object SignInChallengeCognitoActions : SignInChallengeActions {
     private const val KEY_SECRET_HASH = "SECRET_HASH"
     private const val KEY_USERNAME = "USERNAME"
+    private const val KEY_PREFIX_USER_ATTRIBUTE = "userAttributes."
     override fun verifyChallengeAuthAction(
         answer: String,
         metadata: Map<String, String>,
+        attributes: List<AuthUserAttribute>,
         challenge: AuthChallenge
     ): Action = Action<AuthEnvironment>("VerifySignInChallenge") { id, dispatcher ->
         logger.verbose("$id Starting execution")
@@ -50,6 +53,12 @@ internal object SignInChallengeCognitoActions : SignInChallengeActions {
                 challengeResponses[responseKey] = answer
             }
 
+            challengeResponses.putAll(
+                attributes.map {
+                    Pair("${KEY_PREFIX_USER_ATTRIBUTE}${it.key.keyString}", it.value)
+                }
+            )
+
             val secretHash = AuthHelper.getSecretHash(
                 username,
                 configuration.userPool?.appClient,
@@ -90,6 +99,7 @@ internal object SignInChallengeCognitoActions : SignInChallengeActions {
                     SignInChallengeEvent.EventType.RetryVerifyChallengeAnswer(
                         answer,
                         metadata,
+                        attributes,
                         challenge
                     )
                 )
diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/actions/SignInChallengeActions.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/actions/SignInChallengeActions.kt
@@ -15,13 +15,15 @@
 
 package com.amplifyframework.statemachine.codegen.actions
 
+import com.amplifyframework.auth.AuthUserAttribute
 import com.amplifyframework.statemachine.Action
 import com.amplifyframework.statemachine.codegen.data.AuthChallenge
 
 internal interface SignInChallengeActions {
     fun verifyChallengeAuthAction(
         answer: String,
         metadata: Map<String, String>,
+        userAttributes: List<AuthUserAttribute>,
         challenge: AuthChallenge
     ): Action
 }
diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/events/SignInChallengeEvent.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/events/SignInChallengeEvent.kt
@@ -15,18 +15,24 @@
 
 package com.amplifyframework.statemachine.codegen.events
 
+import com.amplifyframework.auth.AuthUserAttribute
 import com.amplifyframework.statemachine.StateMachineEvent
 import com.amplifyframework.statemachine.codegen.data.AuthChallenge
 import java.util.Date
 
 internal class SignInChallengeEvent(val eventType: EventType, override val time: Date? = null) : StateMachineEvent {
     sealed class EventType {
         data class WaitForAnswer(val challenge: AuthChallenge, val hasNewResponse: Boolean = false) : EventType()
-        data class VerifyChallengeAnswer(val answer: String, val metadata: Map<String, String>) : EventType()
+        data class VerifyChallengeAnswer(
+            val answer: String,
+            val metadata: Map<String, String>,
+            val userAttributes: List<AuthUserAttribute>
+        ) : EventType()
 
         data class RetryVerifyChallengeAnswer(
             val answer: String,
             val metadata: Map<String, String>,
+            val userAttributes: List<AuthUserAttribute>,
             val authChallenge: AuthChallenge
         ) : EventType()
         data class FinalizeSignIn(val accessToken: String) : EventType()
diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/states/SignInChallengeState.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/states/SignInChallengeState.kt
@@ -60,7 +60,10 @@ internal sealed class SignInChallengeState : State {
                 is WaitingForAnswer -> when (challengeEvent) {
                     is SignInChallengeEvent.EventType.VerifyChallengeAnswer -> {
                         val action = challengeActions.verifyChallengeAuthAction(
-                            challengeEvent.answer, challengeEvent.metadata, oldState.challenge
+                            challengeEvent.answer,
+                            challengeEvent.metadata,
+                            challengeEvent.userAttributes,
+                            oldState.challenge
                         )
                         StateResolution(Verifying(oldState.challenge.challengeName), listOf(action))
                     }
@@ -78,7 +81,10 @@ internal sealed class SignInChallengeState : State {
                     }
                     is SignInChallengeEvent.EventType.RetryVerifyChallengeAnswer -> {
                         val action = challengeActions.verifyChallengeAuthAction(
-                            challengeEvent.answer, challengeEvent.metadata, challengeEvent.authChallenge
+                            challengeEvent.answer,
+                            challengeEvent.metadata,
+                            challengeEvent.userAttributes,
+                            challengeEvent.authChallenge,
                         )
                         StateResolution(Verifying(challengeEvent.authChallenge.challengeName), listOf(action))
                     }
@@ -92,7 +98,10 @@ internal sealed class SignInChallengeState : State {
                     when (challengeEvent) {
                         is SignInChallengeEvent.EventType.VerifyChallengeAnswer -> {
                             val action = challengeActions.verifyChallengeAuthAction(
-                                challengeEvent.answer, challengeEvent.metadata, oldState.challenge
+                                challengeEvent.answer,
+                                challengeEvent.metadata,
+                                challengeEvent.userAttributes,
+                                oldState.challenge,
                             )
                             StateResolution(Verifying(oldState.challenge.challengeName), listOf(action))
                         }
diff --git a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/StateTransitionTestBase.kt b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/StateTransitionTestBase.kt
@@ -327,6 +327,7 @@ open class StateTransitionTestBase {
 
         Mockito.`when`(
             mockSignInChallengeActions.verifyChallengeAuthAction(
+                MockitoHelper.anyObject(),
                 MockitoHelper.anyObject(),
                 MockitoHelper.anyObject(),
                 MockitoHelper.anyObject()
diff --git a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/StateTransitionTests.kt b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/StateTransitionTests.kt
@@ -387,12 +387,13 @@ class StateTransitionTests : StateTransitionTestBase() {
                         SignInChallengeEvent(
                             SignInChallengeEvent.EventType.RetryVerifyChallengeAnswer(
                                 "test",
-                                mapOf(),
+                                emptyMap(),
+                                emptyList(),
                                 AuthChallenge(
                                     ChallengeNameType.CustomChallenge.toString(),
                                     "Test",
                                     "session_mock_value",
-                                    mapOf()
+                                    emptyMap(),
                                 )
                             )
                         )
@@ -401,7 +402,8 @@ class StateTransitionTests : StateTransitionTestBase() {
                         SignInChallengeEvent(
                             SignInChallengeEvent.EventType.VerifyChallengeAnswer(
                                 "test",
-                                mapOf()
+                                emptyMap(),
+                                emptyList()
                             )
                         )
                     )
@@ -481,7 +483,7 @@ class StateTransitionTests : StateTransitionTestBase() {
                 challengeState?.apply {
                     stateMachine.send(
                         SignInChallengeEvent(
-                            SignInChallengeEvent.EventType.VerifyChallengeAnswer("test", mapOf())
+                            SignInChallengeEvent.EventType.VerifyChallengeAnswer("test", emptyMap(), emptyList())
                         )
                     )
                 }
diff --git a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/actions/SignInChallengeCognitoActionsTest.kt b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/actions/SignInChallengeCognitoActionsTest.kt
@@ -0,0 +1,141 @@
+/*
+ * 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.actions
+
+import androidx.test.core.app.ApplicationProvider
+import aws.sdk.kotlin.services.cognitoidentityprovider.CognitoIdentityProviderClient
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.RespondToAuthChallengeRequest
+import com.amplifyframework.auth.AuthUserAttribute
+import com.amplifyframework.auth.AuthUserAttributeKey
+import com.amplifyframework.auth.cognito.AWSCognitoAuthService
+import com.amplifyframework.auth.cognito.AuthEnvironment
+import com.amplifyframework.auth.cognito.StoreClientBehavior
+import com.amplifyframework.logging.Logger
+import com.amplifyframework.statemachine.EventDispatcher
+import com.amplifyframework.statemachine.StateMachineEvent
+import com.amplifyframework.statemachine.codegen.data.AmplifyCredential
+import com.amplifyframework.statemachine.codegen.data.AuthChallenge
+import com.amplifyframework.statemachine.codegen.data.AuthConfiguration
+import com.amplifyframework.statemachine.codegen.data.CredentialType
+import com.amplifyframework.statemachine.codegen.data.UserPoolConfiguration
+import io.mockk.coEvery
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.slot
+import junit.framework.TestCase.assertTrue
+import kotlin.test.assertEquals
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@RunWith(RobolectricTestRunner::class)
+class SignInChallengeCognitoActionsTest {
+
+    private val pool = mockk<UserPoolConfiguration> {
+        every { appClient } returns "client"
+        every { appClientSecret } returns null
+        every { pinpointAppId } returns null
+    }
+    private val configuration = mockk<AuthConfiguration> {
+        every { userPool } returns pool
+    }
+    private val cognitoAuthService = mockk<AWSCognitoAuthService>()
+    private val credentialStoreClient = mockk<StoreClientBehavior> {
+        coEvery { loadCredentials(CredentialType.ASF) } returns AmplifyCredential.ASFDevice("asf_id")
+    }
+    private val logger = mockk<Logger>(relaxed = true)
+    private val cognitoIdentityProviderClientMock = mockk<CognitoIdentityProviderClient>()
+
+    private val capturedEvent = slot<StateMachineEvent>()
+    private val dispatcher = mockk<EventDispatcher> {
+        every { send(capture(capturedEvent)) }.answers { }
+    }
+
+    private lateinit var authEnvironment: AuthEnvironment
+
+    @Before
+    fun setup() {
+        every { cognitoAuthService.cognitoIdentityProviderClient }.answers { cognitoIdentityProviderClientMock }
+        authEnvironment = AuthEnvironment(
+            ApplicationProvider.getApplicationContext(),
+            configuration,
+            cognitoAuthService,
+            credentialStoreClient,
+            null,
+            null,
+            logger
+        )
+    }
+
+    @Test
+    fun `very auth challenge without user attributes`() = runTest {
+        val expectedChallengeResponses = mapOf(
+            "USERNAME" to "testUser"
+        )
+        val capturedRequest = slot<RespondToAuthChallengeRequest>()
+        coEvery {
+            cognitoIdentityProviderClientMock.respondToAuthChallenge(capture(capturedRequest))
+        }.answers {
+            mockk()
+        }
+
+        SignInChallengeCognitoActions.verifyChallengeAuthAction(
+            "myAnswer",
+            emptyMap(),
+            emptyList(),
+            AuthChallenge(
+                "CONFIRM_SIGN_IN_WITH_NEW_PASSWORD",
+                username = "testUser",
+                session = null,
+                parameters = null
+            )
+        ).execute(dispatcher, authEnvironment)
+
+        assertTrue(capturedRequest.isCaptured)
+        assertEquals(expectedChallengeResponses, capturedRequest.captured.challengeResponses)
+    }
+
+    @Test
+    fun `user attributes are added to auth challenge`() = runTest {
+        val providedUserAttributes = listOf(AuthUserAttribute(AuthUserAttributeKey.phoneNumber(), "+15555555555"))
+        val expectedChallengeResponses = mapOf(
+            "USERNAME" to "testUser",
+            "userAttributes.phone_number" to "+15555555555"
+        )
+        val capturedRequest = slot<RespondToAuthChallengeRequest>()
+        coEvery {
+            cognitoIdentityProviderClientMock.respondToAuthChallenge(capture(capturedRequest))
+        }.answers {
+            mockk()
+        }
+
+        SignInChallengeCognitoActions.verifyChallengeAuthAction(
+            "myAnswer",
+            emptyMap(),
+            providedUserAttributes,
+            AuthChallenge(
+                "CONFIRM_SIGN_IN_WITH_NEW_PASSWORD",
+                username = "testUser",
+                session = null,
+                parameters = null
+            )
+        ).execute(dispatcher, authEnvironment)
+
+        assertTrue(capturedRequest.isCaptured)
+        assertEquals(expectedChallengeResponses, capturedRequest.captured.challengeResponses)
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -785,12 +785,15 @@ internal class RealAWSCognitoAuthPlugin(`
`785`	`785`	`},`
`786`	`786`	`{`
`787`	`787`	`val awsCognitoConfirmSignInOptions = options as? AWSCognitoAuthConfirmSignInOptions`
	`788`	`+ val metadata = awsCognitoConfirmSignInOptions?.metadata ?: emptyMap()`
	`789`	`+ val userAttributes = awsCognitoConfirmSignInOptions?.userAttributes ?: emptyList()`
`788`	`790`	`when (signInState) {`
`789`	`791`	`is SignInState.ResolvingChallenge -> {`
`790`	`792`	`val event = SignInChallengeEvent(`
`791`	`793`	`SignInChallengeEvent.EventType.VerifyChallengeAnswer(`
`792`	`794`	`challengeResponse,`
`793`		`- awsCognitoConfirmSignInOptions?.metadata ?: mapOf()`
	`795`	`+ metadata,`
	`796`	`+ userAttributes`
`794`	`797`	`)`
`795`	`798`	`)`
`796`	`799`	`authStateMachine.send(event)`
Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,10 @@ internal sealed class SignInChallengeState : State {`
`60`	`60`	`is WaitingForAnswer -> when (challengeEvent) {`
`61`	`61`	`is SignInChallengeEvent.EventType.VerifyChallengeAnswer -> {`
`62`	`62`	`val action = challengeActions.verifyChallengeAuthAction(`
`63`		`- challengeEvent.answer, challengeEvent.metadata, oldState.challenge`
	`63`	`+ challengeEvent.answer,`
	`64`	`+ challengeEvent.metadata,`
	`65`	`+ challengeEvent.userAttributes,`
	`66`	`+ oldState.challenge`
`64`	`67`	`)`
`65`	`68`	`StateResolution(Verifying(oldState.challenge.challengeName), listOf(action))`
`66`	`69`	`}`
`@@ -78,7 +81,10 @@ internal sealed class SignInChallengeState : State {`
`78`	`81`	`}`
`79`	`82`	`is SignInChallengeEvent.EventType.RetryVerifyChallengeAnswer -> {`
`80`	`83`	`val action = challengeActions.verifyChallengeAuthAction(`
`81`		`- challengeEvent.answer, challengeEvent.metadata, challengeEvent.authChallenge`
	`84`	`+ challengeEvent.answer,`
	`85`	`+ challengeEvent.metadata,`
	`86`	`+ challengeEvent.userAttributes,`
	`87`	`+ challengeEvent.authChallenge,`
`82`	`88`	`)`
`83`	`89`	`StateResolution(Verifying(challengeEvent.authChallenge.challengeName), listOf(action))`
`84`	`90`	`}`
`@@ -92,7 +98,10 @@ internal sealed class SignInChallengeState : State {`
`92`	`98`	`when (challengeEvent) {`
`93`	`99`	`is SignInChallengeEvent.EventType.VerifyChallengeAnswer -> {`
`94`	`100`	`val action = challengeActions.verifyChallengeAuthAction(`
`95`		`- challengeEvent.answer, challengeEvent.metadata, oldState.challenge`
	`101`	`+ challengeEvent.answer,`
	`102`	`+ challengeEvent.metadata,`
	`103`	`+ challengeEvent.userAttributes,`
	`104`	`+ oldState.challenge,`
`96`	`105`	`)`
`97`	`106`	`StateResolution(Verifying(oldState.challenge.challengeName), listOf(action))`
`98`	`107`	`}`
Original file line number	Diff line number	Diff line change
`@@ -387,12 +387,13 @@ class StateTransitionTests : StateTransitionTestBase() {`
`387`	`387`	`SignInChallengeEvent(`
`388`	`388`	`SignInChallengeEvent.EventType.RetryVerifyChallengeAnswer(`
`389`	`389`	`"test",`
`390`		`- mapOf(),`
	`390`	`+ emptyMap(),`
	`391`	`+ emptyList(),`
`391`	`392`	`AuthChallenge(`
`392`	`393`	`ChallengeNameType.CustomChallenge.toString(),`
`393`	`394`	`"Test",`
`394`	`395`	`"session_mock_value",`
`395`		`- mapOf()`
	`396`	`+ emptyMap(),`
`396`	`397`	`)`
`397`	`398`	`)`
`398`	`399`	`)`
`@@ -401,7 +402,8 @@ class StateTransitionTests : StateTransitionTestBase() {`
`401`	`402`	`SignInChallengeEvent(`
`402`	`403`	`SignInChallengeEvent.EventType.VerifyChallengeAnswer(`
`403`	`404`	`"test",`
`404`		`- mapOf()`
	`405`	`+ emptyMap(),`
	`406`	`+ emptyList()`
`405`	`407`	`)`
`406`	`408`	`)`
`407`	`409`	`)`
`@@ -481,7 +483,7 @@ class StateTransitionTests : StateTransitionTestBase() {`
`481`	`483`	`challengeState?.apply {`
`482`	`484`	`stateMachine.send(`
`483`	`485`	`SignInChallengeEvent(`
`484`		`- SignInChallengeEvent.EventType.VerifyChallengeAnswer("test", mapOf())`
	`486`	`+ SignInChallengeEvent.EventType.VerifyChallengeAnswer("test", emptyMap(), emptyList())`
`485`	`487`	`)`
`486`	`488`	`)`
`487`	`489`	`}`