libutil: BrotliDecompression{Sink -> Source}

Change-Id: I9579dd08f7bd0f927bde9d3128515b0cee15f320

1 files changed, 54 insertions, 40 deletions

diff --git a/src/libutil/compression.cc b/src/libutil/compression.cc
index 678557a58..e78d76500 100644
--- a/src/libutil/compression.cc
+++ b/src/libutil/compression.cc
@@ -137,53 +137,55 @@ struct NoneSink : CompressionSink
     void writeUnbuffered(std::string_view data) override { nextSink(data); }
 };
 
-struct BrotliDecompressionSink : ChunkedCompressionSink
+struct BrotliDecompressionSource : Source
 {
-    Sink & nextSink;
-    BrotliDecoderState * state;
-    bool finished = false;
-
-    BrotliDecompressionSink(Sink & nextSink) : nextSink(nextSink)
+    static constexpr size_t BUF_SIZE = 32 * 1024;
+    std::unique_ptr<char[]> buf;
+    size_t avail_in = 0;
+    const uint8_t * next_in;
+
+    Source * inner;
+    std::unique_ptr<BrotliDecoderState, void (*)(BrotliDecoderState *)> state;
+
+    BrotliDecompressionSource(Source & inner)
+        : buf(std::make_unique<char[]>(BUF_SIZE))
+        , inner(&inner)
+        , state{
+              BrotliDecoderCreateInstance(nullptr, nullptr, nullptr), BrotliDecoderDestroyInstance}
     {
-        state = BrotliDecoderCreateInstance(nullptr, nullptr, nullptr);
-        if (!state)
+        if (!state) {
             throw CompressionError("unable to initialize brotli decoder");
+        }
     }
 
-    ~BrotliDecompressionSink()
-    {
-        BrotliDecoderDestroyInstance(state);
-    }
-
-    void finish() override
-    {
-        flush();
-        writeInternal({});
-    }
-
-    void writeInternal(std::string_view data) override
+    size_t read(char * data, size_t len) override
     {
-        auto next_in = (const uint8_t *) data.data();
-        size_t avail_in = data.size();
-        uint8_t * next_out = outbuf;
-        size_t avail_out = sizeof(outbuf);
-
-        while (!finished && (!data.data() || avail_in)) {
-            checkInterrupt();
-
-            if (!BrotliDecoderDecompressStream(state,
-                    &avail_in, &next_in,
-                    &avail_out, &next_out,
-                    nullptr))
-                throw CompressionError("error while decompressing brotli file");
-
-            if (avail_out < sizeof(outbuf) || avail_in == 0) {
-                nextSink({(char *) outbuf, sizeof(outbuf) - avail_out});
-                next_out = outbuf;
-                avail_out = sizeof(outbuf);
+        uint8_t * out = (uint8_t *) data;
+        const auto * begin = out;
+
+        try {
+            while (len && !BrotliDecoderIsFinished(state.get())) {
+                checkInterrupt();
+
+                while (avail_in == 0) {
+                    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 &) {
+        }
 
-            finished = BrotliDecoderIsFinished(state);
+        if (begin != out) {
+            return out - begin;
+        } else {
+            throw EndOfFile("brotli stream exhausted");
         }
     }
 };
@@ -202,7 +204,19 @@ std::unique_ptr<FinishSink> makeDecompressionSink(const std::string & method, Si
     if (method == "none" || method == "")
         return std::make_unique<NoneSink>(nextSink);
     else if (method == "br")
-        return std::make_unique<BrotliDecompressionSink>(nextSink);
+        return sourceToSink([&](Source & source) {
+            BrotliDecompressionSource wrapped{source};
+            wrapped.drainInto(nextSink);
+            // special handling because sourceToSink is screwy: try
+            // to read the source one final time and fail when that
+            // succeeds (to reject trailing garbage in input data).
+            try {
+                char buf;
+                source(&buf, 1);
+                throw Error("garbage at end of brotli stream detected");
+            } catch (EndOfFile &) {
+            }
+        });
     else
         return sourceToSink([&](Source & source) {
             auto decompressionSource = std::make_unique<ArchiveDecompressionSource>(source);