From 9cefd66c66fccb75070f218c104b376269a74017 Mon Sep 17 00:00:00 2001
From: Osip Fatkullin <osip.fatkullin@gmail.com>
Date: Sun, 26 Feb 2023 12:45:11 +0300
Subject: [PATCH] Fix decryption of non-buffered input stream (#16)

* fix: Add test for decryption of non-buffered input stream
* fix: Fix decryption of non-buffered input stream
---
 CHANGELOG.md                                  |  4 +++
 encrypted-datastore/src/main/kotlin/Aead.kt   |  2 ++
 .../migration/StreamingAeadWithFallback.kt    |  8 ++++--
 .../StreamingAeadWithFallbackTest.kt          | 25 ++++++++++++++++++-
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9ba9f45..af9cf5c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 ## [Unreleased]
 
+### Fixed
+
+- Fixed decryption fallback from StreamingAead to Aead (#15)
+
 ### Housekeeping
 
 - Update Gradle to 7.6.1
diff --git a/encrypted-datastore/src/main/kotlin/Aead.kt b/encrypted-datastore/src/main/kotlin/Aead.kt
index 7ac7c3e..19218d1 100644
--- a/encrypted-datastore/src/main/kotlin/Aead.kt
+++ b/encrypted-datastore/src/main/kotlin/Aead.kt
@@ -4,6 +4,8 @@ import com.google.crypto.tink.Aead
 import java.io.InputStream
 
 internal fun Aead.newDecryptedStream(inputStream: InputStream): InputStream {
+    // Method 'decrypt' throws GeneralSecurityException for empty byte array,
+    // so let's check it is not empty.
     return if (inputStream.available() > 0) {
         decrypt(inputStream.readBytes(), null).inputStream()
     } else {
diff --git a/encrypted-datastore/src/main/kotlin/migration/StreamingAeadWithFallback.kt b/encrypted-datastore/src/main/kotlin/migration/StreamingAeadWithFallback.kt
index 38a3bb5..5ec20b2 100644
--- a/encrypted-datastore/src/main/kotlin/migration/StreamingAeadWithFallback.kt
+++ b/encrypted-datastore/src/main/kotlin/migration/StreamingAeadWithFallback.kt
@@ -13,9 +13,13 @@ internal class StreamingAeadWithFallback(
 ) : StreamingAead by delegate {
 
     override fun newDecryptingStream(ciphertextSource: InputStream, associatedData: ByteArray): InputStream {
+        // Input Stream should support mark and reset to make it reusable in fallback stream.
+        // NOTE: mark is called in delegate.newDecryptingStream
+        val inputStream = if (ciphertextSource.markSupported()) ciphertextSource else ciphertextSource.buffered()
+
         return DecryptingStreamWithFallback(
-            stream = delegate.newDecryptingStream(ciphertextSource, associatedData),
-            fallbackStream = { fallback.newDecryptedStream(ciphertextSource) },
+            stream = delegate.newDecryptingStream(inputStream, associatedData),
+            fallbackStream = { fallback.newDecryptedStream(inputStream) },
         )
     }
 
diff --git a/encrypted-datastore/src/test/kotlin/migration/StreamingAeadWithFallbackTest.kt b/encrypted-datastore/src/test/kotlin/migration/StreamingAeadWithFallbackTest.kt
index 8d88c64..3c637d1 100644
--- a/encrypted-datastore/src/test/kotlin/migration/StreamingAeadWithFallbackTest.kt
+++ b/encrypted-datastore/src/test/kotlin/migration/StreamingAeadWithFallbackTest.kt
@@ -7,7 +7,12 @@ import com.google.crypto.tink.StreamingAead
 import com.google.crypto.tink.aead.AeadConfig
 import com.google.crypto.tink.streamingaead.StreamingAeadConfig
 import org.junit.jupiter.api.assertThrows
+import org.junit.jupiter.api.io.TempDir
 import java.io.IOException
+import java.nio.file.Path
+import kotlin.io.path.div
+import kotlin.io.path.inputStream
+import kotlin.io.path.writeBytes
 import kotlin.test.Test
 import kotlin.test.assertEquals
 
@@ -16,6 +21,9 @@ internal class StreamingAeadWithFallbackTest {
     private var aead: Aead
     private var streamingAead: StreamingAead
 
+    @field:TempDir
+    lateinit var tempDir: Path
+
     init {
         AeadConfig.register()
         StreamingAeadConfig.register()
@@ -37,7 +45,7 @@ internal class StreamingAeadWithFallbackTest {
     }
 
     @Test
-    fun `decrypt encrypted stream with decryption fallback`() {
+    fun `decrypt encrypted bytes stream with decryption fallback`() {
         val plaintext = "Plaintext"
         val encryptedStream = aead.encrypt(plaintext.toByteArray(), null).inputStream()
         val decryptingStream = streamingAead.withDecryptionFallback(aead)
@@ -46,4 +54,19 @@ internal class StreamingAeadWithFallbackTest {
         val decrypted = decryptingStream.readBytes().decodeToString()
         assertEquals(plaintext, decrypted)
     }
+
+    @Test
+    fun `decrypt encrypted file stream with decryption fallback`() {
+        val plaintext = "Plaintext"
+
+        val file = tempDir / "encrypted-file"
+        file.writeBytes(aead.encrypt(plaintext.toByteArray(), null))
+
+        val encryptedStream = file.inputStream()
+        val decryptingStream = streamingAead.withDecryptionFallback(aead)
+            .newDecryptingStream(encryptedStream, byteArrayOf())
+
+        val decrypted = decryptingStream.readBytes().decodeToString()
+        assertEquals(plaintext, decrypted)
+    }
 }