Merge pull request #60 from grote/check-messages

Encrypt values of key/value backups with multiple segments if needed

app/src/main/java/com/stevesoltys/seedvault/crypto/CipherFactory.kt

@@ -6,7 +6,7 @@ import javax.crypto.Cipher.ENCRYPT_MODE
 import javax.crypto.spec.GCMParameterSpec
 
 private const val CIPHER_TRANSFORMATION = "AES/GCM/NoPadding"
-private const val GCM_AUTHENTICATION_TAG_LENGTH = 128
+internal const val GCM_AUTHENTICATION_TAG_LENGTH = 128
 
 interface CipherFactory {
     fun createEncryptionCipher(): Cipher

app/src/main/java/com/stevesoltys/seedvault/crypto/Crypto.kt

@@ -5,6 +5,8 @@ import java.io.EOFException
 import java.io.IOException
 import java.io.InputStream
 import java.io.OutputStream
+import javax.crypto.Cipher
+import kotlin.math.min
 
 /**
  * A backup stream starts with a version byte followed by an encrypted [VersionHeader].
@@ -50,6 +52,14 @@ interface Crypto {
     @Throws(IOException::class)
     fun encryptSegment(outputStream: OutputStream, cleartext: ByteArray)
 
+    /**
+     * Like [encryptSegment],
+     * but if the given cleartext [ByteArray] is larger than [MAX_SEGMENT_CLEARTEXT_LENGTH],
+     * multiple segments will be written.
+     */
+    @Throws(IOException::class)
+    fun encryptMultipleSegments(outputStream: OutputStream, cleartext: ByteArray)
+
     /**
      * Reads and decrypts a [VersionHeader] from the given [InputStream]
      * and ensures that the expected version, package name and key match
@@ -69,6 +79,12 @@ interface Crypto {
      */
     @Throws(EOFException::class, IOException::class, SecurityException::class)
     fun decryptSegment(inputStream: InputStream): ByteArray
+
+    /**
+     * Like [decryptSegment], but decrypts multiple segments and does not throw [EOFException].
+     */
+    @Throws(IOException::class, SecurityException::class)
+    fun decryptMultipleSegments(inputStream: InputStream): ByteArray
 }
 
 internal class CryptoImpl(
@@ -87,9 +103,27 @@ internal class CryptoImpl(
     override fun encryptSegment(outputStream: OutputStream, cleartext: ByteArray) {
         val cipher = cipherFactory.createEncryptionCipher()
 
-        check(cipher.getOutputSize(cleartext.size) <= MAX_SEGMENT_LENGTH)
+        check(cipher.getOutputSize(cleartext.size) <= MAX_SEGMENT_LENGTH) {
+            "Cipher's output size ${cipher.getOutputSize(cleartext.size)} is larger than maximum segment length ($MAX_SEGMENT_LENGTH)"
+        }
+        encryptSegment(cipher, outputStream, cleartext)
+    }
 
-        val encrypted = cipher.doFinal(cleartext)
+    @Throws(IOException::class)
+    override fun encryptMultipleSegments(outputStream: OutputStream, cleartext: ByteArray) {
+        var end = 0
+        while (end < cleartext.size) {
+            val start = end
+            end = min(cleartext.size, start + MAX_SEGMENT_CLEARTEXT_LENGTH)
+            val segment = cleartext.copyOfRange(start, end)
+            val cipher = cipherFactory.createEncryptionCipher()
+            encryptSegment(cipher, outputStream, segment)
+        }
+    }
+
+    @Throws(IOException::class)
+    private fun encryptSegment(cipher: Cipher, outputStream: OutputStream, segment: ByteArray) {
+        val encrypted = cipher.doFinal(segment)
         val segmentHeader = SegmentHeader(encrypted.size.toShort(), cipher.iv)
         headerWriter.writeSegmentHeader(outputStream, segmentHeader)
         outputStream.write(encrypted)
@@ -119,8 +153,21 @@ internal class CryptoImpl(
         return decryptSegment(inputStream, MAX_SEGMENT_LENGTH)
     }
 
+    @Throws(IOException::class, SecurityException::class)
+    override fun decryptMultipleSegments(inputStream: InputStream): ByteArray {
+        var result = ByteArray(0)
+        while (true) {
+            try {
+                result += decryptSegment(inputStream, MAX_SEGMENT_LENGTH)
+            } catch (e: EOFException) {
+                if (result.isEmpty()) throw IOException(e)
+                return result
+            }
+        }
+    }
+
     @Throws(EOFException::class, IOException::class, SecurityException::class)
-    fun decryptSegment(inputStream: InputStream, maxSegmentLength: Int): ByteArray {
+    private fun decryptSegment(inputStream: InputStream, maxSegmentLength: Int): ByteArray {
         val segmentHeader = headerReader.readSegmentHeader(inputStream)
         if (segmentHeader.segmentLength > maxSegmentLength) {
             throw SecurityException("Segment length too long: ${segmentHeader.segmentLength} > $maxSegmentLength")

app/src/main/java/com/stevesoltys/seedvault/header/Header.kt

@@ -1,5 +1,7 @@
 package com.stevesoltys.seedvault.header
 
+import com.stevesoltys.seedvault.crypto.GCM_AUTHENTICATION_TAG_LENGTH
+
 internal const val VERSION: Byte = 0
 internal const val MAX_PACKAGE_LENGTH_SIZE = 255
 internal const val MAX_KEY_LENGTH_SIZE = MAX_PACKAGE_LENGTH_SIZE
@@ -15,14 +17,19 @@ data class VersionHeader(
         internal val key: String? = null      // ?? bytes
 ) {
     init {
-        check(packageName.length <= MAX_PACKAGE_LENGTH_SIZE)
+        check(packageName.length <= MAX_PACKAGE_LENGTH_SIZE) {
-        key?.let { check(key.length <= MAX_KEY_LENGTH_SIZE) }
+            "Package $packageName has name longer than $MAX_PACKAGE_LENGTH_SIZE"
+        }
+        key?.let {
+            check(key.length <= MAX_KEY_LENGTH_SIZE) { "Key $key is longer than $MAX_KEY_LENGTH_SIZE" }
+        }
     }
 }
 
 
 internal const val SEGMENT_LENGTH_SIZE: Int = Short.SIZE_BYTES
 internal const val MAX_SEGMENT_LENGTH: Int = Short.MAX_VALUE.toInt()
+internal const val MAX_SEGMENT_CLEARTEXT_LENGTH: Int = MAX_SEGMENT_LENGTH - GCM_AUTHENTICATION_TAG_LENGTH / 8
 internal const val IV_SIZE: Int = 12
 internal const val SEGMENT_HEADER_SIZE = SEGMENT_LENGTH_SIZE + IV_SIZE
 
@@ -34,6 +41,8 @@ class SegmentHeader(
         internal val nonce: ByteArray         // 12 bytes
 ) {
     init {
-        check(nonce.size == IV_SIZE)
+        check(nonce.size == IV_SIZE) {
+            "Nonce size of ${nonce.size} is not the expected IV size of $IV_SIZE"
+        }
     }
 }

app/src/main/java/com/stevesoltys/seedvault/metadata/MetadataReader.kt

@@ -26,7 +26,7 @@ internal class MetadataReaderImpl(private val crypto: Crypto) : MetadataReader {
         if (version < 0) throw IOException()
         if (version > VERSION) throw UnsupportedVersionException(version)
         val metadataBytes = try {
-            crypto.decryptSegment(inputStream)
+            crypto.decryptMultipleSegments(inputStream)
         } catch (e: AEADBadTagException) {
             throw DecryptionFailedException(e)
         }

app/src/main/java/com/stevesoltys/seedvault/metadata/MetadataWriter.kt

@@ -20,7 +20,7 @@ internal class MetadataWriterImpl(private val crypto: Crypto): MetadataWriter {
     override fun write(outputStream: OutputStream, token: Long) {
         val metadata = BackupMetadata(token = token)
         outputStream.write(ByteArray(1).apply { this[0] = metadata.version })
-        crypto.encryptSegment(outputStream, encode(metadata))
+        crypto.encryptMultipleSegments(outputStream, encode(metadata))
     }
 
     @VisibleForTesting

app/src/main/java/com/stevesoltys/seedvault/restore/RestoreViewModel.kt

@@ -68,7 +68,7 @@ internal class RestoreViewModel(
 
     override fun onRestoreSetClicked(set: RestoreSet) {
         val session = this.session
-        check(session != null)
+        check(session != null) { "Restore set clicked, but no session available" }
         session.restoreAll(set.token, observer, monitor)
 
         mChosenRestoreSet.value = set

app/src/main/java/com/stevesoltys/seedvault/settings/SettingsManager.kt

@@ -35,7 +35,7 @@ class SettingsManager(context: Context) {
     fun getStorage(): Storage? {
         val uriStr = prefs.getString(PREF_KEY_STORAGE_URI, null) ?: return null
         val uri = Uri.parse(uriStr)
-        val name = prefs.getString(PREF_KEY_STORAGE_NAME, null) ?: throw IllegalStateException()
+        val name = prefs.getString(PREF_KEY_STORAGE_NAME, null) ?: throw IllegalStateException("no storage name")
         val isUsb = prefs.getBoolean(PREF_KEY_STORAGE_IS_USB, false)
         return Storage(uri, name, isUsb)
     }

app/src/main/java/com/stevesoltys/seedvault/transport/ConfigurableBackupTransportService.kt

@@ -35,7 +35,7 @@ class ConfigurableBackupTransportService : Service() {
     }
 
     override fun onBind(intent: Intent): IBinder {
-        val transport = this.transport ?: throw IllegalStateException()
+        val transport = this.transport ?: throw IllegalStateException("no transport in onBind()")
         return transport.binder.apply {
             Log.d(TAG, "Transport bound.")
         }

app/src/main/java/com/stevesoltys/seedvault/transport/backup/BackupCoordinator.kt

@@ -178,11 +178,11 @@ internal class BackupCoordinator(
 
     fun finishBackup(): Int = when {
         kv.hasState() -> {
-            check(!full.hasState())
+            check(!full.hasState()) { "K/V backup has state, but full backup has dangling state as well" }
             kv.finishBackup()
         }
         full.hasState() -> {
-            check(!kv.hasState())
+            check(!kv.hasState()) { "Full backup has state, but K/V backup has dangling state as well" }
             full.finishBackup()
         }
         calledInitialize || calledClearBackupData -> {
@@ -190,7 +190,7 @@ internal class BackupCoordinator(
             calledClearBackupData = false
             TRANSPORT_OK
         }
-        else -> throw IllegalStateException()
+        else -> throw IllegalStateException("Unexpected state in finishBackup()")
     }
 
     @Throws(IOException::class)

app/src/main/java/com/stevesoltys/seedvault/transport/backup/KVBackup.kt

@@ -102,7 +102,7 @@ internal class KVBackup(
                     val header = VersionHeader(packageName = packageInfo.packageName, key = op.key)
                     headerWriter.writeVersion(outputStream, header)
                     crypto.encryptHeader(outputStream, header)
-                    crypto.encryptSegment(outputStream, op.value)
+                    crypto.encryptMultipleSegments(outputStream, op.value)
                     outputStream.flush()
                     closeQuietly(outputStream)
                 }

app/src/main/java/com/stevesoltys/seedvault/transport/restore/FullRestore.kt

@@ -77,7 +77,7 @@ internal class FullRestore(
      */
     fun getNextFullRestoreDataChunk(socket: ParcelFileDescriptor): Int {
         Log.i(TAG, "Get next full restore data chunk.")
-        val state = this.state ?: throw IllegalStateException()
+        val state = this.state ?: throw IllegalStateException("no state")
         val packageName = state.packageInfo.packageName
 
         if (state.inputStream == null) {
@@ -103,9 +103,9 @@ internal class FullRestore(
     }
 
     private fun readInputStream(socket: ParcelFileDescriptor): Int = socket.use { fileDescriptor ->
-        val state = this.state ?: throw IllegalStateException()
+        val state = this.state ?: throw IllegalStateException("no state")
         val packageName = state.packageInfo.packageName
-        val inputStream = state.inputStream ?: throw IllegalStateException()
+        val inputStream = state.inputStream ?: throw IllegalStateException("no stream")
         val outputStream = outputFactory.getOutputStream(fileDescriptor)
 
         try {
@@ -144,7 +144,7 @@ internal class FullRestore(
      * with no further attempts to restore app data.
      */
     fun abortFullRestore(): Int {
-        val state = this.state ?: throw IllegalStateException()
+        val state = this.state ?: throw IllegalStateException("no state")
         Log.i(TAG, "Abort full restore of ${state.packageInfo.packageName}!")
 
         resetState()
@@ -156,7 +156,7 @@ internal class FullRestore(
      * freeing any resources and connections used during the restore process.
      */
     fun finishRestore() {
-        val state = this.state ?: throw IllegalStateException()
+        val state = this.state ?: throw IllegalStateException("no state")
         Log.i(TAG, "Finish restore of ${state.packageInfo.packageName}!")
 
         resetState()

app/src/main/java/com/stevesoltys/seedvault/transport/restore/KVRestore.kt

@@ -55,7 +55,7 @@ internal class KVRestore(
      * or [TRANSPORT_ERROR] (an error occurred, the restore should be aborted and rescheduled).
      */
     fun getRestoreData(data: ParcelFileDescriptor): Int {
-        val state = this.state ?: throw IllegalStateException()
+        val state = this.state ?: throw IllegalStateException("no state")
 
         // The restore set is the concatenation of the individual record blobs,
         // each of which is a file in the package's directory.
@@ -124,7 +124,7 @@ internal class KVRestore(
         try {
             val version = headerReader.readVersion(inputStream)
             crypto.decryptHeader(inputStream, version, state.packageInfo.packageName, dKey.key)
-            val value = crypto.decryptSegment(inputStream)
+            val value = crypto.decryptMultipleSegments(inputStream)
             val size = value.size
             Log.v(TAG, "    ... key=${dKey.key} size=$size")
 

app/src/main/java/com/stevesoltys/seedvault/transport/restore/RestoreCoordinator.kt

@@ -41,7 +41,9 @@ internal class RestoreCoordinator(
         val restoreSets = ArrayList<RestoreSet>()
         for (encryptedMetadata in availableBackups) {
             if (encryptedMetadata.error) continue
-            check(encryptedMetadata.inputStream != null)  // if there's no error, there must be a stream
+            check(encryptedMetadata.inputStream != null) {
+                "No error when getting encrypted metadata, but stream is still missing."
+            }
             try {
                 val metadata = metadataReader.readMetadata(encryptedMetadata.inputStream, encryptedMetadata.token)
                 val set = RestoreSet(metadata.deviceName, metadata.deviceName, metadata.token)
@@ -91,7 +93,7 @@ internal class RestoreCoordinator(
      * or [TRANSPORT_ERROR] (an error occurred, the restore should be aborted and rescheduled).
      */
     fun startRestore(token: Long, packages: Array<out PackageInfo>): Int {
-        check(state == null)
+        check(state == null) { "Started new restore with existing state" }
         Log.i(TAG, "Start restore with ${packages.map { info -> info.packageName }}")
         state = RestoreCoordinatorState(token, packages.iterator())
         return TRANSPORT_OK
@@ -125,7 +127,7 @@ internal class RestoreCoordinator(
      */
     fun nextRestorePackage(): RestoreDescription? {
         Log.i(TAG, "Next restore package!")
-        val state = this.state ?: throw IllegalStateException()
+        val state = this.state ?: throw IllegalStateException("no state")
 
         if (!state.packages.hasNext()) return NO_MORE_PACKAGES
         val packageInfo = state.packages.next()

app/src/main/java/com/stevesoltys/seedvault/ui/storage/StorageViewModel.kt

@@ -87,7 +87,7 @@ internal abstract class StorageViewModel(
      */
     protected fun saveStorage(uri: Uri): Boolean {
         // store backup storage location in settings
-        val root = storageRoot ?: throw IllegalStateException()
+        val root = storageRoot ?: throw IllegalStateException("no storage root")
         val name = if (root.isInternal()) {
             "${root.title} (${app.getString(R.string.settings_backup_location_internal)})"
         } else {

app/src/test/java/com/stevesoltys/seedvault/crypto/CryptoImplTest.kt

@@ -7,11 +7,13 @@ import com.stevesoltys.seedvault.header.MAX_SEGMENT_LENGTH
 import io.mockk.every
 import io.mockk.mockk
 import org.junit.jupiter.api.Assertions.assertArrayEquals
+import org.junit.jupiter.api.Assertions.assertThrows
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.TestInstance
 import org.junit.jupiter.api.TestInstance.Lifecycle.PER_METHOD
 import java.io.ByteArrayInputStream
 import java.io.ByteArrayOutputStream
+import java.io.IOException
 import javax.crypto.Cipher
 import kotlin.random.Random
 
@@ -50,4 +52,12 @@ class CryptoImplTest {
         assertArrayEquals(cleartext, crypto.decryptSegment(inputStream))
     }
 
+    @Test
+    fun `decrypting multiple segments on empty stream throws`() {
+        val inputStream = ByteArrayInputStream(ByteArray(0))
+        assertThrows(IOException::class.java) {
+            crypto.decryptMultipleSegments(inputStream)
+        }
+    }
+
 }

app/src/test/java/com/stevesoltys/seedvault/crypto/CryptoIntegrationTest.kt

@@ -2,12 +2,16 @@ package com.stevesoltys.seedvault.crypto
 
 import com.stevesoltys.seedvault.header.HeaderReaderImpl
 import com.stevesoltys.seedvault.header.HeaderWriterImpl
+import com.stevesoltys.seedvault.header.MAX_SEGMENT_CLEARTEXT_LENGTH
+import com.stevesoltys.seedvault.header.MAX_SEGMENT_LENGTH
 import org.junit.jupiter.api.Assertions.assertArrayEquals
+import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.TestInstance
 import org.junit.jupiter.api.TestInstance.Lifecycle.PER_METHOD
 import java.io.ByteArrayInputStream
 import java.io.ByteArrayOutputStream
+import kotlin.random.Random
 
 @TestInstance(PER_METHOD)
 class CryptoIntegrationTest {
@@ -41,4 +45,24 @@ class CryptoIntegrationTest {
         assertArrayEquals(cleartext, crypto.decryptSegment(inputStream))
     }
 
+    @Test
+    fun `multiple segments get encrypted and decrypted as expected`() {
+        val size = Random.nextInt(5) * MAX_SEGMENT_CLEARTEXT_LENGTH + Random.nextInt(0, 1337)
+        val cleartext = ByteArray(size).apply { Random.nextBytes(this) }
+
+        crypto.encryptMultipleSegments(outputStream, cleartext)
+        val inputStream = ByteArrayInputStream(outputStream.toByteArray())
+        assertArrayEquals(cleartext, crypto.decryptMultipleSegments(inputStream))
+    }
+
+    @Test
+    fun `test maximum lengths`() {
+        val cipher = cipherFactory.createEncryptionCipher()
+        val expectedDiff = MAX_SEGMENT_LENGTH - MAX_SEGMENT_CLEARTEXT_LENGTH
+        for (i in 1..(3 * MAX_SEGMENT_LENGTH + 42)) {
+            val outputSize = cipher.getOutputSize(i)
+            assertEquals(expectedDiff, outputSize - i)
+        }
+    }
+
 }

app/src/test/java/com/stevesoltys/seedvault/transport/CoordinatorIntegrationTest.kt

@@ -14,6 +14,7 @@ import com.stevesoltys.seedvault.crypto.KeyManagerTestImpl
 import com.stevesoltys.seedvault.encodeBase64
 import com.stevesoltys.seedvault.header.HeaderReaderImpl
 import com.stevesoltys.seedvault.header.HeaderWriterImpl
+import com.stevesoltys.seedvault.header.MAX_SEGMENT_CLEARTEXT_LENGTH
 import com.stevesoltys.seedvault.metadata.MetadataReaderImpl
 import com.stevesoltys.seedvault.metadata.MetadataWriterImpl
 import com.stevesoltys.seedvault.transport.backup.*
@@ -123,6 +124,53 @@ internal class CoordinatorIntegrationTest : TransportTest() {
         assertEquals(TRANSPORT_OK, restore.getRestoreData(fileDescriptor))
     }
 
+    @Test
+    fun `test key-value backup with huge value`() {
+        val value = CapturingSlot<ByteArray>()
+        val size = Random.nextInt(5) * MAX_SEGMENT_CLEARTEXT_LENGTH + Random.nextInt(0, 1337)
+        val appData = ByteArray(size).apply { Random.nextBytes(this) }
+        val bOutputStream = ByteArrayOutputStream()
+
+        // read one key/value record and write it to output stream
+        every { kvBackupPlugin.hasDataForPackage(packageInfo) } returns false
+        every { kvBackupPlugin.ensureRecordStorageForPackage(packageInfo) } just Runs
+        every { inputFactory.getBackupDataInput(fileDescriptor) } returns backupDataInput
+        every { backupDataInput.readNextHeader() } returns true andThen false
+        every { backupDataInput.key } returns key
+        every { backupDataInput.dataSize } returns appData.size
+        every { backupDataInput.readEntityData(capture(value), 0, appData.size) } answers {
+            appData.copyInto(value.captured) // write the app data into the passed ByteArray
+            appData.size
+        }
+        every { kvBackupPlugin.getOutputStreamForRecord(packageInfo, key64) } returns bOutputStream
+        every { settingsManager.saveNewBackupTime() } just Runs
+
+        // start and finish K/V backup
+        assertEquals(TRANSPORT_OK, backup.performIncrementalBackup(packageInfo, fileDescriptor, 0))
+        assertEquals(TRANSPORT_OK, backup.finishBackup())
+
+        // start restore
+        assertEquals(TRANSPORT_OK, restore.startRestore(token, arrayOf(packageInfo)))
+
+        // find data for K/V backup
+        every { kvRestorePlugin.hasDataForPackage(token, packageInfo) } returns true
+
+        val restoreDescription = restore.nextRestorePackage() ?: fail()
+        assertEquals(packageInfo.packageName, restoreDescription.packageName)
+        assertEquals(RestoreDescription.TYPE_KEY_VALUE, restoreDescription.dataType)
+
+        // restore finds the backed up key and writes the decrypted value
+        val backupDataOutput = mockk<BackupDataOutput>()
+        val rInputStream = ByteArrayInputStream(bOutputStream.toByteArray())
+        every { kvRestorePlugin.listRecords(token, packageInfo) } returns listOf(key64)
+        every { outputFactory.getBackupDataOutput(fileDescriptor) } returns backupDataOutput
+        every { kvRestorePlugin.getInputStreamForRecord(token, packageInfo, key64) } returns rInputStream
+        every { backupDataOutput.writeEntityHeader(key, appData.size) } returns 1137
+        every { backupDataOutput.writeEntityData(appData, appData.size) } returns appData.size
+
+        assertEquals(TRANSPORT_OK, restore.getRestoreData(fileDescriptor))
+    }
+
     @Test
     fun `test full backup and restore with two chunks`() {
         // return streams from plugin and app data

app/src/test/java/com/stevesoltys/seedvault/transport/backup/KVBackupTest.kt

@@ -142,7 +142,7 @@ internal class KVBackupTest : BackupTest() {
         writeHeaderAndEncrypt()
         every { plugin.getOutputStreamForRecord(packageInfo, key64) } returns outputStream
         every { headerWriter.writeVersion(outputStream, versionHeader) } just Runs
-        every { crypto.encryptSegment(outputStream, any()) } throws IOException()
+        every { crypto.encryptMultipleSegments(outputStream, any()) } throws IOException()
 
         assertEquals(TRANSPORT_ERROR, backup.performBackup(packageInfo, data, 0))
         assertFalse(backup.hasState())
@@ -205,7 +205,7 @@ internal class KVBackupTest : BackupTest() {
         every { plugin.getOutputStreamForRecord(packageInfo, key64) } returns outputStream
         every { headerWriter.writeVersion(outputStream, versionHeader) } just Runs
         every { crypto.encryptHeader(outputStream, versionHeader) } just Runs
-        every { crypto.encryptSegment(outputStream, any()) } just Runs
+        every { crypto.encryptMultipleSegments(outputStream, any()) } just Runs
     }
 
 }

app/src/test/java/com/stevesoltys/seedvault/transport/restore/KVRestoreTest.kt

@@ -88,7 +88,7 @@ internal class KVRestoreTest : RestoreTest() {
         every { plugin.getInputStreamForRecord(token, packageInfo, key64) } returns inputStream
         every { headerReader.readVersion(inputStream) } returns VERSION
         every { crypto.decryptHeader(inputStream, VERSION, packageInfo.packageName, key) } returns versionHeader
-        every { crypto.decryptSegment(inputStream) } throws IOException()
+        every { crypto.decryptMultipleSegments(inputStream) } throws IOException()
         streamsGetClosed()
 
         assertEquals(TRANSPORT_ERROR, restore.getRestoreData(fileDescriptor))
@@ -131,7 +131,7 @@ internal class KVRestoreTest : RestoreTest() {
         every { plugin.getInputStreamForRecord(token, packageInfo, key64) } returns inputStream
         every { headerReader.readVersion(inputStream) } returns VERSION
         every { crypto.decryptHeader(inputStream, VERSION, packageInfo.packageName, key) } returns versionHeader
-        every { crypto.decryptSegment(inputStream) } returns data
+        every { crypto.decryptMultipleSegments(inputStream) } returns data
         every { output.writeEntityHeader(key, data.size) } throws IOException()
         streamsGetClosed()
 
@@ -147,7 +147,7 @@ internal class KVRestoreTest : RestoreTest() {
         every { plugin.getInputStreamForRecord(token, packageInfo, key64) } returns inputStream
         every { headerReader.readVersion(inputStream) } returns VERSION
         every { crypto.decryptHeader(inputStream, VERSION, packageInfo.packageName, key) } returns versionHeader
-        every { crypto.decryptSegment(inputStream) } returns data
+        every { crypto.decryptMultipleSegments(inputStream) } returns data
         every { output.writeEntityHeader(key, data.size) } returns 42
         every { output.writeEntityData(data, data.size) } throws IOException()
         streamsGetClosed()
@@ -164,7 +164,7 @@ internal class KVRestoreTest : RestoreTest() {
         every { plugin.getInputStreamForRecord(token, packageInfo, key64) } returns inputStream
         every { headerReader.readVersion(inputStream) } returns VERSION
         every { crypto.decryptHeader(inputStream, VERSION, packageInfo.packageName, key) } returns versionHeader
-        every { crypto.decryptSegment(inputStream) } returns data
+        every { crypto.decryptMultipleSegments(inputStream) } returns data
         every { output.writeEntityHeader(key, data.size) } returns 42
         every { output.writeEntityData(data, data.size) } returns data.size
         streamsGetClosed()
@@ -184,14 +184,14 @@ internal class KVRestoreTest : RestoreTest() {
         every { plugin.getInputStreamForRecord(token, packageInfo, key64) } returns inputStream
         every { headerReader.readVersion(inputStream) } returns VERSION
         every { crypto.decryptHeader(inputStream, VERSION, packageInfo.packageName, key) } returns versionHeader
-        every { crypto.decryptSegment(inputStream) } returns data
+        every { crypto.decryptMultipleSegments(inputStream) } returns data
         every { output.writeEntityHeader(key, data.size) } returns 42
         every { output.writeEntityData(data, data.size) } returns data.size
         // second key/value
         every { plugin.getInputStreamForRecord(token, packageInfo, key264) } returns inputStream2
         every { headerReader.readVersion(inputStream2) } returns VERSION
         every { crypto.decryptHeader(inputStream2, VERSION, packageInfo.packageName, key2) } returns versionHeader2
-        every { crypto.decryptSegment(inputStream2) } returns data2
+        every { crypto.decryptMultipleSegments(inputStream2) } returns data2
         every { output.writeEntityHeader(key2, data2.size) } returns 42
         every { output.writeEntityData(data2, data2.size) } returns data2.size
         every { inputStream2.close() } just Runs