chore(auth): Add typesafe classes for handling auth tokens

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

@@ -99,7 +99,7 @@ internal fun AmplifyCredential.getCognitoSession(
         }
 
         return try {
-            AuthSessionResult.success(userPoolTokens?.accessToken?.let(SessionHelper::getUserSub))
+            AuthSessionResult.success(userPoolTokens?.accessToken?.userSub)
         } catch (e: Exception) {
             AuthSessionResult.failure(UnknownException(cause = e))
         }
@@ -115,9 +115,9 @@ internal fun AmplifyCredential.getCognitoSession(
 
         return AuthSessionResult.success(
             AWSCognitoUserPoolTokens(
-                accessToken = cognitoUserPoolTokens.accessToken,
-                idToken = cognitoUserPoolTokens.idToken,
-                refreshToken = cognitoUserPoolTokens.refreshToken
+                accessToken = cognitoUserPoolTokens.accessToken?.tokenValue,
+                idToken = cognitoUserPoolTokens.idToken?.tokenValue,
+                refreshToken = cognitoUserPoolTokens.refreshToken?.tokenValue
             )
         )
     }

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

@@ -20,7 +20,6 @@ import aws.sdk.kotlin.services.cognitoidentity.model.GetIdRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.getTokensFromRefreshToken
 import aws.smithy.kotlin.runtime.time.Instant
 import com.amplifyframework.auth.cognito.AuthEnvironment
-import com.amplifyframework.auth.cognito.helpers.SessionHelper
 import com.amplifyframework.auth.exceptions.NotAuthorizedException
 import com.amplifyframework.auth.exceptions.SessionExpiredException
 import com.amplifyframework.auth.exceptions.SignedOutException
@@ -48,7 +47,7 @@ internal object FetchAuthSessionCognitoActions : FetchAuthSessionActions {
                 val deviceMetadata: DeviceMetadata.Metadata? = getDeviceMetadata(username)
 
                 val response = cognitoAuthService.cognitoIdentityProviderClient?.getTokensFromRefreshToken {
-                    refreshToken = tokens.refreshToken
+                    refreshToken = tokens.refreshToken?.tokenValue
                     clientId = configuration.userPool?.appClient
                     clientSecret = configuration.userPool?.appClientSecret
                     deviceKey = deviceMetadata?.deviceKey
@@ -58,13 +57,13 @@ internal object FetchAuthSessionCognitoActions : FetchAuthSessionActions {
                 val refreshedUserPoolTokens = CognitoUserPoolTokens(
                     idToken = response?.authenticationResult?.idToken,
                     accessToken = response?.authenticationResult?.accessToken,
-                    refreshToken = response?.authenticationResult?.refreshToken ?: tokens.refreshToken,
+                    refreshToken = response?.authenticationResult?.refreshToken ?: tokens.refreshToken?.tokenValue,
                     expiration = Instant.now().plus(expiresIn.seconds).epochSeconds
                 )
 
                 val updatedSignedInData = signedInData.copy(
-                    userId = refreshedUserPoolTokens.accessToken?.let(SessionHelper::getUserSub) ?: signedInData.userId,
-                    username = refreshedUserPoolTokens.accessToken?.let(SessionHelper::getUsername) ?: username,
+                    userId = refreshedUserPoolTokens.accessToken?.userSub ?: signedInData.userId,
+                    username = refreshedUserPoolTokens.accessToken?.username ?: username,
                     cognitoUserPoolTokens = refreshedUserPoolTokens
                 )
 

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

@@ -17,7 +17,6 @@ package com.amplifyframework.auth.cognito.actions
 
 import com.amplifyframework.auth.cognito.AuthEnvironment
 import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidOauthConfigurationException
-import com.amplifyframework.auth.cognito.helpers.JWTParser
 import com.amplifyframework.statemachine.Action
 import com.amplifyframework.statemachine.codegen.actions.HostedUIActions
 import com.amplifyframework.statemachine.codegen.data.DeviceMetadata
@@ -54,8 +53,8 @@ internal object HostedUICognitoActions : HostedUIActions {
                 if (hostedUIClient == null) throw InvalidOauthConfigurationException()
 
                 val token = hostedUIClient.fetchToken(event.uri)
-                val userId = token.accessToken?.let { JWTParser.getClaim(it, "sub") } ?: ""
-                val username = token.accessToken?.let { JWTParser.getClaim(it, "username") } ?: ""
+                val userId = token.accessToken?.userSub ?: ""
+                val username = token.accessToken?.username ?: ""
 
                 val signedInData = SignedInData(
                     userId,

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

@@ -126,7 +126,7 @@ internal object SignInCognitoActions : SignInActions {
 
                 cognitoAuthService.cognitoIdentityProviderClient?.confirmDevice(
                     ConfirmDeviceRequest.invoke {
-                        this.accessToken = event.signedInData.cognitoUserPoolTokens.accessToken
+                        this.accessToken = event.signedInData.cognitoUserPoolTokens.accessToken?.tokenValue
                         this.deviceKey = deviceKey
                         this.deviceName = Build.MODEL
                         this.deviceSecretVerifierConfig = DeviceSecretVerifierConfigType.invoke {

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

@@ -19,7 +19,6 @@ import aws.sdk.kotlin.services.cognitoidentityprovider.model.GlobalSignOutReques
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.RevokeTokenRequest
 import com.amplifyframework.auth.cognito.AuthEnvironment
 import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidOauthConfigurationException
-import com.amplifyframework.auth.cognito.helpers.JWTParser
 import com.amplifyframework.statemachine.Action
 import com.amplifyframework.statemachine.codegen.actions.SignOutActions
 import com.amplifyframework.statemachine.codegen.data.DeviceMetadata
@@ -76,15 +75,15 @@ internal object SignOutCognitoActions : SignOutActions {
             val accessToken = event.signedInData.cognitoUserPoolTokens.accessToken
             val evt = try {
                 cognitoAuthService.cognitoIdentityProviderClient?.globalSignOut(
-                    GlobalSignOutRequest { this.accessToken = accessToken }
+                    GlobalSignOutRequest { this.accessToken = accessToken?.tokenValue }
                 )
                 SignOutEvent(
                     SignOutEvent.EventType.RevokeToken(event.signedInData, hostedUIErrorData = event.hostedUIErrorData)
                 )
             } catch (e: Exception) {
                 logger.warn("Failed to sign out globally.", e)
                 val globalSignOutErrorData = GlobalSignOutErrorData(
-                    accessToken = accessToken,
+                    accessToken = accessToken?.tokenValue,
                     error = e
                 )
                 SignOutEvent(
@@ -106,19 +105,19 @@ internal object SignOutCognitoActions : SignOutActions {
             val refreshToken = event.signedInData.cognitoUserPoolTokens.refreshToken
             val evt = try {
                 // Check for "origin_jti" claim in access token, else skip revoking
-                if (accessToken?.let { JWTParser.hasClaim(it, "origin_jti") } == true) {
+                if (accessToken?.tokenRevocationId != null) {
                     cognitoAuthService.cognitoIdentityProviderClient?.revokeToken(
                         RevokeTokenRequest {
                             clientId = configuration.userPool?.appClient
                             clientSecret = configuration.userPool?.appClientSecret
-                            token = refreshToken
+                            token = refreshToken?.tokenValue
                         }
                     )
                     SignOutEvent(SignOutEvent.EventType.SignOutLocally(event.signedInData, event.hostedUIErrorData))
                 } else {
                     logger.debug("Access Token does not contain `origin_jti` claim. Skip revoking tokens.")
                     val error = RevokeTokenErrorData(
-                        refreshToken = refreshToken,
+                        refreshToken = refreshToken?.tokenValue,
                         error = Exception("Access Token does not contain `origin_jti` claim. Skip revoking tokens.")
                     )
 
@@ -134,7 +133,7 @@ internal object SignOutCognitoActions : SignOutActions {
             } catch (e: Exception) {
                 logger.warn("Failed to revoke tokens.", e)
                 val error = RevokeTokenErrorData(
-                    refreshToken = refreshToken,
+                    refreshToken = refreshToken?.tokenValue,
                     error = e
                 )
 
@@ -156,7 +155,7 @@ internal object SignOutCognitoActions : SignOutActions {
             logger.verbose("$id Starting execution")
 
             val error = RevokeTokenErrorData(
-                refreshToken = event.signedInData.cognitoUserPoolTokens.refreshToken,
+                refreshToken = event.signedInData.cognitoUserPoolTokens.refreshToken?.tokenValue,
                 error = Exception("RevokeToken not attempted because GlobalSignOut failed.")
             )
 

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

@@ -19,7 +19,6 @@ import android.content.Context
 import androidx.core.content.edit
 import com.amplifyframework.auth.AuthProvider
 import com.amplifyframework.auth.cognito.AuthConfiguration
-import com.amplifyframework.auth.cognito.helpers.SessionHelper
 import com.amplifyframework.auth.cognito.helpers.identityProviderName
 import com.amplifyframework.core.store.KeyValueRepository
 import com.amplifyframework.statemachine.codegen.data.AWSCredentials
@@ -216,13 +215,13 @@ internal class AWSCognitoLegacyCredentialStore(
         val signInMethod = retrieveUserPoolSignInMethod() ?: return null
         val tokenUserId =
             try {
-                cognitoUserPoolTokens.accessToken?.let { SessionHelper.getUserSub(it) } ?: ""
+                cognitoUserPoolTokens.accessToken?.userSub ?: ""
             } catch (e: Exception) {
                 ""
             }
         val tokenUsername =
             try {
-                cognitoUserPoolTokens.accessToken?.let { SessionHelper.getUsername(it) } ?: ""
+                cognitoUserPoolTokens.accessToken?.username ?: ""
             } catch (e: Exception) {
                 ""
             }

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

@@ -21,29 +21,6 @@ import java.time.Instant
 import java.time.temporal.ChronoUnit
 
 internal object SessionHelper {
-    /**
-     * Returns expiration of this id token.
-     * @return id token expiration claim as {@link java.time.Instant} in UTC.
-     */
-    internal fun getExpiration(token: String): Instant? {
-        val claim = JWTParser.getClaim(token, "exp")
-        return claim?.let {
-            Instant.ofEpochSecond(claim.toLong())
-        }
-    }
-
-    /**
-     * Returns the username set in the access token.
-     * @return Username.
-     */
-    fun getUsername(token: String): String? = JWTParser.getClaim(token, "username")
-
-    /**
-     * Returns the usersub set in the access token.
-     * @return usersub
-     */
-    fun getUserSub(token: String): String? = JWTParser.getClaim(token, "sub")
-
     /**
      * Returns true if the access and id tokens have not expired.
      * @return boolean to indicate if the access and id tokens are expired.
@@ -53,10 +30,9 @@ internal object SessionHelper {
         return when {
             userPoolTokens.idToken == null -> false
             userPoolTokens.accessToken == null -> false
-            else -> currentTimeStamp < getExpiration(userPoolTokens.idToken) &&
-                currentTimeStamp < getExpiration(
-                    userPoolTokens.accessToken
-                )
+            else ->
+                currentTimeStamp < userPoolTokens.idToken.expiration &&
+                    currentTimeStamp < userPoolTokens.accessToken.expiration
         }
     }
 

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

@@ -59,9 +59,9 @@ internal object SignInChallengeHelper {
     ): StateMachineEvent = when {
         authenticationResult != null -> {
             authenticationResult.let {
-                val userId = it.accessToken?.let { token -> SessionHelper.getUserSub(token) } ?: ""
-                val expiresIn = Instant.now().plus(it.expiresIn.seconds).epochSeconds
-                val tokens = CognitoUserPoolTokens(it.idToken, it.accessToken, it.refreshToken, expiresIn)
+                val expiration = Instant.now().plus(it.expiresIn.seconds).epochSeconds
+                val tokens = CognitoUserPoolTokens(it.idToken, it.accessToken, it.refreshToken, expiration)
+                val userId = tokens.accessToken?.userSub ?: ""
                 val signedInData = SignedInData(
                     userId,
                     username,