From 4bd6e3daba6e798e5ebd2ff00ed9744deb4c36e7 Mon Sep 17 00:00:00 2001
From: Yorhel <git@yorhel.nl>
Date: Wed, 13 Nov 2024 15:28:07 +0100
Subject: [PATCH] Fix handling short reads when uncompressed json

Turns out that zstd can consume compressed data without returning any
decompressed data when the input buffer isn't full enough. I just
increased the input buffer as a workaround.
---
 src/json_import.zig | 20 ++++++++++++--------
 src/main.zig        |  1 -
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/json_import.zig b/src/json_import.zig
index d119952..6cb627a 100644
--- a/src/json_import.zig
+++ b/src/json_import.zig
@@ -13,9 +13,9 @@ const c = @import("c.zig").c;
 const ZstdReader = struct {
     ctx: ?*c.ZSTD_DStream,
     in: c.ZSTD_inBuffer,
-    // Assumption: ZSTD_decompressStream() can always make progress as long as
-    // 1/8 of this buffer is filled.
-    inbuf: [c.ZSTD_BLOCKSIZE_MAX + 16]u8, // This ZSTD_DStreamInSize() + a little bit extra
+    // Assumption: ZSTD_decompressStream() can always output uncompressed data
+    // as long as half of this buffer is filled.
+    inbuf: [c.ZSTD_BLOCKSIZE_MAX * 2 + 1024]u8, // This twice ZSTD_DStreamInSize() + a little bit extra
 
     fn create(head: []const u8) *ZstdReader {
         const r = main.allocator.create(ZstdReader) catch unreachable;
@@ -39,11 +39,15 @@ const ZstdReader = struct {
     }
 
     fn read(r: *ZstdReader, f: std.fs.File, out: []u8) !usize {
-        if (r.in.size - r.in.pos < r.inbuf.len / 8) {
-            std.mem.copyForwards(u8, &r.inbuf, r.inbuf[r.in.pos..][0..r.in.size - r.in.pos]);
-            r.in.size -= r.in.pos;
-            r.in.pos = 0;
-            r.in.size += try f.read(r.inbuf[r.in.size..]);
+        while (r.in.size - r.in.pos < r.inbuf.len / 2) {
+            if (r.in.pos > 0) {
+                std.mem.copyForwards(u8, &r.inbuf, r.inbuf[r.in.pos..][0..r.in.size - r.in.pos]);
+                r.in.size -= r.in.pos;
+                r.in.pos = 0;
+            }
+            const rd = try f.read(r.inbuf[r.in.size..]);
+            r.in.size += rd;
+            if (rd == 0) break;
         }
         var arg = c.ZSTD_outBuffer{
             .dst = out.ptr,
diff --git a/src/main.zig b/src/main.zig
index 4db4ec8..440ec91 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -461,7 +461,6 @@ fn readImport(path: [:0]const u8) !void {
         else try std.fs.cwd().openFileZ(path, .{});
     errdefer fd.close();
 
-    // TODO: While we're at it, recognize and handle compressed JSON
     var buf: [8]u8 = undefined;
     try fd.reader().readNoEof(&buf);
     if (std.mem.eql(u8, &buf, bin_export.SIGNATURE)) {