From 1cd6b3b8e0f923c55f4c47d7e3072af69ae78e4d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9rgio=20Santos?= <me@sergiosantos.info>
Date: Wed, 27 Sep 2023 12:03:31 +0100
Subject: [PATCH] Make EncryptionInitializationException public. Fix test
 coverage.

---
 .../awaladroid/AndroidPrivateKeyStore.kt      | 29 ++++++++++++-------
 .../awaladroid/AndroidPrivateKeyStoreTest.kt  | 11 +++++++
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/lib/src/main/java/tech/relaycorp/awaladroid/AndroidPrivateKeyStore.kt b/lib/src/main/java/tech/relaycorp/awaladroid/AndroidPrivateKeyStore.kt
index 68d9fd3b..dd39611e 100644
--- a/lib/src/main/java/tech/relaycorp/awaladroid/AndroidPrivateKeyStore.kt
+++ b/lib/src/main/java/tech/relaycorp/awaladroid/AndroidPrivateKeyStore.kt
@@ -11,6 +11,14 @@ import javax.crypto.AEADBadTagException
 internal class AndroidPrivateKeyStore(
     root: FileKeystoreRoot,
     private val context: Context,
+    private val encryptedFileBuilder: (File, MasterKey) -> EncryptedFile = { file, masterKey ->
+        EncryptedFile.Builder(
+            context,
+            file,
+            masterKey,
+            EncryptedFile.FileEncryptionScheme.AES256_GCM_HKDF_4KB,
+        ).build()
+    },
 ) : FilePrivateKeyStore(root) {
 
     @Throws(EncryptionInitializationException::class)
@@ -20,17 +28,15 @@ internal class AndroidPrivateKeyStore(
     override fun makeEncryptedOutputStream(file: File) = buildEncryptedFile(file).openFileOutput()
 
     @Throws(EncryptionInitializationException::class)
-    private fun buildEncryptedFile(file: File) =
+    private fun buildEncryptedFile(file: File): EncryptedFile =
         try {
-            EncryptedFile.Builder(
-                context,
-                file,
-                masterKey,
-                EncryptedFile.FileEncryptionScheme.AES256_GCM_HKDF_4KB,
-            ).build()
-        } catch (e: AEADBadTagException) {
+            encryptedFileBuilder(file, masterKey)
+        } catch (exception: AEADBadTagException) {
             // Known issue: https://issuetracker.google.com/issues/164901843
-            throw EncryptionInitializationException(e)
+            throw EncryptionInitializationException(
+                "Could not build encrypted file due to internal issue",
+                exception,
+            )
         }
 
     private val masterKey by lazy {
@@ -39,9 +45,10 @@ internal class AndroidPrivateKeyStore(
             .build()
     }
 
-    class EncryptionInitializationException(cause: Throwable) : Exception(cause)
-
     companion object {
         private const val MASTER_KEY_ALIAS = "_awaladroid_master_key_"
     }
 }
+
+public class EncryptionInitializationException(message: String, cause: Throwable) :
+    AwaladroidException(message, cause)
diff --git a/lib/src/test/java/tech/relaycorp/awaladroid/AndroidPrivateKeyStoreTest.kt b/lib/src/test/java/tech/relaycorp/awaladroid/AndroidPrivateKeyStoreTest.kt
index 6fe6d06c..1b730134 100644
--- a/lib/src/test/java/tech/relaycorp/awaladroid/AndroidPrivateKeyStoreTest.kt
+++ b/lib/src/test/java/tech/relaycorp/awaladroid/AndroidPrivateKeyStoreTest.kt
@@ -12,6 +12,7 @@ import tech.relaycorp.awaladroid.test.FakeAndroidKeyStore
 import tech.relaycorp.relaynet.testing.pki.KeyPairSet
 import tech.relaycorp.relaynet.testing.pki.PDACertPath
 import java.io.File
+import javax.crypto.AEADBadTagException
 
 @RunWith(RobolectricTestRunner::class)
 public class AndroidPrivateKeyStoreTest {
@@ -33,4 +34,14 @@ public class AndroidPrivateKeyStoreTest {
         val retrievedId = store.retrieveIdentityKey(certificate.subjectId)
         assertEquals(id, retrievedId)
     }
+
+    @Test(expected = EncryptionInitializationException::class)
+    public fun failWithAEADBadTagException(): Unit = runTest {
+        val androidContext = RuntimeEnvironment.getApplication()
+        val root = FileKeystoreRoot(File(androidContext.filesDir, "tmp-keystore"))
+        val store = AndroidPrivateKeyStore(root, androidContext) { _, _ ->
+            throw AEADBadTagException("")
+        }
+        store.saveIdentityKey(KeyPairSet.PRIVATE_ENDPOINT.private)
+    }
 }