From eacae2dff9d46d4d575441f67c7ca0afb57001aa Mon Sep 17 00:00:00 2001
From: Tyler Roach <tjroach@amazon.com>
Date: Tue, 15 Nov 2022 13:58:03 -0500
Subject: [PATCH] Prevent attempting to read backed up
 EncryptedSharedPreferences that are no longer readable

---
 .../data/EncryptedKeyValueRepository.kt       | 69 +++++++++++++++----
 .../data/EncryptedKeyValueRepositoryTest.kt   |  4 +-
 2 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/data/EncryptedKeyValueRepository.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/data/EncryptedKeyValueRepository.kt
index d4be967b38..6993e92063 100644
--- a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/data/EncryptedKeyValueRepository.kt
+++ b/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/data/EncryptedKeyValueRepository.kt
@@ -22,36 +22,81 @@ import androidx.security.crypto.EncryptedSharedPreferences
 import androidx.security.crypto.EncryptedSharedPreferences.PrefKeyEncryptionScheme.AES256_SIV
 import androidx.security.crypto.EncryptedSharedPreferences.PrefValueEncryptionScheme.AES256_GCM
 import androidx.security.crypto.MasterKeys
+import java.io.File
+import java.util.UUID
 
 internal class EncryptedKeyValueRepository(
     private val context: Context,
     private val sharedPreferencesName: String,
 ) : KeyValueRepository {
 
+    @VisibleForTesting
+    internal val sharedPreferences: SharedPreferences by lazy {
+        EncryptedSharedPreferences.create(
+            "$sharedPreferencesName.${getInstallationIdentifier(context, sharedPreferencesName)}",
+            MasterKeys.getOrCreate(MasterKeys.AES256_GCM_SPEC),
+            context,
+            AES256_SIV,
+            AES256_GCM
+        )
+    }
+
+    @VisibleForTesting
+    internal val editor: SharedPreferences.Editor by lazy {
+        sharedPreferences.edit()
+    }
+
     override fun put(dataKey: String, value: String?) {
-        with(getSharedPreferences().edit()) {
+        with(editor) {
             putString(dataKey, value)
             apply()
         }
     }
 
-    override fun get(dataKey: String): String? = getSharedPreferences().getString(dataKey, null)
+    override fun get(dataKey: String): String? = sharedPreferences.getString(dataKey, null)
 
     override fun remove(dataKey: String) {
-        with(getSharedPreferences().edit()) {
+        with(editor) {
             remove(dataKey)
             apply()
         }
     }
 
-    @VisibleForTesting
-    fun getSharedPreferences(): SharedPreferences {
-        return EncryptedSharedPreferences.create(
-            sharedPreferencesName,
-            MasterKeys.getOrCreate(MasterKeys.AES256_GCM_SPEC),
-            context,
-            AES256_SIV,
-            AES256_GCM
-        )
+    /**
+     * EncryptedSharedPreferences may have been backed up by the application, but will be unreadable due to the
+     * KeyStore record being lost. To prevent an unreadable EncryptedSharedPreferences, we append a suffix to the name
+     * with a UUID created in the noBackupFilesDir
+     */
+    @Synchronized
+    private fun getInstallationIdentifier(context: Context, keyValueRepoID: String): String {
+        val identifierFile = File(context.noBackupFilesDir, "$keyValueRepoID.installationIdentifier")
+        val previousIdentifier = getExistingInstallationIdentifier(identifierFile)
+
+        return previousIdentifier ?: createInstallationIdentifier(identifierFile)
+    }
+
+    /**
+     * Gets the existing installation identifier (if exists)
+     */
+    private fun getExistingInstallationIdentifier(identifierFile: File): String? {
+        return if (identifierFile.exists()) {
+            val identifier = identifierFile.readText()
+            identifier.ifBlank { null }
+        } else {
+            null
+        }
+    }
+
+    /**
+     * Creates a new installation identifier for the install
+     */
+    private fun createInstallationIdentifier(identifierFile: File): String {
+        val newIdentifier = UUID.randomUUID().toString()
+        try {
+            identifierFile.writeText(newIdentifier)
+        } catch (e: Exception) {
+            // Failed to write identifier to file, session will be forced to be in memory
+        }
+        return newIdentifier
     }
 }
diff --git a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/data/EncryptedKeyValueRepositoryTest.kt b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/data/EncryptedKeyValueRepositoryTest.kt
index 2cbd1e0ad9..94a4ff4fec 100644
--- a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/data/EncryptedKeyValueRepositoryTest.kt
+++ b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/data/EncryptedKeyValueRepositoryTest.kt
@@ -49,8 +49,8 @@ class EncryptedKeyValueRepositoryTest {
 
     @Before
     fun setup() {
-        Mockito.`when`(repository.getSharedPreferences()).thenReturn(mockPrefs)
-        Mockito.`when`(mockPrefs.edit()).thenReturn(mockPrefsEditor)
+        Mockito.`when`(repository.sharedPreferences).thenReturn(mockPrefs)
+        Mockito.`when`(repository.editor).thenReturn(mockPrefsEditor)
     }
 
     @Test