Merge pull request #1047 from openzim/compression_infinite_loop

Eliminated the possibility for zim::Compressor being trapped in an infinite loop

Author: kelson42

src/compression.h

@@ -227,10 +227,12 @@ template<typename INFO>
 class Compressor
 {
   public:
-    Compressor(size_t initial_size=1024*1024) :
+    explicit Compressor(size_t initial_size=1024*1024) :
       ret_data(new char[initial_size]),
       ret_size(initial_size)
-    {}
+    {
+      ASSERT(initial_size, >, 0u);
+    }
 
     ~Compressor() = default;
 
@@ -244,32 +246,25 @@ class Compressor
       stream.next_in = (unsigned char*)data;
       stream.avail_in = size;
       while (true) {
+        if (stream.avail_out == 0) {
+          ret_size *= 2;
+          std::unique_ptr<char[]> new_ret_data(new char[ret_size]);
+          memcpy(new_ret_data.get(), ret_data.get(), stream.total_out);
+          stream.next_out = (unsigned char*)(new_ret_data.get() + stream.total_out);
+          stream.avail_out = ret_size - stream.total_out;
+          ret_data = std::move(new_ret_data);
+        }
         auto errcode = INFO::stream_run_encode(&stream, step);
         switch (errcode) {
           case CompStatus::OK:
-            if (stream.avail_out == 0) {
-              // lzma return a OK return status the first time it runs out of output memory.
-              // The BUF_ERROR is returned only the second time we call a lzma_code.
-              continue;
-            } else {
-              return RunnerStatus::NEED_MORE;
-            }
+            return RunnerStatus::NEED_MORE;
           case CompStatus::STREAM_END:
             return RunnerStatus::NEED_MORE;
           case CompStatus::BUF_ERROR:
-            if (stream.avail_out == 0) {
-              //Not enought output size
-              ret_size *= 2;
-              std::unique_ptr<char[]> new_ret_data(new char[ret_size]);
-              memcpy(new_ret_data.get(), ret_data.get(), stream.total_out);
-              stream.next_out = (unsigned char*)(new_ret_data.get() + stream.total_out);
-              stream.avail_out = ret_size - stream.total_out;
-              ret_data = std::move(new_ret_data);
-              continue;
-            } else {
+            if (stream.avail_out != 0) {
               return RunnerStatus::ERROR;
             }
-          break;
+            break;
           default:
             // unreachable
             return RunnerStatus::ERROR;

test/compression.cpp

@@ -99,5 +99,32 @@ TYPED_TEST(CompressionTest, compress) {
   }
 }
 
+TEST(CompressionTest, zstdCompressionOutputExactlyFillsTheBuffer) {
+  const size_t DATA_SIZE = 10000000;
+  std::vector<char> data(DATA_SIZE, '\0');
+  for (size_t i = 0; i < DATA_SIZE; ++i ) {
+    data[i] = 1 + (char)(i%255);
+  }
+  std::vector<unsigned char> outputDataBuffer(DATA_SIZE, '\0');
+
+  ZSTD_INFO::stream_t zstdStream;
+  ZSTD_INFO::init_stream_encoder(&zstdStream, &data[0]);
+  zstdStream.next_in   = (const unsigned char*)&data[0];
+  zstdStream.avail_in  = DATA_SIZE;
+  zstdStream.next_out  = &outputDataBuffer[0];
+  zstdStream.avail_out = DATA_SIZE;
+  ZSTD_INFO::stream_run_encode(&zstdStream, CompStep::STEP);
+  const size_t sizeOfIntermediateOutput = DATA_SIZE - zstdStream.avail_out;
+  ASSERT_GT(sizeOfIntermediateOutput, 0);
+  ASSERT_LT(sizeOfIntermediateOutput, DATA_SIZE);
+
+  // Setup the initial size of the output data buffer to exactly match the size
+  // of output data from feeding in the first chunk of input data
+  zim::Compressor<ZSTD_INFO> compressor(sizeOfIntermediateOutput);
+  compressor.init(&data[0]);
+  compressor.feed(&data[0], DATA_SIZE);
+  zim::zsize_t compressedDataSize2;
+  const auto out2 = compressor.get_data(&compressedDataSize2);
+}
 
 }  // namespace