fix(auth): Fix Device Metadata migration if alised userId was used

aws-auth-cognito/src/androidTest/java/com/amplifyframework/auth/cognito/AWSCognitoLegacyCredentialStoreInstrumentationTest.kt

@@ -20,7 +20,8 @@ import androidx.test.platform.app.InstrumentationRegistry
 import com.amplifyframework.auth.cognito.data.AWSCognitoLegacyCredentialStore
 import com.amplifyframework.auth.cognito.testutils.AuthConfigurationProvider
 import com.amplifyframework.auth.cognito.testutils.CredentialStoreUtil
-import org.junit.Assert.assertTrue
+import kotlin.test.assertEquals
+import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -31,20 +32,42 @@ class AWSCognitoLegacyCredentialStoreInstrumentationTest {
 
     private val configuration: AuthConfiguration = AuthConfigurationProvider.getAuthConfiguration()
 
-    private val credential = CredentialStoreUtil.getDefaultCredential()
+    private val credentialStoreUtil = CredentialStoreUtil()
+    private val credential = credentialStoreUtil.getDefaultCredential()
 
     private lateinit var store: AWSCognitoLegacyCredentialStore
 
     @Before
     fun setup() {
         store = AWSCognitoLegacyCredentialStore(context, configuration)
         // TODO: Pull the appClientID from the configuration instead of hardcoding
-        CredentialStoreUtil.setupLegacyStore(context, "userPoolAppClientId", "userPoolId", "identityPoolId")
+        credentialStoreUtil.setupLegacyStore(context, "userPoolAppClientId", "userPoolId", "identityPoolId")
+    }
+
+    @After
+    fun tearDown() {
+        credentialStoreUtil.clearSharedPreferences(context)
     }
 
     @Test
     fun test_legacy_store_implementation_can_retrieve_credentials_stored_using_aws_sdk() {
-        val creds = store.retrieveCredential()
-        assertTrue(creds == credential)
+        assertEquals(credential, store.retrieveCredential())
+    }
+
+    @Test
+    fun test_legacy_store_implementation_can_retrieve_device_metadata_using_aws_sdk() {
+        val user1DeviceMetadata = store.retrieveDeviceMetadata(credentialStoreUtil.user1Username)
+        val user2DeviceMetadata = store.retrieveDeviceMetadata(credentialStoreUtil.user2Username)
+
+        assertEquals(credentialStoreUtil.getUser1DeviceMetadata(), user1DeviceMetadata)
+        assertEquals(credentialStoreUtil.getUser2DeviceMetadata(), user2DeviceMetadata)
+    }
+
+    @Test
+    fun test_legacy_store_implementation_can_retrieve_usernames_for_device_metadata() {
+        val expectedUsernames = listOf(credentialStoreUtil.user1Username, credentialStoreUtil.user2Username)
+        val deviceMetadataUsernames = store.retrieveDeviceMetadataUsernameList()
+
+        assertEquals(expectedUsernames, deviceMetadataUsernames)
     }
 }

aws-auth-cognito/src/androidTest/java/com/amplifyframework/auth/cognito/CredentialStoreStateMachineInstrumentationTest.kt

@@ -20,27 +20,33 @@ import androidx.test.platform.app.InstrumentationRegistry
 import com.amplifyframework.auth.cognito.data.AWSCognitoAuthCredentialStore
 import com.amplifyframework.auth.cognito.testutils.AuthConfigurationProvider
 import com.amplifyframework.auth.cognito.testutils.CredentialStoreUtil
+import com.amplifyframework.statemachine.codegen.data.DeviceMetadata
 import com.google.gson.Gson
 import junit.framework.TestCase.assertEquals
 import org.json.JSONObject
+import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
 @RunWith(AndroidJUnit4::class)
 class CredentialStoreStateMachineInstrumentationTest {
     private val context = InstrumentationRegistry.getInstrumentation().context
+    private val credentialStoreUtil = CredentialStoreUtil()
 
     private val configuration = AuthConfigurationProvider.getAuthConfigurationObject()
     private val userPoolId = configuration.userPool.userPool.PoolId
     private val identityPoolId = configuration.credentials.cognitoIdentity.identityData.PoolId
     private val userPoolAppClientId = configuration.userPool.userPool.AppClientId
 
-    private val credential = CredentialStoreUtil.getDefaultCredential()
-
     @Before
     fun setup() {
-        CredentialStoreUtil.setupLegacyStore(context, userPoolAppClientId, userPoolId, identityPoolId)
+        credentialStoreUtil.setupLegacyStore(context, userPoolAppClientId, userPoolId, identityPoolId)
+    }
+
+    @After
+    fun tearDown() {
+        credentialStoreUtil.clearSharedPreferences(context)
     }
 
     private val authConfigJson = JSONObject(Gson().toJson(configuration))
@@ -51,11 +57,83 @@ class CredentialStoreStateMachineInstrumentationTest {
         plugin.configure(authConfigJson, context)
         plugin.initialize(context)
 
-        val receivedCredentials = AWSCognitoAuthCredentialStore(
+        val credentialStore = AWSCognitoAuthCredentialStore(
+            context,
+            AuthConfiguration.fromJson(authConfigJson)
+        )
+
+        assertEquals(credentialStoreUtil.getDefaultCredential(), credentialStore.retrieveCredential())
+        assertEquals(
+            credentialStoreUtil.getUser1DeviceMetadata(),
+            credentialStore.retrieveDeviceMetadata(credentialStoreUtil.user1Username)
+        )
+        assertEquals(
+            credentialStoreUtil.getUser2DeviceMetadata(),
+            credentialStore.retrieveDeviceMetadata(credentialStoreUtil.user2Username)
+        )
+    }
+
+    @Test
+    fun test_CredentialStore_Missing_DeviceMetadata_Migration_Succeeds_On_Plugin_Configuration() {
+        // GIVEN
+        val userAUsername = "userA"
+        val expectedUserADeviceMetadata = DeviceMetadata.Metadata("A", "B", "C")
+        val userBUsername = "userB"
+        val expectedUserBDeviceMetadata = DeviceMetadata.Metadata("1", "2", "3")
+
+        AWSCognitoAuthPlugin().apply {
+            configure(authConfigJson, context)
+            initialize(context)
+        }
+
+        AWSCognitoAuthCredentialStore(
+            context,
+            AuthConfiguration.fromJson(authConfigJson)
+        ).apply {
+            saveDeviceMetadata("userA", expectedUserADeviceMetadata)
+        }
+
+        // WHEN
+        // Simulating missed device metadata migration from issue 2929
+        // We expect this to not migrate as it will conflict with existing metadata already saved
+        credentialStoreUtil.saveLegacyDeviceMetadata(
+            context,
+            userPoolId,
+            userAUsername,
+            DeviceMetadata.Metadata("X", "Y", "Z")
+        )
+
+        // We expect this to migrate as it does not conflict with any existing saved metadata
+        credentialStoreUtil.saveLegacyDeviceMetadata(
+            context,
+            userPoolId,
+            userBUsername,
+            expectedUserBDeviceMetadata
+        )
+
+        // THEN
+
+        // Initialize plugin again to complete migration of missing device metadata
+        AWSCognitoAuthPlugin().apply {
+            configure(authConfigJson, context)
+            initialize(context)
+        }
+
+        // WHEN
+        val credentialStore = AWSCognitoAuthCredentialStore(
             context,
             AuthConfiguration.fromJson(authConfigJson)
-        ).retrieveCredential()
+        )
 
-        assertEquals(credential, receivedCredentials)
+        // Expect the device metadata for user A to have not changed from data that was already saved in v2 store
+        assertEquals(
+            expectedUserADeviceMetadata,
+            credentialStore.retrieveDeviceMetadata(userAUsername)
+        )
+        // Expect the device metadata for user A to have not changed from data that was already saved in v2 store
+        assertEquals(
+            expectedUserBDeviceMetadata,
+            credentialStore.retrieveDeviceMetadata(userBUsername)
+        )
     }
 }

aws-auth-cognito/src/androidTest/java/com/amplifyframework/auth/cognito/testutils/CredentialStoreUtil.kt

@@ -20,13 +20,14 @@ import com.amazonaws.internal.keyvaluestore.AWSKeyValueStore
 import com.amplifyframework.statemachine.codegen.data.AWSCredentials
 import com.amplifyframework.statemachine.codegen.data.AmplifyCredential
 import com.amplifyframework.statemachine.codegen.data.CognitoUserPoolTokens
+import com.amplifyframework.statemachine.codegen.data.DeviceMetadata
 import com.amplifyframework.statemachine.codegen.data.SignInMethod
 import com.amplifyframework.statemachine.codegen.data.SignedInData
+import java.io.File
 import java.util.Date
 
-internal object CredentialStoreUtil {
-
-    private const val accessToken = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwidXNlcm5hbWUiO" +
+internal class CredentialStoreUtil {
+    private val accessToken = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwidXNlcm5hbWUiO" +
         "iJhbXBsaWZ5X3VzZXIiLCJpYXQiOjE1MTYyMzkwMjJ9.zBiQ0guLRX34pUEYLPyDxQAyDDlXmL0JY7kgPWAHZos"
 
     private val credential = AmplifyCredential.UserAndIdentityPool(
@@ -50,8 +51,29 @@ internal object CredentialStoreUtil {
         return credential
     }
 
+    val user1Username = "2924030b-54c0-48bc-8bff-948418fba949"
+    val user2Username = "7e001127-5f11-41fb-9d10-ab9d6cf41dba"
+
+    fun getUser1DeviceMetadata(): DeviceMetadata.Metadata {
+        return DeviceMetadata.Metadata(
+            "DeviceKey1",
+            "DeviceGroupKey1",
+            "DeviceSecret1"
+        )
+    }
+
+    fun getUser2DeviceMetadata(): DeviceMetadata.Metadata {
+        return DeviceMetadata.Metadata(
+            "DeviceKey2",
+            "DeviceGroupKey2",
+            "DeviceSecret2"
+        )
+    }
+
     fun setupLegacyStore(context: Context, appClientId: String, userPoolId: String, identityPoolId: String) {
 
+        clearSharedPreferences(context)
+
         AWSKeyValueStore(context, "CognitoIdentityProviderCache", true).apply {
             put("CognitoIdentityProvider.$appClientId.testuser.idToken", "idToken")
             put("CognitoIdentityProvider.$appClientId.testuser.accessToken", accessToken)
@@ -60,10 +82,16 @@ internal object CredentialStoreUtil {
             put("CognitoIdentityProvider.$appClientId.LastAuthUser", "testuser")
         }
 
-        AWSKeyValueStore(context, "CognitoIdentityProviderDeviceCache.$userPoolId.testuser", true).apply {
-            put("DeviceKey", "someDeviceKey")
-            put("DeviceGroupKey", "someDeviceGroupKey")
-            put("DeviceSecret", "someSecret")
+        AWSKeyValueStore(context, "CognitoIdentityProviderDeviceCache.$userPoolId.$user1Username", true).apply {
+            put("DeviceKey", "DeviceKey1")
+            put("DeviceGroupKey", "DeviceGroupKey1")
+            put("DeviceSecret", "DeviceSecret1")
+        }
+
+        AWSKeyValueStore(context, "CognitoIdentityProviderDeviceCache.$userPoolId.$user2Username", true).apply {
+            put("DeviceKey", "DeviceKey2")
+            put("DeviceGroupKey", "DeviceGroupKey2")
+            put("DeviceSecret", "DeviceSecret2")
         }
 
         AWSKeyValueStore(context, "com.amazonaws.android.auth", true).apply {
@@ -73,5 +101,46 @@ internal object CredentialStoreUtil {
             put("$identityPoolId.expirationDate", "1212")
             put("$identityPoolId.identityId", "identityId")
         }
+
+        // we need to wait for shared prefs to actually hit filesystem as we always use apply instead of commit
+        val beginWait = System.currentTimeMillis()
+        while (System.currentTimeMillis() - beginWait < 3000) {
+            if ((File(context.dataDir, "shared_prefs").listFiles()?.size ?: 0) >= 4) {
+                break
+            } else {
+                Thread.sleep(50)
+            }
+        }
+    }
+
+    fun saveLegacyDeviceMetadata(
+        context: Context,
+        userPoolId: String,
+        username: String,
+        deviceMetadata: DeviceMetadata.Metadata
+    ) {
+        val prefsName = "CognitoIdentityProviderDeviceCache.$userPoolId.$username"
+        AWSKeyValueStore(
+            context,
+            "CognitoIdentityProviderDeviceCache.$userPoolId.$username", true
+        ).apply {
+            put("DeviceKey", deviceMetadata.deviceKey)
+            put("DeviceGroupKey", deviceMetadata.deviceGroupKey)
+            put("DeviceSecret", deviceMetadata.deviceSecret)
+        }
+
+        // we need to wait for shared prefs to actually hit filesystem as we always use apply instead of commit
+        val beginWait = System.currentTimeMillis()
+        while (System.currentTimeMillis() - beginWait < 3000) {
+            if (File(context.dataDir, "shared_prefs/$prefsName.xml").exists()) {
+                break
+            } else {
+                Thread.sleep(50)
+            }
+        }
+    }
+
+    fun clearSharedPreferences(context: Context) {
+        File(context.dataDir, "shared_prefs").listFiles()?.forEach { it.delete() }
     }
 }

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

@@ -36,13 +36,27 @@ internal object CredentialStoreCognitoActions : CredentialStoreActions {
                     legacyCredentialStore.deleteCredential()
                 }
 
-                // migrate device data
-                val lastAuthUserId = legacyCredentialStore.retrieveLastAuthUserId()
-                lastAuthUserId?.let {
-                    val deviceMetaData = legacyCredentialStore.retrieveDeviceMetadata(lastAuthUserId)
+                /*
+                Migrate Device Metadata
+                1. We first need to get the list of usernames that contain device metadata on the device.
+                2. For each username, we check to see if the current credential store has device metadata for that user.
+                3. If the current user does not have device metadata in the current store, migrate from legacy.
+                This is a possibility because of a bug where we were previously attempting to migrate using an aliased
+                username lookup.
+                4. If the current user has device metadata in the current credential store, do not migrate from legacy.
+                This situation would happen if a user updated from legacy, signed out, then signed back in. Upon
+                signing back in, they would be granted new device metadata. Since that new metadata is what is
+                associated with the refresh token, we do not want to overwrite it with legacy metadata.
+                5. Upon completed migration, we delete the legacy device metadata.
+                 */
+                legacyCredentialStore.retrieveDeviceMetadataUsernameList().forEach { username ->
+                    val deviceMetaData = legacyCredentialStore.retrieveDeviceMetadata(username)
                     if (deviceMetaData != DeviceMetadata.Empty) {
-                        credentialStore.saveDeviceMetadata(lastAuthUserId, deviceMetaData)
-                        legacyCredentialStore.deleteDeviceKeyCredential(lastAuthUserId)
+                        credentialStore.retrieveDeviceMetadata(username)
+                        if (credentialStore.retrieveDeviceMetadata(username) == DeviceMetadata.Empty) {
+                            credentialStore.saveDeviceMetadata(username, deviceMetaData)
+                        }
+                        legacyCredentialStore.deleteDeviceKeyCredential(username)
                     }
                 }
 

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

@@ -29,6 +29,7 @@ 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.io.File
 import java.time.Instant
 import java.time.temporal.ChronoUnit
 import java.util.Date
@@ -74,6 +75,7 @@ internal class AWSCognitoLegacyCredentialStore(
         const val TOKEN_KEY = "token"
     }
 
+    private val userDeviceDetailsCacheKeyPrefix = "$APP_DEVICE_INFO_CACHE.${authConfiguration.userPool?.poolId}."
     private val userDeviceDetailsCacheKey = "$APP_DEVICE_INFO_CACHE.${authConfiguration.userPool?.poolId}.%s"
 
     private val idAndCredentialsKeyValue: KeyValueRepository by lazy {
@@ -229,6 +231,22 @@ internal class AWSCognitoLegacyCredentialStore(
         )
     }
 
+    /*
+    During migration away from the legacy credential store, we need to find all shared preference files that store
+    device metadata. These filenames contain the real username (not aliased) for the tracked device metadata.
+     */
+    fun retrieveDeviceMetadataUsernameList(): List<String> {
+        return try {
+            val sharedPrefsSuffix = ".xml"
+            File(context.dataDir, "shared_prefs").listFiles { _, filename ->
+                filename.startsWith(userDeviceDetailsCacheKeyPrefix) && filename.endsWith(sharedPrefsSuffix)
+            }?.map { it.name.substringAfter(userDeviceDetailsCacheKeyPrefix).substringBefore(sharedPrefsSuffix) }
+                ?.filter { it.isNotBlank() } ?: emptyList()
+        } catch (e: Exception) {
+            return emptyList()
+        }
+    }
+
     @Synchronized
     override fun retrieveDeviceMetadata(username: String): DeviceMetadata {
         val deviceDetailsCacheKey = String.format(userDeviceDetailsCacheKey, username)