Add test for BlockDeviceOutputStream and fix bugs found testing

app/src/main/java/eu/depau/etchdroid/utils/blockdevice/BlockDeviceOutputStream.kt

@@ -7,7 +7,7 @@ import java.nio.ByteBuffer
 
 class BlockDeviceOutputStream(
         private val blockDev: BlockDeviceDriver,
-        bufferBlocks: Int = 2048
+        private val bufferBlocks: Int = 2048
 ) : OutputStream() {
 
     private val byteBuffer = ByteBuffer.allocate(blockDev.blockSize * bufferBlocks)
@@ -16,15 +16,14 @@ class BlockDeviceOutputStream(
     private val currentByteOffset: Long
         get() = currentBlockOffset * blockDev.blockSize + byteBuffer.position()
 
-    private val sizeBytes: Long
-        get() = blockDev.size.toLong() * blockDev.blockSize
-
     private val bytesUntilEOF: Long
-        get() = blockDev.size.toLong() * blockDev.size - currentByteOffset
+        get() = blockDev.size.toLong() * blockDev.blockSize - currentByteOffset
 
     override fun write(b: Int) {
-        if (bytesUntilEOF < 1)
+        if (bytesUntilEOF < 1) {
+            flush()
             throw IOException("No space left on device")
+        }
 
         byteBuffer.put(b.toByte())
 
@@ -37,7 +36,7 @@ class BlockDeviceOutputStream(
     }
 
     override fun write(b: ByteArray, off: Int, len: Int) {
-        val maxPos = (off + len) % (b.size + 1)
+        val maxPos = Math.min(off + len, b.size)
 
         if (len <= 0 || off > b.size)
             return
@@ -61,12 +60,16 @@ class BlockDeviceOutputStream(
             blockDev.read(currentBlockOffset + fullBlocks, incompleteBlockBuffer)
 
             // Add it to the incomplete block
-            byteBuffer.limit(fullBlocks * blockDev.blockSize)
+            byteBuffer.apply {
-            byteBuffer.put(
+                position(toWrite)
-                    incompleteBlockBuffer.array(),
+                limit((fullBlocks + 1) * blockDev.blockSize)
-                    toWrite, blockDev.blockSize - incompleteBlockFullBytes
+                put(
-            )
+                        incompleteBlockBuffer.array(),
-            byteBuffer.position(0)
+                        incompleteBlockFullBytes,
+                        blockDev.blockSize - incompleteBlockFullBytes
+                )
+                position(0)
+            }
         }
 
         // Flush to device
@@ -80,6 +83,13 @@ class BlockDeviceOutputStream(
             clear()
             position(incompleteBlockFullBytes)
         }
-        currentBlockOffset += fullBlocks;
+
+        // Ensure the buffer is limited on EOF
+        if (blockDev.size - currentBlockOffset < bufferBlocks)
+            byteBuffer.limit(
+                    (blockDev.size - currentBlockOffset).toInt() * blockDev.blockSize
+            )
+
+        currentBlockOffset += fullBlocks
     }
 }

app/src/test/java/eu/depau/etchdroid/utils/blockdevice/BlockDeviceOutputStreamTest.kt

@@ -0,0 +1,167 @@
+package eu.depau.etchdroid.utils.blockdevice
+
+import com.github.mjdev.libaums.driver.BlockDeviceDriver
+import com.github.mjdev.libaums.driver.file.FileBlockDeviceDriver
+import org.junit.Assert
+import org.junit.Test
+import java.io.File
+import java.io.IOException
+import java.io.RandomAccessFile
+
+class BlockDeviceOutputStreamTest {
+    @Test
+    fun testWithCommonParams() {
+        runWriteTest(
+                10L * 1024 * 1024,  // 10 MiB
+                512,
+                2048
+        )
+    }
+
+    @Test
+    fun testWithWeirdBlockSize() {
+        runWriteTest(
+                10L * 666 * 2 * 666 * 2,  // 10 MegaDevils
+                666,
+                2048
+        )
+    }
+
+    @Test
+    fun testUnbuffered() {
+        runWriteTest(
+                10L * 1024 * 1024,  // 10 MiB
+                512,
+                1
+        )
+    }
+
+    fun runWriteTest(testDevSize: Long, testBlockSize: Int, testBufferBlocks: Int) {
+        val testFile = createTestFile(testDevSize)
+
+        try {
+            val testDev = createTestBlockDevice(testFile, testBlockSize)
+            val outputStream = BlockDeviceOutputStream(testDev, bufferBlocks = testBufferBlocks)
+
+            val deviceRandomAccessFile: RandomAccessFile = testDev::class.java
+                    .getDeclaredField("file")
+                    .apply { isAccessible = true }
+                    .get(testDev) as RandomAccessFile
+            var byteArray: ByteArray
+
+
+            // Write some bytes
+            outputStream.write(0)
+            outputStream.write(1)
+            outputStream.write(2)
+            outputStream.write(3)
+            outputStream.flush()
+            deviceRandomAccessFile.fd.sync()
+
+            byteArray = ByteArray(8)
+            RandomAccessFile(testFile, "r").use {
+                it.read(byteArray)
+            }
+
+            Assert.assertArrayEquals(
+                    byteArrayOf(
+                            0, 1, 2, 3, 0xFA.toByte(), 0xF9.toByte(), 0xF8.toByte(), 0xF7.toByte()
+                    ),
+                    byteArray
+            )
+
+            // Fill current block
+            byteArray = ByteArray(testBlockSize - 4)
+            outputStream.write(byteArray)
+            outputStream.flush()
+            deviceRandomAccessFile.fd.sync()
+
+            byteArray = ByteArray(testBlockSize + 8)
+            RandomAccessFile(testFile, "r").use {
+                it.read(byteArray)
+            }
+
+            Assert.assertArrayEquals(
+                    byteArrayOf(
+                            0, 1, 2, 3
+                    ),
+                    byteArray.copyOfRange(0, 4)
+            )
+            Assert.assertArrayEquals(
+                    (4 until testBlockSize).map { 0.toByte() }.toByteArray(),
+                    byteArray.copyOfRange(4, testBlockSize)
+            )
+            Assert.assertArrayEquals(
+                    (testBlockSize until testBlockSize + 8).map { (0xFE - it % 0xFF).toByte() }.toByteArray(),
+                    byteArray.copyOfRange(testBlockSize, testBlockSize + 8)
+            )
+
+            // Fill the buffer except for the last 8 bytes
+            // Note that flush successfully wrote one full block, so the block offset is now 1
+            byteArray = ByteArray(testBufferBlocks * testBlockSize - 8)
+            outputStream.write(byteArray)
+
+            // Now do a write that goes out of the buffer
+            byteArray = (0 until 16).map { 42.toByte() }.toByteArray()
+            outputStream.write(byteArray)
+            outputStream.flush()
+            deviceRandomAccessFile.fd.sync()
+
+            val fullBufferOffset = (testBufferBlocks + 1) * testBlockSize
+
+            val byteArray1 = ByteArray(20)
+            RandomAccessFile(testFile, "r").use {
+                it.seek(fullBufferOffset - 8L)
+                it.read(byteArray1)
+            }
+
+            Assert.assertArrayEquals(
+                    byteArray,
+                    byteArray1.copyOfRange(0, 16)
+            )
+            Assert.assertArrayEquals(
+                    (fullBufferOffset + 8 until fullBufferOffset + 12)
+                            .map { (0xFE - it % 0xFF).toByte() }
+                            .toByteArray(),
+                    byteArray1.copyOfRange(16, 20)
+            )
+
+            // Go to end of file - 4 bytes
+            val remainingBytes = testDevSize - (fullBufferOffset + 8)
+            byteArray = ByteArray(remainingBytes.toInt() - 4)
+            outputStream.write(byteArray)
+
+            // This should work
+            byteArray = ByteArray(4)
+            outputStream.write(byteArray)
+
+            // This should not
+            try {
+                outputStream.write(0)
+                Assert.fail("Did not throw exception on EOF")
+            } catch (e: IOException) {
+                Assert.assertEquals("No space left on device", e.message)
+            }
+
+        } finally {
+            testFile.delete()
+        }
+    }
+
+    companion object {
+        fun createTestFile(size: Long): File {
+            val file = File.createTempFile("etchdroid-test-", ".img")
+            val fileOutputStream = file.outputStream()
+
+            // Writes 0xFE, 0xFD, ..., 0x00
+            for (i in 0..size)
+                fileOutputStream.write(0xFE - i.toInt() % 0xFF)
+            fileOutputStream.flush()
+            file.deleteOnExit()
+            return file
+        }
+
+        fun createTestBlockDevice(file: File, blockSize: Int): BlockDeviceDriver =
+                FileBlockDeviceDriver(file, blockSize, 0)
+    }
+}