fix(auth): Fix issue with extracting username during USER_PASSWORD_AUTH

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

@@ -22,6 +22,7 @@ import aws.sdk.kotlin.services.cognitoidentityprovider.endpoints.CognitoIdentity
 import aws.smithy.kotlin.runtime.client.RequestInterceptorContext
 import aws.smithy.kotlin.runtime.client.endpoints.Endpoint
 import aws.smithy.kotlin.runtime.http.interceptors.HttpInterceptor
+import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidUserPoolConfigurationException
 import com.amplifyframework.util.setHttpEngine
 
 interface AWSCognitoAuthService {
@@ -75,3 +76,6 @@ interface AWSCognitoAuthService {
         }
     }
 }
+
+internal fun AWSCognitoAuthService.requireIdentityClient() =
+    cognitoIdentityProviderClient ?: throw InvalidUserPoolConfigurationException()

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

@@ -18,14 +18,12 @@ package com.amplifyframework.auth.cognito.actions
 import aws.sdk.kotlin.services.cognitoidentityprovider.initiateAuth
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChallengeNameType
 import aws.sdk.kotlin.services.cognitoidentityprovider.respondToAuthChallenge
-import aws.smithy.kotlin.runtime.util.type
-import com.amplifyframework.AmplifyException
 import com.amplifyframework.auth.cognito.AuthEnvironment
 import com.amplifyframework.auth.cognito.helpers.AuthHelper
 import com.amplifyframework.auth.cognito.helpers.SignInChallengeHelper
 import com.amplifyframework.auth.cognito.helpers.toCognitoType
 import com.amplifyframework.auth.cognito.options.AuthFlowType
-import com.amplifyframework.auth.exceptions.ServiceException
+import com.amplifyframework.auth.cognito.requireIdentityClient
 import com.amplifyframework.statemachine.Action
 import com.amplifyframework.statemachine.codegen.actions.MigrateAuthActions
 import com.amplifyframework.statemachine.codegen.data.SignInMethod
@@ -50,6 +48,8 @@ internal object MigrateAuthCognitoActions : MigrateAuthActions {
                     configuration.userPool?.appClientSecret
                 )
 
+                val client = cognitoAuthService.requireIdentityClient()
+
                 val authParams = mutableMapOf(KEY_USERNAME to event.username, KEY_PASSWORD to event.password)
                 secretHash?.let { authParams[KEY_SECRET_HASH] = it }
 
@@ -59,7 +59,7 @@ internal object MigrateAuthCognitoActions : MigrateAuthActions {
                 if (event.respondToAuthChallenge?.session != null) {
                     authParams[KEY_ANSWER] = ChallengeNameType.Password.value
 
-                    val response = cognitoAuthService.cognitoIdentityProviderClient?.respondToAuthChallenge {
+                    val response = client.respondToAuthChallenge {
                         clientId = configuration.userPool?.appClient
                         challengeName = ChallengeNameType.SelectChallenge
                         this.challengeResponses = authParams
@@ -69,21 +69,20 @@ internal object MigrateAuthCognitoActions : MigrateAuthActions {
                         encodedContextData?.let { this.userContextData { encodedData = it } }
                     }
 
-                    response?.let {
-                        SignInChallengeHelper.evaluateNextStep(
-                            username = event.username,
-                            challengeNameType = response.challengeName,
-                            session = response.session,
-                            challengeParameters = response.challengeParameters,
-                            authenticationResult = response.authenticationResult,
-                            signInMethod = SignInMethod.ApiBased(SignInMethod.ApiBased.AuthType.USER_AUTH)
-                        )
-                    } ?: throw ServiceException("Sign in failed", AmplifyException.TODO_RECOVERY_SUGGESTION)
+                    val username = AuthHelper.getActiveUsername(event.username, response)
+                    SignInChallengeHelper.evaluateNextStep(
+                        username = username,
+                        challengeNameType = response.challengeName,
+                        session = response.session,
+                        challengeParameters = response.challengeParameters,
+                        authenticationResult = response.authenticationResult,
+                        signInMethod = SignInMethod.ApiBased(SignInMethod.ApiBased.AuthType.USER_AUTH)
+                    )
                 } else {
                     if (event.authFlowType == AuthFlowType.USER_AUTH) {
                         authParams[KEY_PREFERRED_CHALLENGE] = KEY_PASSWORD
                     }
-                    val response = cognitoAuthService.cognitoIdentityProviderClient?.initiateAuth {
+                    val response = client.initiateAuth {
                         authFlow = event.authFlowType.toCognitoType()
                         clientId = configuration.userPool?.appClient
                         authParameters = authParams
@@ -92,28 +91,20 @@ internal object MigrateAuthCognitoActions : MigrateAuthActions {
                         encodedContextData?.let { userContextData { encodedData = it } }
                     }
 
-                    response?.let {
-                        val username = AuthHelper.getActiveUsername(
-                            username = event.username,
-                            alternateUsername = response.challengeParameters?.get(KEY_USERNAME),
-                            userIDForSRP = response.challengeParameters?.get(
-                                KEY_USERID_FOR_SRP
-                            )
-                        )
-                        val signInMethod = if (event.authFlowType == AuthFlowType.USER_AUTH) {
-                            SignInMethod.ApiBased(SignInMethod.ApiBased.AuthType.USER_AUTH)
-                        } else {
-                            SignInMethod.ApiBased(SignInMethod.ApiBased.AuthType.USER_SRP_AUTH)
-                        }
-                        SignInChallengeHelper.evaluateNextStep(
-                            username = username,
-                            challengeNameType = response.challengeName,
-                            session = response.session,
-                            challengeParameters = response.challengeParameters,
-                            authenticationResult = response.authenticationResult,
-                            signInMethod = signInMethod
-                        )
-                    } ?: throw ServiceException("Sign in failed", AmplifyException.TODO_RECOVERY_SUGGESTION)
+                    val username = AuthHelper.getActiveUsername(event.username, response)
+                    val signInMethod = if (event.authFlowType == AuthFlowType.USER_AUTH) {
+                        SignInMethod.ApiBased(SignInMethod.ApiBased.AuthType.USER_AUTH)
+                    } else {
+                        SignInMethod.ApiBased(SignInMethod.ApiBased.AuthType.USER_SRP_AUTH)
+                    }
+                    SignInChallengeHelper.evaluateNextStep(
+                        username = username,
+                        challengeNameType = response.challengeName,
+                        session = response.session,
+                        challengeParameters = response.challengeParameters,
+                        authenticationResult = response.authenticationResult,
+                        signInMethod = signInMethod
+                    )
                 }
             } catch (e: Exception) {
                 val errorEvent = SignInEvent(SignInEvent.EventType.ThrowError(e))

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

@@ -26,6 +26,7 @@ import com.amplifyframework.auth.cognito.helpers.SRPHelper
 import com.amplifyframework.auth.cognito.helpers.SignInChallengeHelper
 import com.amplifyframework.auth.cognito.helpers.toCognitoType
 import com.amplifyframework.auth.cognito.options.AuthFlowType
+import com.amplifyframework.auth.cognito.requireIdentityClient
 import com.amplifyframework.auth.exceptions.ServiceException
 import com.amplifyframework.auth.exceptions.UnknownException
 import com.amplifyframework.statemachine.Action
@@ -62,6 +63,8 @@ internal object SRPCognitoActions : SRPActions {
             val evt = try {
                 srpHelper = SRPHelper(event.password)
 
+                val client = cognitoAuthService.requireIdentityClient()
+
                 val secretHash = AuthHelper.getSecretHash(
                     event.username,
                     configuration.userPool?.appClient,
@@ -79,7 +82,7 @@ internal object SRPCognitoActions : SRPActions {
                 if (event.respondToAuthChallenge?.session != null) {
                     authParams[KEY_ANSWER] = ChallengeNameType.PasswordSrp.value
 
-                    val response = cognitoAuthService.cognitoIdentityProviderClient?.respondToAuthChallenge {
+                    val response = client.respondToAuthChallenge {
                         clientId = configuration.userPool?.appClient
                         challengeName = ChallengeNameType.SelectChallenge
                         this.challengeResponses = authParams
@@ -90,49 +93,36 @@ internal object SRPCognitoActions : SRPActions {
                     }
 
                     val updatedDeviceMetadata = getDeviceMetadata(
-                        AuthHelper.getActiveUsername(
-                            username = event.username,
-                            alternateUsername = response?.challengeParameters?.get(KEY_USERNAME),
-                            userIDForSRP = response?.challengeParameters?.get(
-                                KEY_USERID_FOR_SRP
-                            )
-                        )
+                        AuthHelper.getActiveUsername(event.username, response)
                     )
 
                     parseResponseChallenge(
-                        challengeNameType = response?.challengeName,
-                        challengeParams = response?.challengeParameters,
-                        session = response?.session,
+                        challengeNameType = response.challengeName,
+                        challengeParams = response.challengeParameters,
+                        session = response.session,
                         updatedDeviceMetadata = updatedDeviceMetadata,
                         metadata = event.metadata
                     )
                 } else {
                     if (event.authFlowType == AuthFlowType.USER_AUTH) {
                         authParams[KEY_PREFERRED_CHALLENGE] = KEY_PASSWORD_SRP
                     }
-                    val initiateAuthResponse = cognitoAuthService.cognitoIdentityProviderClient?.initiateAuth {
+                    val response = client.initiateAuth {
                         authFlow = event.authFlowType.toCognitoType()
                         clientId = configuration.userPool?.appClient
                         authParameters = authParams
                         clientMetadata = event.metadata
                         pinpointEndpointId?.let { analyticsMetadata { analyticsEndpointId = it } }
                         encodedContextData?.let { userContextData { encodedData = it } }
                     }
-
                     val updatedDeviceMetadata = getDeviceMetadata(
-                        AuthHelper.getActiveUsername(
-                            username = event.username,
-                            alternateUsername = initiateAuthResponse?.challengeParameters?.get(KEY_USERNAME),
-                            userIDForSRP = initiateAuthResponse?.challengeParameters?.get(
-                                KEY_USERID_FOR_SRP
-                            )
-                        )
+                        AuthHelper.getActiveUsername(event.username, response)
                     )
 
                     parseResponseChallenge(
-                        challengeNameType = initiateAuthResponse?.challengeName,
-                        challengeParams = initiateAuthResponse?.challengeParameters,
-                        session = initiateAuthResponse?.session,
+                        challengeNameType = response.challengeName,
+                        challengeParams = response.challengeParameters,
+                        session = response.session,
                         updatedDeviceMetadata = updatedDeviceMetadata,
                         metadata = event.metadata
                     )

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

@@ -64,13 +64,7 @@ internal object SignInCustomCognitoActions : CustomSignInActions {
                 if (initiateAuthResponse?.challengeName == ChallengeNameType.CustomChallenge &&
                     initiateAuthResponse.challengeParameters != null
                 ) {
-                    val activeUserName = AuthHelper.getActiveUsername(
-                        username = event.username,
-                        alternateUsername = initiateAuthResponse.challengeParameters?.get(KEY_USERNAME),
-                        userIDForSRP = initiateAuthResponse.challengeParameters?.get(
-                            KEY_USERID_FOR_SRP
-                        )
-                    )
+                    val activeUserName = AuthHelper.getActiveUsername(event.username, initiateAuthResponse)
                     SignInChallengeHelper.evaluateNextStep(
                         username = activeUserName,
                         challengeNameType = initiateAuthResponse.challengeName,

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

@@ -72,13 +72,7 @@ internal object UserAuthSignInCognitoActions : UserAuthSignInActions {
                     resolvedSession != null &&
                     resolvedChallenges != null
                 ) {
-                    val activeUserName = AuthHelper.getActiveUsername(
-                        username = event.username,
-                        alternateUsername = initiateAuthResponse.challengeParameters?.get(KEY_USERNAME),
-                        userIDForSRP = initiateAuthResponse.challengeParameters?.get(
-                            KEY_USERID_FOR_SRP
-                        )
-                    )
+                    val activeUserName = AuthHelper.getActiveUsername(event.username, initiateAuthResponse)
 
                     val listOfChallenges = resolvedChallenges.map { it.value }
 
@@ -95,13 +89,7 @@ internal object UserAuthSignInCognitoActions : UserAuthSignInActions {
                     initiateAuthResponse?.challengeParameters != null &&
                     resolvedSession != null
                 ) {
-                    val activeUserName = AuthHelper.getActiveUsername(
-                        username = event.username,
-                        alternateUsername = initiateAuthResponse.challengeParameters?.get(KEY_USERNAME),
-                        userIDForSRP = initiateAuthResponse.challengeParameters?.get(
-                            KEY_USERID_FOR_SRP
-                        )
-                    )
+                    val activeUserName = AuthHelper.getActiveUsername(event.username, initiateAuthResponse)
 
                     SignInChallengeHelper.evaluateNextStep(
                         username = activeUserName,

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

@@ -15,11 +15,19 @@
 
 package com.amplifyframework.auth.cognito.helpers
 
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.AuthenticationResultType
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.InitiateAuthResponse
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.RespondToAuthChallengeResponse
 import com.amplifyframework.auth.cognito.exceptions.service.InvalidParameterException
 import com.amplifyframework.auth.exceptions.UnknownException
+import com.amplifyframework.statemachine.codegen.data.ChallengeParameter
+import com.amplifyframework.statemachine.codegen.data.asAccessToken
+import com.amplifyframework.statemachine.codegen.data.asIdToken
 import javax.crypto.Mac
 import javax.crypto.spec.SecretKeySpec
 
+private typealias ChallengeParameters = Map<String, String>
+
 internal open class AuthHelper {
 
     companion object {
@@ -55,13 +63,30 @@ internal open class AuthHelper {
         }
 
         /**
-         * This method is used to capture the activeUsername we should be using to store and retreive device Metadata.
+         * This method is used to capture the activeUsername we should be using to store and retrieve device Metadata.
          * This becomes more important when an alias is used to signin instead of a username where cognito then
          * generates the username on the customers behalf and returns it. We need to then use that username to store
          * and retrieve device information as the username that the customer entered will no longer be available when
          * respondToAuthChallenge/confirmSignIn calls are made.
          * */
-        fun getActiveUsername(username: String, alternateUsername: String?, userIDForSRP: String?): String =
-            alternateUsername ?: (userIDForSRP ?: username)
+        fun getActiveUsername(usernameSentInRequest: String, response: InitiateAuthResponse): String =
+            getActiveUsername(usernameSentInRequest, response.authenticationResult, response.challengeParameters)
+        fun getActiveUsername(usernameSentInRequest: String, response: RespondToAuthChallengeResponse): String =
+            getActiveUsername(usernameSentInRequest, response.authenticationResult, response.challengeParameters)
+
+        private fun getActiveUsername(
+            usernameSentInRequest: String,
+            authenticationResult: AuthenticationResultType?,
+            challengeParameters: ChallengeParameters?
+        ): String = authenticationResult?.getUsername() ?: challengeParameters?.getUsername() ?: usernameSentInRequest
+
+        private fun AuthenticationResultType.getUsername(): String? {
+            val idToken = idToken?.asIdToken()
+            val accessToken = accessToken?.asAccessToken()
+            return idToken?.username ?: accessToken?.username
+        }
+
+        private fun ChallengeParameters.getUsername(): String? =
+            get(ChallengeParameter.Username.key) ?: get(ChallengeParameter.UserIdForSrp.key)
     }
 }

aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/data/ChallengeParameter.kt

@@ -23,5 +23,7 @@ internal enum class ChallengeParameter(val key: String) {
     CodeDeliveryMedium("CODE_DELIVERY_DELIVERY_MEDIUM"),
     CredentialRequestOptions("CREDENTIAL_REQUEST_OPTIONS"),
     MfasCanChoose("MFAS_CAN_CHOOSE"),
-    MfasCanSetup("MFAS_CAN_SETUP")
+    MfasCanSetup("MFAS_CAN_SETUP"),
+    Username("USERNAME"),
+    UserIdForSrp("USER_ID_FOR_SRP")
 }

aws-auth-cognito/src/main/java/com/amplifyframework/statemachine/codegen/data/Tokens.kt

@@ -65,14 +65,19 @@ internal abstract class Jwt {
     enum class Claim(val key: String) {
         Expiration("exp"),
         UserSub("sub"),
-        Username("username"),
-        TokenRevocationId("origin_jti")
+        Username("username"), // username claim in the access token
+        TokenRevocationId("origin_jti"),
+        CognitoUsername("cognito:username") // username claim in the identity token
     }
 }
 
 // See https://docs.aws.amazon.com/cognito/latest/developerguide/amazon-cognito-user-pools-using-the-id-token.html
 @Serializable
 internal class IdToken(override val tokenValue: String) : Jwt() {
+    val userSub: String?
+        get() = getClaim(Claim.UserSub)
+    val username: String?
+        get() = getClaim(Claim.CognitoUsername)
     val expiration: Instant? by lazy {
         getClaim(Claim.Expiration)?.let { Instant.ofEpochSecond(it.toLong()) }
     }