fix(auth): Fix timestamp issue with v1 to v2 migration (#2799)

Author: vincetran

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

@@ -29,6 +29,8 @@ import com.amplifyframework.statemachine.codegen.data.DeviceMetadata
 import com.amplifyframework.statemachine.codegen.data.FederatedToken
 import com.amplifyframework.statemachine.codegen.data.SignInMethod
 import com.amplifyframework.statemachine.codegen.data.SignedInData
+import java.time.Instant
+import java.time.temporal.ChronoUnit
 import java.util.Date
 import java.util.Locale
 
@@ -180,7 +182,17 @@ internal class AWSCognitoLegacyCredentialStore(
         val accessKey = idAndCredentialsKeyValue.get(namespace(AK_KEY))
         val secretKey = idAndCredentialsKeyValue.get(namespace(SK_KEY))
         val sessionToken = idAndCredentialsKeyValue.get(namespace(ST_KEY))
-        val expiration = idAndCredentialsKeyValue.get(namespace(EXP_KEY))?.toLongOrNull()
+
+        // V2 AWSCredential expiration is in epoch seconds whereas legacy expiration may be in epoch milliseconds
+        // Session expiration should be within 24 hours so if we see a date in the far future, we can assume the
+        // timestamp is encoded in milliseconds.
+        val expiration = idAndCredentialsKeyValue.get(namespace(EXP_KEY))?.toLongOrNull()?.let {
+            if (Instant.ofEpochSecond(it).isAfter(Instant.now().plus(365, ChronoUnit.DAYS))) {
+                it / 1000
+            } else {
+                it
+            }
+        }
 
         return if (accessKey == null && secretKey == null && sessionToken == null) {
             null

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

@@ -18,6 +18,7 @@ package com.amplifyframework.auth.cognito.helpers
 import com.amplifyframework.statemachine.codegen.data.AWSCredentials
 import com.amplifyframework.statemachine.codegen.data.CognitoUserPoolTokens
 import java.time.Instant
+import java.time.temporal.ChronoUnit
 
 internal object SessionHelper {
     /**
@@ -68,6 +69,14 @@ internal object SessionHelper {
      */
     fun isValidSession(awsCredentials: AWSCredentials): Boolean {
         val currentTimeStamp = Instant.now()
-        return currentTimeStamp < awsCredentials.expiration?.let { Instant.ofEpochSecond(it) }
+        val credentialsExpirationInSecond = awsCredentials.expiration?.let { Instant.ofEpochSecond(it) }
+
+        // Check if current timestamp is BEFORE expiration && next year is AFTER expiration
+        // The latter check is to fix v1 > v2 migration issues as found in:
+        // https://github.com/aws-amplify/amplify-android/issues/2789
+        return (
+            currentTimeStamp < credentialsExpirationInSecond &&
+                currentTimeStamp.plus(365, ChronoUnit.DAYS) > credentialsExpirationInSecond
+            )
     }
 }

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

@@ -93,6 +93,9 @@ class AWSCognitoLegacyCredentialStoreTest {
 
         private const val userDeviceDetailsCacheKey = "$deviceCachePrefix.$USER_POOL_ID.%s"
         private val deviceDetailsCacheKey = String.format(userDeviceDetailsCacheKey, userId)
+
+        private const val expirationTimestampInSec: Long = 1714431706
+        private const val expirationTimestampInMs: Long = 1714431706486
     }
 
     @Mock
@@ -181,7 +184,7 @@ class AWSCognitoLegacyCredentialStoreTest {
         `when`(mockKeyValue.get(cachedIdTokenKey)).thenReturn("idToken")
         `when`(mockKeyValue.get(cachedAccessTokenKey)).thenReturn(dummyToken)
         `when`(mockKeyValue.get(cachedRefreshTokenKey)).thenReturn("refreshToken")
-        `when`(mockKeyValue.get(cachedTokenExpirationKey)).thenReturn("123123")
+        `when`(mockKeyValue.get(cachedTokenExpirationKey)).thenReturn(expirationTimestampInMs.toString())
 
         // Device Metadata
         `when`(mockKeyValue.get("DeviceKey")).thenReturn("someDeviceKey")
@@ -192,7 +195,7 @@ class AWSCognitoLegacyCredentialStoreTest {
         `when`(mockKeyValue.get("$IDENTITY_POOL_ID.${"accessKey"}")).thenReturn("accessKeyId")
         `when`(mockKeyValue.get("$IDENTITY_POOL_ID.${"secretKey"}")).thenReturn("secretAccessKey")
         `when`(mockKeyValue.get("$IDENTITY_POOL_ID.${"sessionToken"}")).thenReturn("sessionToken")
-        `when`(mockKeyValue.get("$IDENTITY_POOL_ID.${"expirationDate"}")).thenReturn("123123")
+        `when`(mockKeyValue.get("$IDENTITY_POOL_ID.${"expirationDate"}")).thenReturn(expirationTimestampInMs.toString())
 
         // Identity ID
         `when`(mockKeyValue.get("$IDENTITY_POOL_ID.${"identityId"}")).thenReturn("identityPool")
@@ -225,10 +228,10 @@ class AWSCognitoLegacyCredentialStoreTest {
                 "amplify_user",
                 Date(0),
                 SignInMethod.ApiBased(SignInMethod.ApiBased.AuthType.USER_SRP_AUTH),
-                CognitoUserPoolTokens("idToken", dummyToken, "refreshToken", 123123)
+                CognitoUserPoolTokens("idToken", dummyToken, "refreshToken", expirationTimestampInSec)
             ),
             "identityPool",
-            AWSCredentials("accessKeyId", "secretAccessKey", "sessionToken", 123123)
+            AWSCredentials("accessKeyId", "secretAccessKey", "sessionToken", expirationTimestampInSec)
         )
     }
 
@@ -249,10 +252,10 @@ class AWSCognitoLegacyCredentialStoreTest {
                 "amplify_user",
                 Date(0),
                 SignInMethod.HostedUI(),
-                CognitoUserPoolTokens("idToken", dummyToken, "refreshToken", 123123)
+                CognitoUserPoolTokens("idToken", dummyToken, "refreshToken", expirationTimestampInSec)
             ),
             "identityPool",
-            AWSCredentials("accessKeyId", "secretAccessKey", "sessionToken", 123123)
+            AWSCredentials("accessKeyId", "secretAccessKey", "sessionToken", expirationTimestampInSec)
         )
     }
 }

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

@@ -18,8 +18,10 @@ package com.amplifyframework.auth.cognito.helpers
 import com.amplifyframework.statemachine.codegen.data.AWSCredentials
 import com.amplifyframework.statemachine.codegen.data.CognitoUserPoolTokens
 import java.time.Instant
+import java.time.temporal.ChronoUnit
 import kotlin.test.assertEquals
 import kotlin.test.assertFalse
+import kotlin.test.assertTrue
 import org.junit.Test
 
 class SessionHelperTests {
@@ -61,4 +63,47 @@ class SessionHelperTests {
     fun testsIsInvalidSession() {
         assertFalse(SessionHelper.isValidSession(AWSCredentials.empty))
     }
+
+    @Test
+    fun `Pulling a V1 credential should fail isValidSession check`() {
+        // Expiration is encoded in ms to simulate v1 > v2 migration issue
+        assertFalse(
+            SessionHelper.isValidSession(
+                AWSCredentials(
+                    accessKeyId = dummyToken,
+                    secretAccessKey = dummyToken,
+                    sessionToken = dummyToken,
+                    expiration = Instant.now().plus(30, ChronoUnit.MINUTES).toEpochMilli()
+                )
+            )
+        )
+    }
+
+    @Test
+    fun `Session with an expiration in the past should fail isValidSession check`() {
+        assertFalse(
+            SessionHelper.isValidSession(
+                AWSCredentials(
+                    accessKeyId = dummyToken,
+                    secretAccessKey = dummyToken,
+                    sessionToken = dummyToken,
+                    expiration = Instant.now().minus(1, ChronoUnit.MINUTES).epochSecond
+                )
+            )
+        )
+    }
+
+    @Test
+    fun `Session with an expiration in the future should pass isValidSession check`() {
+        assertTrue(
+            SessionHelper.isValidSession(
+                AWSCredentials(
+                    accessKeyId = dummyToken,
+                    secretAccessKey = dummyToken,
+                    sessionToken = dummyToken,
+                    expiration = Instant.now().plus(1, ChronoUnit.MINUTES).epochSecond
+                )
+            )
+        )
+    }
 }