Donate to e Foundation | Murena handsets with /e/OS | Own a part of Murena! Learn more

Commit a2a32b0e authored by Christopher Ferris's avatar Christopher Ferris
Browse files

Fix bug in writing zips.

The code does not handle an edge case where writing a compressed
image can overflow the avail_out data when doing a flush. Add a
loop to keep writing the data while deflate indicates that it doesn't
have enough space to write out the compressed data during the flush.

Change-Id: I1f1d1646457ed9b67ed24f6582688c300186c23c
parent 77cdfbd5
Loading
Loading
Loading
Loading
+24 −10
Original line number Diff line number Diff line
@@ -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;
+38 −0
Original line number Diff line number Diff line
@@ -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);
}