Fix decompression of truncated data

[folly.git] / folly / io / Compression.cpp

diff --git a/folly/io/Compression.cpp b/folly/io/Compression.cpp
index 8dd19bc79ae80db4f05a64d20df4961a9b43e637..b1671a5645d29cce6c530a99d30cfe3aa4d68a13 100644 (file)
--- a/folly/io/Compression.cpp
+++ b/folly/io/Compression.cpp
@@ -341,6 +341,8 @@ std::unique_ptr<IOBuf> StreamCodec::doCompress(IOBuf const* data) {
     if (output.empty()) {
       buffer->prependChain(addOutputBuffer(output, kDefaultBufferLength));
     }
     if (output.empty()) {
       buffer->prependChain(addOutputBuffer(output, kDefaultBufferLength));
     }

+    size_t const inputSize = input.size();
+    size_t const outputSize = output.size();

     bool const done = compressStream(input, output, flushOp);
     if (done) {
       DCHECK(input.empty());
     bool const done = compressStream(input, output, flushOp);
     if (done) {
       DCHECK(input.empty());


@@ -348,6 +350,9 @@ std::unique_ptr<IOBuf> StreamCodec::doCompress(IOBuf const* data) {
       DCHECK_EQ(current->next(), data);
       break;
     }
       DCHECK_EQ(current->next(), data);
       break;
     }

+    if (inputSize == input.size() && outputSize == output.size()) {
+      throw std::runtime_error("Codec: No forward progress made");
+    }

   }
   buffer->prev()->trimEnd(output.size());
   return buffer;
   }
   buffer->prev()->trimEnd(output.size());
   return buffer;


@@ -395,10 +400,15 @@ std::unique_ptr<IOBuf> StreamCodec::doUncompress(
     if (output.empty()) {
       buffer->prependChain(addOutputBuffer(output, defaultBufferLength));
     }
     if (output.empty()) {
       buffer->prependChain(addOutputBuffer(output, defaultBufferLength));
     }

+    size_t const inputSize = input.size();
+    size_t const outputSize = output.size();

     bool const done = uncompressStream(input, output, flushOp);
     if (done) {
       break;
     }
     bool const done = uncompressStream(input, output, flushOp);
     if (done) {
       break;
     }

+    if (inputSize == input.size() && outputSize == output.size()) {
+      throw std::runtime_error("Codec: Truncated data");
+    }

   }
   if (!input.empty()) {
     throw std::runtime_error("Codec: Junk after end of data");
   }
   if (!input.empty()) {
     throw std::runtime_error("Codec: Junk after end of data");


@@ -2008,8 +2018,11 @@ std::unique_ptr<IOBuf> Bzip2Codec::doUncompress(
     if (stream.avail_out == 0) {
       out->prependChain(addOutputBuffer(&stream, kDefaultBufferLength));
     }
     if (stream.avail_out == 0) {
       out->prependChain(addOutputBuffer(&stream, kDefaultBufferLength));
     }

-
+    size_t const outputSize = stream.avail_out;

     rc = bzCheck(BZ2_bzDecompress(&stream));
     rc = bzCheck(BZ2_bzDecompress(&stream));

+    if (outputSize == stream.avail_out) {
+      throw std::runtime_error("Bzip2Codec: Truncated input");
+    }

   }
 
   out->prev()->trimEnd(stream.avail_out);
   }
 
   out->prev()->trimEnd(stream.avail_out);