Merge "Fix bug in writing zips."

libziparchive/zip_writer.cc

@@ -267,12 +267,12 @@ int32_t ZipWriter::CompressBytes(FileInfo* file, const void* data, size_t len) {
 
     if (z_stream_->avail_out == 0) {
       // The output is full, let's write it to disk.
-      size_t dataToWrite = z_stream_->next_out - buffer_.data();
-      if (fwrite(buffer_.data(), 1, dataToWrite, file_) != dataToWrite) {
+      size_t write_bytes = z_stream_->next_out - buffer_.data();
+      if (fwrite(buffer_.data(), 1, write_bytes, file_) != write_bytes) {
         return HandleError(kIoError);
       }
-      file->compressed_size += dataToWrite;
-      current_offset_ += dataToWrite;
+      file->compressed_size += write_bytes;
+      current_offset_ += write_bytes;
 
       // Reset the output buffer for the next input.
       z_stream_->next_out = buffer_.data();
@@ -288,18 +288,32 @@ int32_t ZipWriter::FlushCompressedBytes(FileInfo* file) {
   assert(z_stream_->next_out != nullptr);
   assert(z_stream_->avail_out != 0);
 
-  int zerr = deflate(z_stream_.get(), Z_FINISH);
+  // Keep deflating while there isn't enough space in the buffer to
+  // to complete the compress.
+  int zerr;
+  while ((zerr = deflate(z_stream_.get(), Z_FINISH)) == Z_OK) {
+    assert(z_stream_->avail_out == 0);
+    size_t write_bytes = z_stream_->next_out - buffer_.data();
+    if (fwrite(buffer_.data(), 1, write_bytes, file_) != write_bytes) {
+      return HandleError(kIoError);
+    }
+    file->compressed_size += write_bytes;
+    current_offset_ += write_bytes;
+
+    z_stream_->next_out = buffer_.data();
+    z_stream_->avail_out = buffer_.size();
+  }
   if (zerr != Z_STREAM_END) {
     return HandleError(kZlibError);
   }
 
-  size_t dataToWrite = z_stream_->next_out - buffer_.data();
-  if (dataToWrite != 0) {
-    if (fwrite(buffer_.data(), 1, dataToWrite, file_) != dataToWrite) {
+  size_t write_bytes = z_stream_->next_out - buffer_.data();
+  if (write_bytes != 0) {
+    if (fwrite(buffer_.data(), 1, write_bytes, file_) != write_bytes) {
       return HandleError(kIoError);
     }
-    file->compressed_size += dataToWrite;
-    current_offset_ += dataToWrite;
+    file->compressed_size += write_bytes;
+    current_offset_ += write_bytes;
   }
   z_stream_.reset();
   return kNoError;

libziparchive/zip_writer_test.cc

@@ -20,6 +20,7 @@
 #include <base/test_utils.h>
 #include <gtest/gtest.h>
 #include <memory>
+#include <vector>
 
 struct zipwriter : public ::testing::Test {
   TemporaryFile* temp_file_;
@@ -168,3 +169,40 @@ TEST_F(zipwriter, WriteCompressedZipWithOneFile) {
 
   CloseArchive(handle);
 }
+
+TEST_F(zipwriter, WriteCompressedZipFlushFull) {
+  // This exact data will cause the Finish() to require multiple calls
+  // to deflate() because the ZipWriter buffer isn't big enough to hold
+  // the entire compressed data buffer.
+  constexpr size_t kBufSize = 10000000;
+  std::vector<uint8_t> buffer(kBufSize);
+  size_t prev = 1;
+  for (size_t i = 0; i < kBufSize; i++) {
+    buffer[i] = i + prev;
+    prev = i;
+  }
+
+  ZipWriter writer(file_);
+  ASSERT_EQ(0, writer.StartEntry("file.txt", ZipWriter::kCompress));
+  ASSERT_EQ(0, writer.WriteBytes(buffer.data(), buffer.size()));
+  ASSERT_EQ(0, writer.FinishEntry());
+  ASSERT_EQ(0, writer.Finish());
+
+  ASSERT_GE(0, lseek(fd_, 0, SEEK_SET));
+
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveFd(fd_, "temp", &handle, false));
+
+  ZipEntry data;
+  ASSERT_EQ(0, FindEntry(handle, ZipString("file.txt"), &data));
+  EXPECT_EQ(kCompressDeflated, data.method);
+  EXPECT_EQ(kBufSize, data.uncompressed_length);
+
+  std::vector<uint8_t> decompress(kBufSize);
+  memset(decompress.data(), 0, kBufSize);
+  ASSERT_EQ(0, ExtractToMemory(handle, &data, decompress.data(), decompress.size()));
+  EXPECT_EQ(0, memcmp(decompress.data(), buffer.data(), kBufSize))
+      << "Input buffer and output buffer are different.";
+
+  CloseArchive(handle);
+}