4 files changed, 55 insertions, 6 deletions

diff --git a/src/libstore/filetransfer.cc b/src/libstore/filetransfer.cc
index f3e8a5312..10c810e49 100644
--- a/src/libstore/filetransfer.cc
+++ b/src/libstore/filetransfer.cc
@@ -337,7 +337,7 @@ struct curlFileTransfer : public FileTransfer
             // wrapping user `callback`s instead is not possible because the
             // Callback api expects std::functions, and copying Callbacks is
             // not possible due the promises they hold.
-            if (code == CURLE_OK && !dataCallback) {
+            if (code == CURLE_OK && !dataCallback && result.data.length() > 0) {
                 result.data = decompress(encoding, result.data);
             }
 
diff --git a/src/libutil/compression.cc b/src/libutil/compression.cc
index 5152a2146..51c820d55 100644
--- a/src/libutil/compression.cc
+++ b/src/libutil/compression.cc
@@ -144,6 +144,7 @@ struct BrotliDecompressionSource : Source
     std::unique_ptr<char[]> buf;
     size_t avail_in = 0;
     const uint8_t * next_in;
+    std::exception_ptr inputEofException = nullptr;
 
     Source * inner;
     std::unique_ptr<BrotliDecoderState, void (*)(BrotliDecoderState *)> state;
@@ -167,23 +168,42 @@ struct BrotliDecompressionSource : Source
         while (len && !BrotliDecoderIsFinished(state.get())) {
             checkInterrupt();
 
-            while (avail_in == 0) {
+            while (avail_in == 0 && inputEofException == nullptr) {
                 try {
                     avail_in = inner->read(buf.get(), BUF_SIZE);
                 } catch (EndOfFile &) {
+                    // No more data, but brotli may still have output remaining
+                    // from the last call.
+                    inputEofException = std::current_exception();
                     break;
                 }
                 next_in = charptr_cast<const uint8_t *>(buf.get());
             }
 
-            if (!BrotliDecoderDecompressStream(
-                    state.get(), &avail_in, &next_in, &len, &out, nullptr
-                ))
-            {
+            BrotliDecoderResult res = BrotliDecoderDecompressStream(
+                state.get(), &avail_in, &next_in, &len, &out, nullptr
+            );
+
+            switch (res) {
+            case BROTLI_DECODER_RESULT_SUCCESS:
+                // We're done here!
+                goto finish;
+            case BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT:
+                // Grab more input. Don't try if we already have exhausted our input stream.
+                if (inputEofException != nullptr) {
+                    std::rethrow_exception(inputEofException);
+                } else {
+                    continue;
+                }
+            case BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT:
+                // Need more output space: we can only get another buffer by someone calling us again, so get out.
+                goto finish;
+            case BROTLI_DECODER_RESULT_ERROR:
                 throw CompressionError("error while decompressing brotli file");
             }
         }
 
+finish:
         if (begin != out) {
             return out - begin;
         } else {
diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index 612658b2d..3a9685e0e 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -77,6 +77,11 @@ struct Source
      * Store up to ‘len’ in the buffer pointed to by ‘data’, and
      * return the number of bytes stored.  It blocks until at least
      * one byte is available.
+     *
+     * Should not return 0 (generally you want to throw EndOfFile), but nothing
+     * stops that.
+     *
+     * \throws EndOfFile if there is no more data.
      */
     virtual size_t read(char * data, size_t len) = 0;
 
diff --git a/tests/unit/libutil/compression.cc b/tests/unit/libutil/compression.cc
index f629fb0a3..8d181f53d 100644
--- a/tests/unit/libutil/compression.cc
+++ b/tests/unit/libutil/compression.cc
@@ -136,4 +136,28 @@ TEST_P(PerTypeNonNullCompressionTest, bogusInputDecompression)
     auto bogus = "this data is bogus and should throw when decompressing";
     ASSERT_THROW(decompress(param, bogus), CompressionError);
 }
+
+TEST_P(PerTypeNonNullCompressionTest, truncatedValidInput)
+{
+    auto method = GetParam();
+
+    auto inputString = "the quick brown fox jumps over the lazy doggos";
+    auto compressed = compress(method, inputString);
+
+    /* n.b. This also tests zero-length input, which is also invalid.
+     * As of the writing of this comment, it returns empty output, but is
+     * allowed to throw a compression error instead. */
+    for (int i = 0; i < compressed.length(); ++i) {
+        auto newCompressed = compressed.substr(compressed.length() - i);
+        try {
+            decompress(method, newCompressed);
+            // Success is acceptable as well, even though it is corrupt data.
+            // The compression method is not expected to provide integrity,
+            // just, not break explosively on bad input.
+        } catch (CompressionError &) {
+            // Acceptable
+        }
+    }
+}
+
 }