diff --git a/libziparchive/zip_writer.cc b/libziparchive/zip_writer.cc index 22a7c53b2..f117cc5af 100644 --- a/libziparchive/zip_writer.cc +++ b/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; diff --git a/libziparchive/zip_writer_test.cc b/libziparchive/zip_writer_test.cc index 046f19572..f752b7e6e 100644 --- a/libziparchive/zip_writer_test.cc +++ b/libziparchive/zip_writer_test.cc @@ -20,6 +20,7 @@ #include #include #include +#include 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 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 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); +}