From 7d52d74bbe39376f331887759020cff097dace55 Mon Sep 17 00:00:00 2001
From: K900 <me@0upti.me>
Date: Thu, 20 Jun 2024 08:54:43 +0300
Subject: [PATCH] BrotliDecompressionSource: don't bail out too early

If we've consumed the entire input, that doesn't actually mean we're
done decompressing - there might be more output left. This worked (?)
in most cases because the input and output sizes are pretty comparable,
but sometimes they're not and then things get very funny.

Change-Id: I73435a654a911b8ce25119f713b80706c5783c1b
---
 src/libutil/compression.cc        | 27 ++++++++++++++-------------
 tests/unit/libutil/compression.cc | 10 ++++++++++
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/libutil/compression.cc b/src/libutil/compression.cc
index d93a1b1d6..a66069e52 100644
--- a/src/libutil/compression.cc
+++ b/src/libutil/compression.cc
@@ -163,23 +163,24 @@ struct BrotliDecompressionSource : Source
         uint8_t * out = (uint8_t *) data;
         const auto * begin = out;
 
-        try {
-            while (len && !BrotliDecoderIsFinished(state.get())) {
-                checkInterrupt();
+        while (len && !BrotliDecoderIsFinished(state.get())) {
+            checkInterrupt();
 
-                while (avail_in == 0) {
+            while (avail_in == 0) {
+                try {
                     avail_in = inner->read(buf.get(), BUF_SIZE);
-                    next_in = (const uint8_t *) buf.get();
-                }
-
-                if (!BrotliDecoderDecompressStream(
-                        state.get(), &avail_in, &next_in, &len, &out, nullptr
-                    ))
-                {
-                    throw CompressionError("error while decompressing brotli file");
+                } catch (EndOfFile &) {
+                    break;
                 }
+                next_in = (const uint8_t *) buf.get();
+            }
+
+            if (!BrotliDecoderDecompressStream(
+                    state.get(), &avail_in, &next_in, &len, &out, nullptr
+                ))
+            {
+                throw CompressionError("error while decompressing brotli file");
             }
-        } catch (EndOfFile &) {
         }
 
         if (begin != out) {
diff --git a/tests/unit/libutil/compression.cc b/tests/unit/libutil/compression.cc
index 0542e7d33..3b40db0cd 100644
--- a/tests/unit/libutil/compression.cc
+++ b/tests/unit/libutil/compression.cc
@@ -66,6 +66,16 @@ namespace nix {
         ASSERT_THROW(decompress(method, str), CompressionError);
     }
 
+    TEST(decompress, veryLongBrotli) {
+        auto method = "br";
+        auto str = std::string(65536, 'a');
+        auto o = decompress(method, compress(method, str));
+
+        // This is just to not print 64k of "a" for most failures
+        ASSERT_EQ(o.length(), str.length());
+        ASSERT_EQ(o, str);
+    }
+
     /* ----------------------------------------------------------------------------
      * compression sinks
      * --------------------------------------------------------------------------*/