PatchCompressor: write the decompressed size into the zstd frame header

generator/common/src/jvmMain/kotlin/io/github/kdroidfilter/seforimlibrary/common/patch/PatchCompressor.kt

@@ -1,9 +1,14 @@
 package io.github.kdroidfilter.seforimlibrary.common.patch
 
+import com.github.luben.zstd.EndDirective
+import com.github.luben.zstd.ZstdCompressCtx
 import com.github.luben.zstd.ZstdOutputStream
+import java.nio.ByteBuffer
+import java.nio.channels.FileChannel
 import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.StandardCopyOption
+import java.nio.file.StandardOpenOption
 import java.security.MessageDigest
 
 /**
@@ -12,6 +17,16 @@ import java.security.MessageDigest
  *
  * Output sits next to the input (`<patch>.zst`). Caller decides whether
  * to keep or delete the original `.db`.
+ *
+ * The frame header **must** carry the decompressed content size. The Otzaria
+ * client decompresses through the `zstandard` FFI plugin, which allocates the
+ * destination buffer from `ZSTD_getFrameContentSize`; when the size is
+ * missing it falls back to `compressedSize * 20` in a single contiguous
+ * allocation. A ~560 MB patch compressed by the old `ZstdOutputStream` path
+ * (which never pledges a size) therefore asked for ~11 GB and failed with
+ * `Could not allocate 11694234840 bytes`. Streaming via [ZstdCompressCtx] with
+ * [ZstdCompressCtx.setPledgedSrcSize] keeps memory bounded here and writes the
+ * exact size into the header.
  */
 object PatchCompressor {
 
@@ -28,6 +43,9 @@ object PatchCompressor {
         val uncompressedSize: Long,
     )
 
+    /** 1 MiB input chunks; output buffer sized per zstd's recommendation. */
+    private const val IN_CHUNK = 1 shl 20
+
     /**
      * Compresses [patchDb] to `<patchDb>.zst` at level [level].
      *
@@ -43,27 +61,57 @@ object PatchCompressor {
         val uncompressedSha = MessageDigest.getInstance("SHA-256")
         val uncompressedSize = Files.size(patchDb)
 
-        Files.newInputStream(patchDb).use { input ->
-            Files.newOutputStream(
-                target,
-                java.nio.file.StandardOpenOption.CREATE,
-                java.nio.file.StandardOpenOption.TRUNCATE_EXISTING,
-                java.nio.file.StandardOpenOption.WRITE,
-            ).use { rawOut ->
-                ZstdOutputStream(rawOut, level).use { zstd ->
-                    // Enable multithreaded compression when workers > 1.
-                    if (workers > 1) zstd.setWorkers(workers)
-                    val buf = ByteArray(1 shl 20) // 1 MiB
+        ZstdCompressCtx().use { ctx ->
+            ctx.setLevel(level)
+            // Enable multithreaded compression when workers > 1.
+            if (workers > 1) ctx.setWorkers(workers)
+            // Written into the frame header -> the client can size its buffer exactly.
+            ctx.setPledgedSrcSize(uncompressedSize)
+
+            val inBuf = ByteBuffer.allocateDirect(IN_CHUNK)
+            val outBuf = ByteBuffer.allocateDirect(ZstdOutputStream.recommendedCOutSize().toInt())
+            val hashView = ByteArray(IN_CHUNK)
+
+            FileChannel.open(patchDb, StandardOpenOption.READ).use { input ->
+                FileChannel.open(
+                    target,
+                    StandardOpenOption.CREATE,
+                    StandardOpenOption.TRUNCATE_EXISTING,
+                    StandardOpenOption.WRITE,
+                ).use { output ->
+                    var remaining = uncompressedSize
                     while (true) {
-                        val n = input.read(buf)
+                        inBuf.clear()
+                        val n = input.read(inBuf)
                         if (n <= 0) break
-                        uncompressedSha.update(buf, 0, n)
-                        zstd.write(buf, 0, n)
+                        remaining -= n
+                        check(remaining >= 0L) {
+                            "patch.db grew while compressing: expected $uncompressedSize bytes"
+                        }
+                        inBuf.flip()
+                        inBuf.duplicate().get(hashView, 0, n)
+                        uncompressedSha.update(hashView, 0, n)
+
+                        // END on the last chunk finalises the frame in the same call,
+                        // so the pledged size is honoured exactly.
+                        val directive = if (remaining == 0L) EndDirective.END else EndDirective.CONTINUE
+                        drain(ctx, inBuf, outBuf, output, directive)
+                    }
+                    check(remaining == 0L) {
+                        "patch.db changed while compressing: expected $uncompressedSize bytes, " +
+                            "read ${uncompressedSize - remaining}"
+                    }
+                    if (uncompressedSize == 0L) {
+                        // Empty patch: still emit a well-formed frame with size 0.
+                        inBuf.clear().flip()
+                        drain(ctx, inBuf, outBuf, output, EndDirective.END)
                     }
                 }
             }
         }
 
+        check(Files.size(target) > 0) { "zstd produced an empty file for $patchDb" }
+
         val compressedSha = sha256(target)
         return Result(
             compressedFile = target,
@@ -74,6 +122,31 @@ object PatchCompressor {
         )
     }
 
+    /**
+     * Feeds [inBuf] to [ctx] until it is fully consumed, writing every produced
+     * block to [output]. With [EndDirective.END] it also loops until zstd reports
+     * the frame is complete (`compressDirectByteBufferStream` returns `true`).
+     */
+    private fun drain(
+        ctx: ZstdCompressCtx,
+        inBuf: ByteBuffer,
+        outBuf: ByteBuffer,
+        output: FileChannel,
+        directive: EndDirective,
+    ) {
+        while (true) {
+            outBuf.clear()
+            val done = ctx.compressDirectByteBufferStream(outBuf, inBuf, directive)
+            outBuf.flip()
+            while (outBuf.hasRemaining()) output.write(outBuf)
+            if (directive == EndDirective.END) {
+                if (done) return
+            } else if (!inBuf.hasRemaining()) {
+                return
+            }
+        }
+    }
+
     private fun sha256(path: Path): ByteArray {
         val md = MessageDigest.getInstance("SHA-256")
         Files.newInputStream(path).use { stream ->

generator/common/src/jvmTest/kotlin/io/github/kdroidfilter/seforimlibrary/common/patch/PatchCompressorTest.kt

@@ -0,0 +1,80 @@
+package io.github.kdroidfilter.seforimlibrary.common.patch
+
+import com.github.luben.zstd.Zstd
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TemporaryFolder
+import java.nio.file.Files
+import java.security.MessageDigest
+import kotlin.random.Random
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+
+class PatchCompressorTest {
+    @JvmField @Rule
+    val tmp = TemporaryFolder()
+
+    private fun sha256Hex(bytes: ByteArray): String =
+        MessageDigest.getInstance("SHA-256").digest(bytes).joinToString("") { "%02x".format(it) }
+
+    /** Multi-chunk input (> 1 MiB) so the streaming loop crosses chunk boundaries. */
+    private fun sampleBytes(size: Int): ByteArray {
+        val rnd = Random(42)
+        // Mostly repetitive (compressible) with random noise sprinkled in.
+        return ByteArray(size) { i -> if (i % 7 == 0) rnd.nextInt().toByte() else (i % 13).toByte() }
+    }
+
+    @Test
+    fun `frame header carries the exact decompressed size`() {
+        val payload = sampleBytes(3 * (1 shl 20) + 12345)
+        val patch = tmp.newFile("patch.db").toPath()
+        Files.write(patch, payload)
+
+        val result = PatchCompressor.compress(patch, level = 3, workers = 1)
+        val compressed = Files.readAllBytes(result.compressedFile)
+
+        // The Otzaria client sizes its buffer from this field; without it, the
+        // fallback is compressedSize*20 in one allocation (~11GB for a 560MB patch).
+        assertEquals(payload.size.toLong(), Zstd.getFrameContentSize(compressed))
+        assertEquals(payload.size.toLong(), result.uncompressedSize)
+        assertEquals(compressed.size.toLong(), result.compressedSize)
+    }
+
+    @Test
+    fun `round-trips and reports matching hashes`() {
+        val payload = sampleBytes(2 * (1 shl 20) + 1)
+        val patch = tmp.newFile("patch.db").toPath()
+        Files.write(patch, payload)
+
+        val result = PatchCompressor.compress(patch, level = 3, workers = 1)
+        val compressed = Files.readAllBytes(result.compressedFile)
+        val restored = Zstd.decompress(compressed, payload.size)
+
+        assertContentEquals(payload, restored)
+        assertEquals(sha256Hex(payload), result.uncompressedSha256)
+        assertEquals(sha256Hex(compressed), result.compressedSha256)
+    }
+
+    @Test
+    fun `multithreaded compression also pledges the size`() {
+        val payload = sampleBytes(4 * (1 shl 20))
+        val patch = tmp.newFile("patch.db").toPath()
+        Files.write(patch, payload)
+
+        val result = PatchCompressor.compress(patch, level = 3, workers = 2)
+        val compressed = Files.readAllBytes(result.compressedFile)
+
+        assertEquals(payload.size.toLong(), Zstd.getFrameContentSize(compressed))
+        assertContentEquals(payload, Zstd.decompress(compressed, payload.size))
+    }
+
+    @Test
+    fun `empty patch yields a valid frame with size zero`() {
+        val patch = tmp.newFile("patch.db").toPath()
+        val result = PatchCompressor.compress(patch, level = 3, workers = 1)
+        val compressed = Files.readAllBytes(result.compressedFile)
+
+        assertEquals(0L, Zstd.getFrameContentSize(compressed))
+        assertEquals(0, Zstd.decompress(compressed, 0).size)
+    }
+}