4 files changed, 653 insertions, 618 deletions

diff --git a/applypatch/Android.mk b/applypatch/Android.mk
index 47ddad6c8..8be5c36be 100644
--- a/applypatch/Android.mk
+++ b/applypatch/Android.mk
@@ -136,6 +136,9 @@ libimgdiff_static_libraries := \
     libbsdiff \
     libdivsufsort \
     libdivsufsort64 \
+    libziparchive \
+    libutils \
+    liblog \
     libbase \
     libz
 
diff --git a/applypatch/imgdiff.cpp b/applypatch/imgdiff.cpp
index 2f0e1651c..fba74e836 100644
--- a/applypatch/imgdiff.cpp
+++ b/applypatch/imgdiff.cpp
@@ -132,502 +132,611 @@
 #include <sys/types.h>
 #include <unistd.h>
 
+#include <algorithm>
+#include <string>
+#include <vector>
+
 #include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/memory.h>
 #include <android-base/unique_fd.h>
+#include <ziparchive/zip_archive.h>
 
 #include <bsdiff.h>
 #include <zlib.h>
 
 #include "utils.h"
 
-typedef struct {
-  int type;             // CHUNK_NORMAL, CHUNK_DEFLATE
-  size_t start;         // offset of chunk in original image file
+using android::base::get_unaligned;
+
+static constexpr auto BUFFER_SIZE = 0x8000;
+
+class ImageChunk {
+ public:
+  static constexpr auto WINDOWBITS = -15;  // 32kb window; negative to indicate a raw stream.
+  static constexpr auto MEMLEVEL = 8;      // the default value.
+  static constexpr auto METHOD = Z_DEFLATED;
+  static constexpr auto STRATEGY = Z_DEFAULT_STRATEGY;
+
+  ImageChunk(int type, size_t start, const std::vector<uint8_t>* file_content, size_t raw_data_len)
+      : type_(type),
+        start_(start),
+        input_file_ptr_(file_content),
+        raw_data_len_(raw_data_len),
+        entry_name_(""),
+        compress_level_(6),
+        source_start_(0),
+        source_len_(0),
+        source_uncompressed_len_(0) {}
+
+  int GetType() const {
+    return type_;
+  }
+  size_t GetRawDataLength() const {
+    return raw_data_len_;
+  }
+  const std::string& GetEntryName() const {
+    return entry_name_;
+  }
 
-  size_t len;
-  unsigned char* data;  // data to be patched (uncompressed, for deflate chunks)
+  // CHUNK_DEFLATE will return the uncompressed data for diff, while other types will simply return
+  // the raw data.
+  const uint8_t * DataForPatch() const;
+  size_t DataLengthForPatch() const;
 
-  size_t source_start;
-  size_t source_len;
+  void Dump() const {
+    printf("type %d start %zu len %zu\n", type_, start_, DataLengthForPatch());
+  }
 
-  // --- for CHUNK_DEFLATE chunks only: ---
+  void SetSourceInfo(const ImageChunk& other);
+  void SetEntryName(std::string entryname);
+  void SetUncompressedData(std::vector<uint8_t> data);
+  bool SetBonusData(const std::vector<uint8_t>& bonus_data);
+
+  bool operator==(const ImageChunk& other) const;
+  bool operator!=(const ImageChunk& other) const {
+    return !(*this == other);
+  }
 
-  // original (compressed) deflate data
-  size_t deflate_len;
-  unsigned char* deflate_data;
+  size_t GetHeaderSize(size_t patch_size) const;
+  size_t WriteHeaderToFile(FILE* f, const std::vector<uint8_t> patch, size_t offset);
+
+  /*
+   * Cause a gzip chunk to be treated as a normal chunk (ie, as a blob
+   * of uninterpreted data).  The resulting patch will likely be about
+   * as big as the target file, but it lets us handle the case of images
+   * where some gzip chunks are reconstructible but others aren't (by
+   * treating the ones that aren't as normal chunks).
+   */
+  void ChangeDeflateChunkToNormal();
+  bool ChangeChunkToRaw(size_t patch_size);
+
+  /*
+   * Verify that we can reproduce exactly the same compressed data that
+   * we started with.  Sets the level, method, windowBits, memLevel, and
+   * strategy fields in the chunk to the encoding parameters needed to
+   * produce the right output.
+   */
+  bool ReconstructDeflateChunk();
+  bool IsAdjacentNormal(const ImageChunk& other) const;
+  void MergeAdjacentNormal(const ImageChunk& other);
+
+ private:
+  int type_;                                   // CHUNK_NORMAL, CHUNK_DEFLATE, CHUNK_RAW
+  size_t start_;                               // offset of chunk in the original input file
+  const std::vector<uint8_t>* input_file_ptr_; // pointer to the full content of original input file
+  size_t raw_data_len_;
 
-  char* filename;       // used for zip entries
+  // --- for CHUNK_DEFLATE chunks only: ---
+  std::vector<uint8_t> uncompressed_data_;
+  std::string entry_name_;  // used for zip entries
 
   // deflate encoder parameters
-  int level, method, windowBits, memLevel, strategy;
-
-  size_t source_uncompressed_len;
-} ImageChunk;
-
-typedef struct {
-  int data_offset;
-  int deflate_len;
-  int uncomp_len;
-  char* filename;
-} ZipFileEntry;
-
-static int fileentry_compare(const void* a, const void* b) {
-  int ao = ((ZipFileEntry*)a)->data_offset;
-  int bo = ((ZipFileEntry*)b)->data_offset;
-  if (ao < bo) {
-    return -1;
-  } else if (ao > bo) {
-    return 1;
-  } else {
-    return 0;
-  }
+  int compress_level_;
+
+  size_t source_start_;
+  size_t source_len_;
+  size_t source_uncompressed_len_;
+
+  const uint8_t* GetRawData() const;
+  bool TryReconstruction(int level);
+};
+
+const uint8_t* ImageChunk::GetRawData() const {
+  CHECK_LE(start_ + raw_data_len_, input_file_ptr_->size());
+  return input_file_ptr_->data() + start_;
 }
 
-unsigned char* ReadZip(const char* filename,
-                       int* num_chunks, ImageChunk** chunks,
-                       int include_pseudo_chunk) {
-  struct stat st;
-  if (stat(filename, &st) != 0) {
-    printf("failed to stat \"%s\": %s\n", filename, strerror(errno));
-    return NULL;
+const uint8_t * ImageChunk::DataForPatch() const {
+  if (type_ == CHUNK_DEFLATE) {
+    return uncompressed_data_.data();
   }
+  return GetRawData();
+}
 
-  size_t sz = static_cast<size_t>(st.st_size);
-  unsigned char* img = static_cast<unsigned char*>(malloc(sz));
-  FILE* f = fopen(filename, "rb");
-  if (fread(img, 1, sz, f) != sz) {
-    printf("failed to read \"%s\" %s\n", filename, strerror(errno));
-    fclose(f);
-    free(img);
-    return NULL;
+size_t ImageChunk::DataLengthForPatch() const {
+  if (type_ == CHUNK_DEFLATE) {
+    return uncompressed_data_.size();
   }
-  fclose(f);
+  return raw_data_len_;
+}
 
-  // look for the end-of-central-directory record.
+bool ImageChunk::operator==(const ImageChunk& other) const {
+  if (type_ != other.type_) {
+    return false;
+  }
+  return (raw_data_len_ == other.raw_data_len_ &&
+          memcmp(GetRawData(), other.GetRawData(), raw_data_len_) == 0);
+}
 
-  int i;
-  for (i = st.st_size-20; i >= 0 && i > st.st_size - 65600; --i) {
-    if (img[i] == 0x50 && img[i+1] == 0x4b &&
-        img[i+2] == 0x05 && img[i+3] == 0x06) {
-      break;
-    }
+void ImageChunk::SetSourceInfo(const ImageChunk& src) {
+  source_start_ = src.start_;
+  if (type_ == CHUNK_NORMAL) {
+    source_len_ = src.raw_data_len_;
+  } else if (type_ == CHUNK_DEFLATE) {
+    source_len_ = src.raw_data_len_;
+    source_uncompressed_len_ = src.uncompressed_data_.size();
   }
-  // double-check: this archive consists of a single "disk"
-  if (!(img[i+4] == 0 && img[i+5] == 0 && img[i+6] == 0 && img[i+7] == 0)) {
-    printf("can't process multi-disk archive\n");
-    return NULL;
+}
+
+void ImageChunk::SetEntryName(std::string entryname) {
+  entry_name_ = entryname;
+}
+
+void ImageChunk::SetUncompressedData(std::vector<uint8_t> data) {
+  uncompressed_data_ = data;
+}
+
+bool ImageChunk::SetBonusData(const std::vector<uint8_t>& bonus_data) {
+  if (type_ != CHUNK_DEFLATE) {
+    return false;
   }
+  uncompressed_data_.insert(uncompressed_data_.end(), bonus_data.begin(), bonus_data.end());
+  return true;
+}
 
-  int cdcount = Read2(img+i+8);
-  int cdoffset = Read4(img+i+16);
+// Convert CHUNK_NORMAL & CHUNK_DEFLATE to CHUNK_RAW if the terget size is
+// smaller. Also take the header size into account during size comparison.
+bool ImageChunk::ChangeChunkToRaw(size_t patch_size) {
+  if (type_ == CHUNK_RAW) {
+    return true;
+  } else if (type_ == CHUNK_NORMAL && (raw_data_len_ <= 160 || raw_data_len_ < patch_size)) {
+    type_ = CHUNK_RAW;
+    return true;
+  }
+  return false;
+}
 
-  ZipFileEntry* temp_entries = static_cast<ZipFileEntry*>(malloc(
-      cdcount * sizeof(ZipFileEntry)));
-  int entrycount = 0;
+void ImageChunk::ChangeDeflateChunkToNormal() {
+  if (type_ != CHUNK_DEFLATE) return;
+  type_ = CHUNK_NORMAL;
+  uncompressed_data_.clear();
+}
 
-  unsigned char* cd = img+cdoffset;
-  for (i = 0; i < cdcount; ++i) {
-    if (!(cd[0] == 0x50 && cd[1] == 0x4b && cd[2] == 0x01 && cd[3] == 0x02)) {
-      printf("bad central directory entry %d\n", i);
-      free(temp_entries);
-      return NULL;
-    }
+// Header size:
+// header_type    4 bytes
+// CHUNK_NORMAL   8*3 = 24 bytes
+// CHUNK_DEFLATE  8*5 + 4*5 = 60 bytes
+// CHUNK_RAW      4 bytes
+size_t ImageChunk::GetHeaderSize(size_t patch_size) const {
+  switch (type_) {
+    case CHUNK_NORMAL:
+      return 4 + 8 * 3;
+    case CHUNK_DEFLATE:
+      return 4 + 8 * 5 + 4 * 5;
+    case CHUNK_RAW:
+      return 4 + 4 + patch_size;
+    default:
+      printf("unexpected chunk type: %d\n", type_);  // should not reach here.
+      CHECK(false);
+      return 0;
+  }
+}
 
-    int clen = Read4(cd+20);   // compressed len
-    int ulen = Read4(cd+24);   // uncompressed len
-    int nlen = Read2(cd+28);   // filename len
-    int xlen = Read2(cd+30);   // extra field len
-    int mlen = Read2(cd+32);   // file comment len
-    int hoffset = Read4(cd+42);   // local header offset
+size_t ImageChunk::WriteHeaderToFile(FILE* f, const std::vector<uint8_t> patch, size_t offset) {
+  Write4(type_, f);
+  switch (type_) {
+    case CHUNK_NORMAL:
+      printf("normal   (%10zu, %10zu)  %10zu\n", start_, raw_data_len_, patch.size());
+      Write8(source_start_, f);
+      Write8(source_len_, f);
+      Write8(offset, f);
+      return offset + patch.size();
+    case CHUNK_DEFLATE:
+      printf("deflate  (%10zu, %10zu)  %10zu  %s\n", start_, raw_data_len_, patch.size(),
+             entry_name_.c_str());
+      Write8(source_start_, f);
+      Write8(source_len_, f);
+      Write8(offset, f);
+      Write8(source_uncompressed_len_, f);
+      Write8(uncompressed_data_.size(), f);
+      Write4(compress_level_, f);
+      Write4(METHOD, f);
+      Write4(WINDOWBITS, f);
+      Write4(MEMLEVEL, f);
+      Write4(STRATEGY, f);
+      return offset + patch.size();
+    case CHUNK_RAW:
+      printf("raw      (%10zu, %10zu)\n", start_, raw_data_len_);
+      Write4(patch.size(), f);
+      fwrite(patch.data(), 1, patch.size(), f);
+      return offset;
+    default:
+      printf("unexpected chunk type: %d\n", type_);
+      CHECK(false);
+      return offset;
+  }
+}
 
-    char* filename = static_cast<char*>(malloc(nlen+1));
-    memcpy(filename, cd+46, nlen);
-    filename[nlen] = '\0';
+bool ImageChunk::IsAdjacentNormal(const ImageChunk& other) const {
+  if (type_ != CHUNK_NORMAL || other.type_ != CHUNK_NORMAL) {
+    return false;
+  }
+  return (other.start_ == start_ + raw_data_len_);
+}
 
-    int method = Read2(cd+10);
+void ImageChunk::MergeAdjacentNormal(const ImageChunk& other) {
+  CHECK(IsAdjacentNormal(other));
+  raw_data_len_ = raw_data_len_ + other.raw_data_len_;
+}
 
-    cd += 46 + nlen + xlen + mlen;
+bool ImageChunk::ReconstructDeflateChunk() {
+  if (type_ != CHUNK_DEFLATE) {
+    printf("attempt to reconstruct non-deflate chunk\n");
+    return false;
+  }
 
-    if (method != 8) {  // 8 == deflate
-      free(filename);
-      continue;
+  // We only check two combinations of encoder parameters:  level 6
+  // (the default) and level 9 (the maximum).
+  for (int level = 6; level <= 9; level += 3) {
+    if (TryReconstruction(level)) {
+      compress_level_ = level;
+      return true;
     }
+  }
+
+  return false;
+}
 
-    unsigned char* lh = img + hoffset;
+/*
+ * Takes the uncompressed data stored in the chunk, compresses it
+ * using the zlib parameters stored in the chunk, and checks that it
+ * matches exactly the compressed data we started with (also stored in
+ * the chunk).
+ */
+bool ImageChunk::TryReconstruction(int level) {
+  z_stream strm;
+  strm.zalloc = Z_NULL;
+  strm.zfree = Z_NULL;
+  strm.opaque = Z_NULL;
+  strm.avail_in = uncompressed_data_.size();
+  strm.next_in = uncompressed_data_.data();
+  int ret = deflateInit2(&strm, level, METHOD, WINDOWBITS, MEMLEVEL, STRATEGY);
+  if (ret < 0) {
+    printf("failed to initialize deflate: %d\n", ret);
+    return false;
+  }
 
-    if (!(lh[0] == 0x50 && lh[1] == 0x4b && lh[2] == 0x03 && lh[3] == 0x04)) {
-      printf("bad local file header entry %d\n", i);
-      return NULL;
+  std::vector<uint8_t> buffer(BUFFER_SIZE);
+  size_t offset = 0;
+  do {
+    strm.avail_out = buffer.size();
+    strm.next_out = buffer.data();
+    ret = deflate(&strm, Z_FINISH);
+    if (ret < 0) {
+      printf("failed to deflate: %d\n", ret);
+      return false;
     }
 
-    if (Read2(lh+26) != nlen || memcmp(lh+30, filename, nlen) != 0) {
-      printf("central dir filename doesn't match local header\n");
-      return NULL;
+    size_t compressed_size = buffer.size() - strm.avail_out;
+    if (memcmp(buffer.data(), input_file_ptr_->data() + start_ + offset, compressed_size) != 0) {
+      // mismatch; data isn't the same.
+      deflateEnd(&strm);
+      return false;
     }
+    offset += compressed_size;
+  } while (ret != Z_STREAM_END);
+  deflateEnd(&strm);
 
-    xlen = Read2(lh+28);   // extra field len; might be different from CD entry?
+  if (offset != raw_data_len_) {
+    // mismatch; ran out of data before we should have.
+    return false;
+  }
+  return true;
+}
 
-    temp_entries[entrycount].data_offset = hoffset+30+nlen+xlen;
-    temp_entries[entrycount].deflate_len = clen;
-    temp_entries[entrycount].uncomp_len = ulen;
-    temp_entries[entrycount].filename = filename;
-    ++entrycount;
+// EOCD record
+// offset 0: signature 0x06054b50, 4 bytes
+// offset 4: number of this disk, 2 bytes
+// ...
+// offset 20: comment length, 2 bytes
+// offset 22: comment, n bytes
+static bool GetZipFileSize(const std::vector<uint8_t>& zip_file, size_t* input_file_size) {
+  if (zip_file.size() < 22) {
+    printf("file is too small to be a zip file\n");
+    return false;
   }
 
-  qsort(temp_entries, entrycount, sizeof(ZipFileEntry), fileentry_compare);
-
-#if 0
-  printf("found %d deflated entries\n", entrycount);
-  for (i = 0; i < entrycount; ++i) {
-    printf("off %10d  len %10d unlen %10d   %p %s\n",
-           temp_entries[i].data_offset,
-           temp_entries[i].deflate_len,
-           temp_entries[i].uncomp_len,
-           temp_entries[i].filename,
-           temp_entries[i].filename);
+  // Look for End of central directory record of the zip file, and calculate the actual
+  // zip_file size.
+  for (int i = zip_file.size() - 22; i >= 0; i--) {
+    if (zip_file[i] == 0x50) {
+      if (get_unaligned<uint32_t>(&zip_file[i]) == 0x06054b50) {
+        // double-check: this archive consists of a single "disk".
+        CHECK_EQ(get_unaligned<uint16_t>(&zip_file[i + 4]), 0);
+
+        uint16_t comment_length = get_unaligned<uint16_t>(&zip_file[i + 20]);
+        size_t file_size = i + 22 + comment_length;
+        CHECK_LE(file_size, zip_file.size());
+        *input_file_size = file_size;
+        return true;
+      }
+    }
   }
-#endif
 
-  *num_chunks = 0;
-  *chunks = static_cast<ImageChunk*>(malloc((entrycount*2+2) * sizeof(ImageChunk)));
-  ImageChunk* curr = *chunks;
+  // EOCD not found, this file is likely not a valid zip file.
+  return false;
+}
 
-  if (include_pseudo_chunk) {
-    curr->type = CHUNK_NORMAL;
-    curr->start = 0;
-    curr->len = st.st_size;
-    curr->data = img;
-    curr->filename = NULL;
-    ++curr;
-    ++*num_chunks;
+static bool ReadZip(const char* filename, std::vector<ImageChunk>* chunks,
+                    std::vector<uint8_t>* zip_file, bool include_pseudo_chunk) {
+  CHECK(zip_file != nullptr);
+  struct stat st;
+  if (stat(filename, &st) != 0) {
+    printf("failed to stat \"%s\": %s\n", filename, strerror(errno));
+    return false;
   }
 
-  int pos = 0;
-  int nextentry = 0;
+  size_t sz = static_cast<size_t>(st.st_size);
+  zip_file->resize(sz);
+  android::base::unique_fd fd(open(filename, O_RDONLY));
+  if (fd == -1) {
+    printf("failed to open \"%s\" %s\n", filename, strerror(errno));
+    return false;
+  }
+  if (!android::base::ReadFully(fd, zip_file->data(), sz)) {
+    printf("failed to read \"%s\" %s\n", filename, strerror(errno));
+    return false;
+  }
+  fd.reset();
 
-  while (pos < st.st_size) {
-    if (nextentry < entrycount && pos == temp_entries[nextentry].data_offset) {
-      curr->type = CHUNK_DEFLATE;
-      curr->start = pos;
-      curr->deflate_len = temp_entries[nextentry].deflate_len;
-      curr->deflate_data = img + pos;
-      curr->filename = temp_entries[nextentry].filename;
+  // Trim the trailing zeros before we pass the file to ziparchive handler.
+  size_t zipfile_size;
+  if (!GetZipFileSize(*zip_file, &zipfile_size)) {
+    printf("failed to parse the actual size of %s\n", filename);
+    return false;
+  }
+  ZipArchiveHandle handle;
+  int err = OpenArchiveFromMemory(zip_file->data(), zipfile_size, filename, &handle);
+  if (err != 0) {
+    printf("failed to open zip file %s: %s\n", filename, ErrorCodeString(err));
+    CloseArchive(handle);
+    return false;
+  }
 
-      curr->len = temp_entries[nextentry].uncomp_len;
-      curr->data = static_cast<unsigned char*>(malloc(curr->len));
+  // Create a list of deflated zip entries, sorted by offset.
+  std::vector<std::pair<std::string, ZipEntry>> temp_entries;
+  void* cookie;
+  int ret = StartIteration(handle, &cookie, nullptr, nullptr);
+  if (ret != 0) {
+    printf("failed to iterate over entries in %s: %s\n", filename, ErrorCodeString(ret));
+    CloseArchive(handle);
+    return false;
+  }
 
-      z_stream strm;
-      strm.zalloc = Z_NULL;
-      strm.zfree = Z_NULL;
-      strm.opaque = Z_NULL;
-      strm.avail_in = curr->deflate_len;
-      strm.next_in = curr->deflate_data;
+  ZipString name;
+  ZipEntry entry;
+  while ((ret = Next(cookie, &entry, &name)) == 0) {
+    if (entry.method == kCompressDeflated) {
+      std::string entryname(name.name, name.name + name.name_length);
+      temp_entries.push_back(std::make_pair(entryname, entry));
+    }
+  }
 
-      // -15 means we are decoding a 'raw' deflate stream; zlib will
-      // not expect zlib headers.
-      int ret = inflateInit2(&strm, -15);
-      if (ret < 0) {
-        printf("failed to initialize inflate: %d\n", ret);
-        return NULL;
-      }
+  if (ret != -1) {
+    printf("Error while iterating over zip entries: %s\n", ErrorCodeString(ret));
+    CloseArchive(handle);
+    return false;
+  }
+  std::sort(temp_entries.begin(), temp_entries.end(),
+            [](auto& entry1, auto& entry2) {
+              return entry1.second.offset < entry2.second.offset;
+            });
+
+  EndIteration(cookie);
+
+  if (include_pseudo_chunk) {
+    chunks->emplace_back(CHUNK_NORMAL, 0, zip_file, zip_file->size());
+  }
 
-      strm.avail_out = curr->len;
-      strm.next_out = curr->data;
-      ret = inflate(&strm, Z_NO_FLUSH);
-      if (ret != Z_STREAM_END) {
-        printf("failed to inflate \"%s\"; %d\n", curr->filename, ret);
-        return NULL;
+  size_t pos = 0;
+  size_t nextentry = 0;
+  while (pos < zip_file->size()) {
+    if (nextentry < temp_entries.size() &&
+        static_cast<off64_t>(pos) == temp_entries[nextentry].second.offset) {
+      // compose the next deflate chunk.
+      std::string entryname = temp_entries[nextentry].first;
+      size_t uncompressed_len = temp_entries[nextentry].second.uncompressed_length;
+      std::vector<uint8_t> uncompressed_data(uncompressed_len);
+      if ((ret = ExtractToMemory(handle, &temp_entries[nextentry].second, uncompressed_data.data(),
+                                 uncompressed_len)) != 0) {
+        printf("failed to extract %s with size %zu: %s\n", entryname.c_str(), uncompressed_len,
+               ErrorCodeString(ret));
+        CloseArchive(handle);
+        return false;
       }
 
-      inflateEnd(&strm);
+      size_t compressed_len = temp_entries[nextentry].second.compressed_length;
+      ImageChunk curr(CHUNK_DEFLATE, pos, zip_file, compressed_len);
+      curr.SetEntryName(std::move(entryname));
+      curr.SetUncompressedData(std::move(uncompressed_data));
+      chunks->push_back(curr);
 
-      pos += curr->deflate_len;
+      pos += compressed_len;
       ++nextentry;
-      ++*num_chunks;
-      ++curr;
       continue;
     }
 
-    // use a normal chunk to take all the data up to the start of the
-    // next deflate section.
-
-    curr->type = CHUNK_NORMAL;
-    curr->start = pos;
-    if (nextentry < entrycount) {
-      curr->len = temp_entries[nextentry].data_offset - pos;
+    // Use a normal chunk to take all the data up to the start of the next deflate section.
+    size_t raw_data_len;
+    if (nextentry < temp_entries.size()) {
+      raw_data_len = temp_entries[nextentry].second.offset - pos;
     } else {
-      curr->len = st.st_size - pos;
+      raw_data_len = zip_file->size() - pos;
     }
-    curr->data = img + pos;
-    curr->filename = NULL;
-    pos += curr->len;
+    chunks->emplace_back(CHUNK_NORMAL, pos, zip_file, raw_data_len);
 
-    ++*num_chunks;
-    ++curr;
+    pos += raw_data_len;
   }
 
-  free(temp_entries);
-  return img;
+  CloseArchive(handle);
+  return true;
 }
 
-/*
- * Read the given file and break it up into chunks, putting the number
- * of chunks and their info in *num_chunks and **chunks,
- * respectively.  Returns a malloc'd block of memory containing the
- * contents of the file; various pointers in the output chunk array
- * will point into this block of memory.  The caller should free the
- * return value when done with all the chunks.  Returns NULL on
- * failure.
- */
-unsigned char* ReadImage(const char* filename, int* num_chunks, ImageChunk** chunks) {
+// Read the given file and break it up into chunks, and putting the data in to a vector.
+static bool ReadImage(const char* filename, std::vector<ImageChunk>* chunks,
+                      std::vector<uint8_t>* img) {
+  CHECK(img != nullptr);
   struct stat st;
   if (stat(filename, &st) != 0) {
     printf("failed to stat \"%s\": %s\n", filename, strerror(errno));
-    return NULL;
+    return false;
   }
 
   size_t sz = static_cast<size_t>(st.st_size);
-  unsigned char* img = static_cast<unsigned char*>(malloc(sz));
+  img->resize(sz);
   android::base::unique_fd fd(open(filename, O_RDONLY));
-  if (!android::base::ReadFully(fd, img, sz)) {
+  if (fd == -1) {
+    printf("failed to open \"%s\" %s\n", filename, strerror(errno));
+    return false;
+  }
+  if (!android::base::ReadFully(fd, img->data(), sz)) {
     printf("failed to read \"%s\" %s\n", filename, strerror(errno));
-    return nullptr;
+    return false;
   }
 
   size_t pos = 0;
 
-  *num_chunks = 0;
-  *chunks = NULL;
-
   while (pos < sz) {
-    unsigned char* p = img + pos;
-
-    if (sz - pos >= 4 &&
-        p[0] == 0x1f && p[1] == 0x8b &&
-        p[2] == 0x08 &&    // deflate compression
-        p[3] == 0x00) {    // no header flags
+    if (sz - pos >= 4 && img->at(pos) == 0x1f && img->at(pos + 1) == 0x8b &&
+        img->at(pos + 2) == 0x08 &&  // deflate compression
+        img->at(pos + 3) == 0x00) {  // no header flags
       // 'pos' is the offset of the start of a gzip chunk.
       size_t chunk_offset = pos;
 
-      *num_chunks += 3;
-      *chunks = static_cast<ImageChunk*>(realloc(*chunks, *num_chunks * sizeof(ImageChunk)));
-      ImageChunk* curr = *chunks + (*num_chunks-3);
-
-      // create a normal chunk for the header.
-      curr->start = pos;
-      curr->type = CHUNK_NORMAL;
-      curr->len = GZIP_HEADER_LEN;
-      curr->data = p;
-
-      pos += curr->len;
-      p += curr->len;
-      ++curr;
-
-      curr->type = CHUNK_DEFLATE;
-      curr->filename = NULL;
+      // The remaining data is too small to be a gzip chunk; treat them as a normal chunk.
+      if (sz - pos < GZIP_HEADER_LEN + GZIP_FOOTER_LEN) {
+        chunks->emplace_back(CHUNK_NORMAL, pos, img, sz - pos);
+        break;
+      }
 
-      // We must decompress this chunk in order to discover where it
-      // ends, and so we can put the uncompressed data and its length
-      // into curr->data and curr->len.
+      // We need three chunks for the deflated image in total, one normal chunk for the header,
+      // one deflated chunk for the body, and another normal chunk for the footer.
+      chunks->emplace_back(CHUNK_NORMAL, pos, img, GZIP_HEADER_LEN);
+      pos += GZIP_HEADER_LEN;
 
-      size_t allocated = 32768;
-      curr->len = 0;
-      curr->data = static_cast<unsigned char*>(malloc(allocated));
-      curr->start = pos;
-      curr->deflate_data = p;
+      // We must decompress this chunk in order to discover where it ends, and so we can update
+      // the uncompressed_data of the image body and its length.
 
       z_stream strm;
       strm.zalloc = Z_NULL;
       strm.zfree = Z_NULL;
       strm.opaque = Z_NULL;
       strm.avail_in = sz - pos;
-      strm.next_in = p;
+      strm.next_in = img->data() + pos;
 
       // -15 means we are decoding a 'raw' deflate stream; zlib will
       // not expect zlib headers.
       int ret = inflateInit2(&strm, -15);
       if (ret < 0) {
         printf("failed to initialize inflate: %d\n", ret);
-        return NULL;
+        return false;
       }
 
+      size_t allocated = BUFFER_SIZE;
+      std::vector<uint8_t> uncompressed_data(allocated);
+      size_t uncompressed_len = 0, raw_data_len = 0;
       do {
-        strm.avail_out = allocated - curr->len;
-        strm.next_out = curr->data + curr->len;
+        strm.avail_out = allocated - uncompressed_len;
+        strm.next_out = uncompressed_data.data() + uncompressed_len;
         ret = inflate(&strm, Z_NO_FLUSH);
         if (ret < 0) {
-          printf("Warning: inflate failed [%s] at offset [%zu],"
-                 " treating as a normal chunk\n",
+          printf("Warning: inflate failed [%s] at offset [%zu], treating as a normal chunk\n",
                  strm.msg, chunk_offset);
           break;
         }
-        curr->len = allocated - strm.avail_out;
+        uncompressed_len = allocated - strm.avail_out;
         if (strm.avail_out == 0) {
           allocated *= 2;
-          curr->data = static_cast<unsigned char*>(realloc(curr->data, allocated));
+          uncompressed_data.resize(allocated);
         }
       } while (ret != Z_STREAM_END);
 
-      curr->deflate_len = sz - strm.avail_in - pos;
+      raw_data_len = sz - strm.avail_in - pos;
       inflateEnd(&strm);
 
       if (ret < 0) {
-        free(curr->data);
-        *num_chunks -= 2;
         continue;
       }
 
-      pos += curr->deflate_len;
-      p += curr->deflate_len;
-      ++curr;
+      ImageChunk body(CHUNK_DEFLATE, pos, img, raw_data_len);
+      uncompressed_data.resize(uncompressed_len);
+      body.SetUncompressedData(std::move(uncompressed_data));
+      chunks->push_back(body);
 
-      // create a normal chunk for the footer
+      pos += raw_data_len;
 
-      curr->type = CHUNK_NORMAL;
-      curr->start = pos;
-      curr->len = GZIP_FOOTER_LEN;
-      curr->data = img+pos;
+      // create a normal chunk for the footer
+      chunks->emplace_back(CHUNK_NORMAL, pos, img, GZIP_FOOTER_LEN);
 
-      pos += curr->len;
-      p += curr->len;
-      ++curr;
+      pos += GZIP_FOOTER_LEN;
 
       // The footer (that we just skipped over) contains the size of
       // the uncompressed data.  Double-check to make sure that it
       // matches the size of the data we got when we actually did
       // the decompression.
-      size_t footer_size = Read4(p-4);
-      if (footer_size != curr[-2].len) {
-        printf("Error: footer size %zu != decompressed size %zu\n", footer_size, curr[-2].len);
-        free(img);
-        return NULL;
+      size_t footer_size = Read4(img->data() + pos - 4);
+      if (footer_size != body.DataLengthForPatch()) {
+        printf("Error: footer size %zu != decompressed size %zu\n", footer_size,
+               body.GetRawDataLength());
+        return false;
       }
     } else {
-      // Reallocate the list for every chunk; we expect the number of
-      // chunks to be small (5 for typical boot and recovery images).
-      ++*num_chunks;
-      *chunks = static_cast<ImageChunk*>(realloc(*chunks, *num_chunks * sizeof(ImageChunk)));
-      ImageChunk* curr = *chunks + (*num_chunks-1);
-      curr->start = pos;
-
-      // 'pos' is not the offset of the start of a gzip chunk, so scan
-      // forward until we find a gzip header.
-      curr->type = CHUNK_NORMAL;
-      curr->data = p;
-
-      for (curr->len = 0; curr->len < (sz - pos); ++curr->len) {
-        if (sz - pos >= 4 && p[curr->len] == 0x1f && p[curr->len + 1] == 0x8b &&
-            p[curr->len + 2] == 0x08 && p[curr->len + 3] == 0x00) {
+      // Use a normal chunk to take all the contents until the next gzip chunk (or EOF); we expect
+      // the number of chunks to be small (5 for typical boot and recovery images).
+
+      // Scan forward until we find a gzip header.
+      size_t data_len = 0;
+      while (data_len + pos < sz) {
+        if (data_len + pos + 4 <= sz && img->at(pos + data_len) == 0x1f &&
+            img->at(pos + data_len + 1) == 0x8b && img->at(pos + data_len + 2) == 0x08 &&
+            img->at(pos + data_len + 3) == 0x00) {
           break;
         }
+        data_len++;
       }
-      pos += curr->len;
-    }
-  }
-
-  return img;
-}
-
-#define BUFFER_SIZE 32768
-
-/*
- * Takes the uncompressed data stored in the chunk, compresses it
- * using the zlib parameters stored in the chunk, and checks that it
- * matches exactly the compressed data we started with (also stored in
- * the chunk).  Return 0 on success.
- */
-int TryReconstruction(ImageChunk* chunk, unsigned char* out) {
-  size_t p = 0;
-
-#if 0
-  printf("trying %d %d %d %d %d\n",
-          chunk->level, chunk->method, chunk->windowBits,
-          chunk->memLevel, chunk->strategy);
-#endif
-
-  z_stream strm;
-  strm.zalloc = Z_NULL;
-  strm.zfree = Z_NULL;
-  strm.opaque = Z_NULL;
-  strm.avail_in = chunk->len;
-  strm.next_in = chunk->data;
-  int ret;
-  ret = deflateInit2(&strm, chunk->level, chunk->method, chunk->windowBits,
-                     chunk->memLevel, chunk->strategy);
-  if (ret < 0) {
-    printf("failed to initialize deflate: %d\n", ret);
-    return -1;
-  }
-  do {
-    strm.avail_out = BUFFER_SIZE;
-    strm.next_out = out;
-    ret = deflate(&strm, Z_FINISH);
-    if (ret < 0) {
-      printf("failed to deflate: %d\n", ret);
-      return -1;
-    }
-    size_t have = BUFFER_SIZE - strm.avail_out;
-
-    if (memcmp(out, chunk->deflate_data+p, have) != 0) {
-      // mismatch; data isn't the same.
-      deflateEnd(&strm);
-      return -1;
-    }
-    p += have;
-  } while (ret != Z_STREAM_END);
-  deflateEnd(&strm);
-  if (p != chunk->deflate_len) {
-    // mismatch; ran out of data before we should have.
-    return -1;
-  }
-  return 0;
-}
-
-/*
- * Verify that we can reproduce exactly the same compressed data that
- * we started with.  Sets the level, method, windowBits, memLevel, and
- * strategy fields in the chunk to the encoding parameters needed to
- * produce the right output.  Returns 0 on success.
- */
-int ReconstructDeflateChunk(ImageChunk* chunk) {
-  if (chunk->type != CHUNK_DEFLATE) {
-    printf("attempt to reconstruct non-deflate chunk\n");
-    return -1;
-  }
+      chunks->emplace_back(CHUNK_NORMAL, pos, img, data_len);
 
-  unsigned char* out = static_cast<unsigned char*>(malloc(BUFFER_SIZE));
-
-  // We only check two combinations of encoder parameters:  level 6
-  // (the default) and level 9 (the maximum).
-  for (chunk->level = 6; chunk->level <= 9; chunk->level += 3) {
-    chunk->windowBits = -15;  // 32kb window; negative to indicate a raw stream.
-    chunk->memLevel = 8;      // the default value.
-    chunk->method = Z_DEFLATED;
-    chunk->strategy = Z_DEFAULT_STRATEGY;
-
-    if (TryReconstruction(chunk, out) == 0) {
-      free(out);
-      return 0;
+      pos += data_len;
     }
   }
 
-  free(out);
-  return -1;
+  return true;
 }
 
 /*
  * Given source and target chunks, compute a bsdiff patch between them.
- * Return the patch data, placing its length in *size. Return NULL on failure.
+ * Store the result in the patch_data.
  * |bsdiff_cache| can be used to cache the suffix array if the same |src| chunk
  * is used repeatedly, pass nullptr if not needed.
  */
-unsigned char* MakePatch(ImageChunk* src, ImageChunk* tgt, size_t* size, saidx_t** bsdiff_cache) {
-  if (tgt->type == CHUNK_NORMAL) {
-    if (tgt->len <= 160) {
-      tgt->type = CHUNK_RAW;
-      *size = tgt->len;
-      return tgt->data;
-    }
+static bool MakePatch(const ImageChunk* src, ImageChunk* tgt, std::vector<uint8_t>* patch_data,
+                      saidx_t** bsdiff_cache) {
+  if (tgt->ChangeChunkToRaw(0)) {
+    size_t patch_size = tgt->DataLengthForPatch();
+    patch_data->resize(patch_size);
+    std::copy(tgt->DataForPatch(), tgt->DataForPatch() + patch_size, patch_data->begin());
+    return true;
   }
 
 #if defined(__ANDROID__)
@@ -635,104 +744,51 @@ unsigned char* MakePatch(ImageChunk* src, ImageChunk* tgt, size_t* size, saidx_t
 #else
   char ptemp[] = "/tmp/imgdiff-patch-XXXXXX";
 #endif
-  int fd = mkstemp(ptemp);
 
+  int fd = mkstemp(ptemp);
   if (fd == -1) {
-    printf("MakePatch failed to create a temporary file: %s\n",
-           strerror(errno));
-    return NULL;
+    printf("MakePatch failed to create a temporary file: %s\n", strerror(errno));
+    return false;
   }
-  close(fd); // temporary file is created and we don't need its file
-             // descriptor
+  close(fd);
 
-  int r = bsdiff::bsdiff(src->data, src->len, tgt->data, tgt->len, ptemp, bsdiff_cache);
+  int r = bsdiff::bsdiff(src->DataForPatch(), src->DataLengthForPatch(), tgt->DataForPatch(),
+                         tgt->DataLengthForPatch(), ptemp, bsdiff_cache);
   if (r != 0) {
     printf("bsdiff() failed: %d\n", r);
-    return NULL;
+    return false;
   }
 
   struct stat st;
   if (stat(ptemp, &st) != 0) {
-    printf("failed to stat patch file %s: %s\n",
-            ptemp, strerror(errno));
-    return NULL;
+    printf("failed to stat patch file %s: %s\n", ptemp, strerror(errno));
+    return false;
   }
 
   size_t sz = static_cast<size_t>(st.st_size);
-  // TODO: Memory leak on error return.
-  unsigned char* data = static_cast<unsigned char*>(malloc(sz));
-
-  if (tgt->type == CHUNK_NORMAL && tgt->len <= sz) {
+  if (tgt->ChangeChunkToRaw(sz)) {
     unlink(ptemp);
-
-    tgt->type = CHUNK_RAW;
-    *size = tgt->len;
-    return tgt->data;
+    size_t patch_size = tgt->DataLengthForPatch();
+    patch_data->resize(patch_size);
+    std::copy(tgt->DataForPatch(), tgt->DataForPatch() + patch_size, patch_data->begin());
+    return true;
   }
 
-  *size = sz;
-
-  FILE* f = fopen(ptemp, "rb");
-  if (f == NULL) {
-    printf("failed to open patch %s: %s\n", ptemp, strerror(errno));
-    return NULL;
+  android::base::unique_fd patch_fd(open(ptemp, O_RDONLY));
+  if (patch_fd == -1) {
+    printf("failed to open %s: %s\n", ptemp, strerror(errno));
+    return false;
   }
-  if (fread(data, 1, sz, f) != sz) {
-    printf("failed to read patch %s: %s\n", ptemp, strerror(errno));
-    return NULL;
+  patch_data->resize(sz);
+  if (!android::base::ReadFully(patch_fd, patch_data->data(), sz)) {
+    printf("failed to read \"%s\" %s\n", ptemp, strerror(errno));
+    return false;
   }
-  fclose(f);
 
   unlink(ptemp);
+  tgt->SetSourceInfo(*src);
 
-  tgt->source_start = src->start;
-  switch (tgt->type) {
-    case CHUNK_NORMAL:
-      tgt->source_len = src->len;
-      break;
-    case CHUNK_DEFLATE:
-      tgt->source_len = src->deflate_len;
-      tgt->source_uncompressed_len = src->len;
-      break;
-  }
-
-  return data;
-}
-
-/*
- * Cause a gzip chunk to be treated as a normal chunk (ie, as a blob
- * of uninterpreted data).  The resulting patch will likely be about
- * as big as the target file, but it lets us handle the case of images
- * where some gzip chunks are reconstructible but others aren't (by
- * treating the ones that aren't as normal chunks).
- */
-void ChangeDeflateChunkToNormal(ImageChunk* ch) {
-  if (ch->type != CHUNK_DEFLATE) return;
-  ch->type = CHUNK_NORMAL;
-  free(ch->data);
-  ch->data = ch->deflate_data;
-  ch->len = ch->deflate_len;
-}
-
-/*
- * Return true if the data in the chunk is identical (including the
- * compressed representation, for gzip chunks).
- */
-int AreChunksEqual(ImageChunk* a, ImageChunk* b) {
-    if (a->type != b->type) return 0;
-
-    switch (a->type) {
-        case CHUNK_NORMAL:
-            return a->len == b->len && memcmp(a->data, b->data, a->len) == 0;
-
-        case CHUNK_DEFLATE:
-            return a->deflate_len == b->deflate_len &&
-                memcmp(a->deflate_data, b->deflate_data, a->deflate_len) == 0;
-
-        default:
-            printf("unknown chunk type %d\n", a->type);
-            return 0;
-    }
+  return true;
 }
 
 /*
@@ -740,71 +796,42 @@ int AreChunksEqual(ImageChunk* a, ImageChunk* b) {
  * a single chunk.  (Such runs can be produced when deflate chunks are
  * changed to normal chunks.)
  */
-void MergeAdjacentNormalChunks(ImageChunk* chunks, int* num_chunks) {
-  int out = 0;
-  int in_start = 0, in_end;
-  while (in_start < *num_chunks) {
-    if (chunks[in_start].type != CHUNK_NORMAL) {
-      in_end = in_start+1;
-    } else {
-      // in_start is a normal chunk.  Look for a run of normal chunks
-      // that constitute a solid block of data (ie, each chunk begins
-      // where the previous one ended).
-      for (in_end = in_start+1;
-           in_end < *num_chunks && chunks[in_end].type == CHUNK_NORMAL &&
-             (chunks[in_end].start ==
-              chunks[in_end-1].start + chunks[in_end-1].len &&
-              chunks[in_end].data ==
-              chunks[in_end-1].data + chunks[in_end-1].len);
-           ++in_end);
+static void MergeAdjacentNormalChunks(std::vector<ImageChunk>* chunks) {
+  size_t merged_last = 0, cur = 0;
+  while (cur < chunks->size()) {
+    // Look for normal chunks adjacent to the current one. If such chunk exists, extend the
+    // length of the current normal chunk.
+    size_t to_check = cur + 1;
+    while (to_check < chunks->size() && chunks->at(cur).IsAdjacentNormal(chunks->at(to_check))) {
+      chunks->at(cur).MergeAdjacentNormal(chunks->at(to_check));
+      to_check++;
     }
 
-    if (in_end == in_start+1) {
-#if 0
-      printf("chunk %d is now %d\n", in_start, out);
-#endif
-      if (out != in_start) {
-        memcpy(chunks+out, chunks+in_start, sizeof(ImageChunk));
-      }
-    } else {
-#if 0
-      printf("collapse normal chunks %d-%d into %d\n", in_start, in_end-1, out);
-#endif
-
-      // Merge chunks [in_start, in_end-1] into one chunk.  Since the
-      // data member of each chunk is just a pointer into an in-memory
-      // copy of the file, this can be done without recopying (the
-      // output chunk has the first chunk's start location and data
-      // pointer, and length equal to the sum of the input chunk
-      // lengths).
-      chunks[out].type = CHUNK_NORMAL;
-      chunks[out].start = chunks[in_start].start;
-      chunks[out].data = chunks[in_start].data;
-      chunks[out].len = chunks[in_end-1].len +
-        (chunks[in_end-1].start - chunks[in_start].start);
+    if (merged_last != cur) {
+      chunks->at(merged_last) = std::move(chunks->at(cur));
     }
-
-    ++out;
-    in_start = in_end;
+    merged_last++;
+    cur = to_check;
+  }
+  if (merged_last < chunks->size()) {
+    chunks->erase(chunks->begin() + merged_last, chunks->end());
   }
-  *num_chunks = out;
 }
 
-ImageChunk* FindChunkByName(const char* name, ImageChunk* chunks, int num_chunks) {
-  for (int i = 0; i < num_chunks; ++i) {
-    if (chunks[i].type == CHUNK_DEFLATE && chunks[i].filename &&
-        strcmp(name, chunks[i].filename) == 0) {
-      return chunks+i;
+static ImageChunk* FindChunkByName(const std::string& name, std::vector<ImageChunk>& chunks) {
+  for (size_t i = 0; i < chunks.size(); ++i) {
+    if (chunks[i].GetType() == CHUNK_DEFLATE && chunks[i].GetEntryName() == name) {
+      return &chunks[i];
     }
   }
-  return NULL;
+  return nullptr;
 }
 
-void DumpChunks(ImageChunk* chunks, int num_chunks) {
-    for (int i = 0; i < num_chunks; ++i) {
-        printf("chunk %d: type %d start %zu len %zu\n",
-               i, chunks[i].type, chunks[i].start, chunks[i].len);
-    }
+static void DumpChunks(const std::vector<ImageChunk>& chunks) {
+  for (size_t i = 0; i < chunks.size(); ++i) {
+    printf("chunk %zu: ", i);
+    chunks[i].Dump();
+  }
 }
 
 int imgdiff(int argc, const char** argv) {
@@ -816,26 +843,24 @@ int imgdiff(int argc, const char** argv) {
     ++argv;
   }
 
-  size_t bonus_size = 0;
-  unsigned char* bonus_data = NULL;
+  std::vector<uint8_t> bonus_data;
   if (argc >= 3 && strcmp(argv[1], "-b") == 0) {
     struct stat st;
     if (stat(argv[2], &st) != 0) {
       printf("failed to stat bonus file %s: %s\n", argv[2], strerror(errno));
       return 1;
     }
-    bonus_size = st.st_size;
-    bonus_data = static_cast<unsigned char*>(malloc(bonus_size));
-    FILE* f = fopen(argv[2], "rb");
-    if (f == NULL) {
+    size_t bonus_size = st.st_size;
+    bonus_data.resize(bonus_size);
+    android::base::unique_fd fd(open(argv[2], O_RDONLY));
+    if (fd == -1) {
       printf("failed to open bonus file %s: %s\n", argv[2], strerror(errno));
       return 1;
     }
-    if (fread(bonus_data, 1, bonus_size, f) != bonus_size) {
+    if (!android::base::ReadFully(fd, bonus_data.data(), bonus_size)) {
       printf("failed to read bonus file %s: %s\n", argv[2], strerror(errno));
       return 1;
     }
-    fclose(f);
 
     argc -= 2;
     argv += 2;
@@ -847,27 +872,26 @@ int imgdiff(int argc, const char** argv) {
     return 2;
   }
 
-  int num_src_chunks;
-  ImageChunk* src_chunks;
-  int num_tgt_chunks;
-  ImageChunk* tgt_chunks;
-  int i;
+  std::vector<ImageChunk> src_chunks;
+  std::vector<ImageChunk> tgt_chunks;
+  std::vector<uint8_t> src_file;
+  std::vector<uint8_t> tgt_file;
 
   if (zip_mode) {
-    if (ReadZip(argv[1], &num_src_chunks, &src_chunks, 1) == NULL) {
+    if (!ReadZip(argv[1], &src_chunks, &src_file, true)) {
       printf("failed to break apart source zip file\n");
       return 1;
     }
-    if (ReadZip(argv[2], &num_tgt_chunks, &tgt_chunks, 0) == NULL) {
+    if (!ReadZip(argv[2], &tgt_chunks, &tgt_file, false)) {
       printf("failed to break apart target zip file\n");
       return 1;
     }
   } else {
-    if (ReadImage(argv[1], &num_src_chunks, &src_chunks) == NULL) {
+    if (!ReadImage(argv[1], &src_chunks, &src_file)) {
       printf("failed to break apart source image\n");
       return 1;
     }
-    if (ReadImage(argv[2], &num_tgt_chunks, &tgt_chunks) == NULL) {
+    if (!ReadImage(argv[2], &tgt_chunks, &tgt_file)) {
       printf("failed to break apart target image\n");
       return 1;
     }
@@ -875,48 +899,47 @@ int imgdiff(int argc, const char** argv) {
     // Verify that the source and target images have the same chunk
     // structure (ie, the same sequence of deflate and normal chunks).
 
-    // Merge the gzip header and footer in with any adjacent
-    // normal chunks.
-    MergeAdjacentNormalChunks(tgt_chunks, &num_tgt_chunks);
-    MergeAdjacentNormalChunks(src_chunks, &num_src_chunks);
+    // Merge the gzip header and footer in with any adjacent normal chunks.
+    MergeAdjacentNormalChunks(&tgt_chunks);
+    MergeAdjacentNormalChunks(&src_chunks);
 
-    if (num_src_chunks != num_tgt_chunks) {
+    if (src_chunks.size() != tgt_chunks.size()) {
       printf("source and target don't have same number of chunks!\n");
       printf("source chunks:\n");
-      DumpChunks(src_chunks, num_src_chunks);
+      DumpChunks(src_chunks);
       printf("target chunks:\n");
-      DumpChunks(tgt_chunks, num_tgt_chunks);
+      DumpChunks(tgt_chunks);
       return 1;
     }
-    for (i = 0; i < num_src_chunks; ++i) {
-      if (src_chunks[i].type != tgt_chunks[i].type) {
-        printf("source and target don't have same chunk structure! (chunk %d)\n", i);
+    for (size_t i = 0; i < src_chunks.size(); ++i) {
+      if (src_chunks[i].GetType() != tgt_chunks[i].GetType()) {
+        printf("source and target don't have same chunk structure! (chunk %zu)\n", i);
         printf("source chunks:\n");
-        DumpChunks(src_chunks, num_src_chunks);
+        DumpChunks(src_chunks);
         printf("target chunks:\n");
-        DumpChunks(tgt_chunks, num_tgt_chunks);
+        DumpChunks(tgt_chunks);
         return 1;
       }
     }
   }
 
-  for (i = 0; i < num_tgt_chunks; ++i) {
-    if (tgt_chunks[i].type == CHUNK_DEFLATE) {
+  for (size_t i = 0; i < tgt_chunks.size(); ++i) {
+    if (tgt_chunks[i].GetType() == CHUNK_DEFLATE) {
       // Confirm that given the uncompressed chunk data in the target, we
       // can recompress it and get exactly the same bits as are in the
       // input target image.  If this fails, treat the chunk as a normal
       // non-deflated chunk.
-      if (ReconstructDeflateChunk(tgt_chunks+i) < 0) {
-        printf("failed to reconstruct target deflate chunk %d [%s]; "
-               "treating as normal\n", i, tgt_chunks[i].filename);
-        ChangeDeflateChunkToNormal(tgt_chunks+i);
+      if (!tgt_chunks[i].ReconstructDeflateChunk()) {
+        printf("failed to reconstruct target deflate chunk %zu [%s]; treating as normal\n", i,
+               tgt_chunks[i].GetEntryName().c_str());
+        tgt_chunks[i].ChangeDeflateChunkToNormal();
         if (zip_mode) {
-          ImageChunk* src = FindChunkByName(tgt_chunks[i].filename, src_chunks, num_src_chunks);
-          if (src) {
-            ChangeDeflateChunkToNormal(src);
+          ImageChunk* src = FindChunkByName(tgt_chunks[i].GetEntryName(), src_chunks);
+          if (src != nullptr) {
+            src->ChangeDeflateChunkToNormal();
           }
         } else {
-          ChangeDeflateChunkToNormal(src_chunks+i);
+          src_chunks[i].ChangeDeflateChunkToNormal();
         }
         continue;
       }
@@ -929,16 +952,16 @@ int imgdiff(int argc, const char** argv) {
       // data.
       ImageChunk* src;
       if (zip_mode) {
-        src = FindChunkByName(tgt_chunks[i].filename, src_chunks, num_src_chunks);
+        src = FindChunkByName(tgt_chunks[i].GetEntryName(), src_chunks);
       } else {
-        src = src_chunks+i;
+        src = &src_chunks[i];
       }
 
-      if (src == NULL || AreChunksEqual(tgt_chunks+i, src)) {
-        ChangeDeflateChunkToNormal(tgt_chunks+i);
-        if (src) {
-          ChangeDeflateChunkToNormal(src);
-        }
+      if (src == nullptr) {
+        tgt_chunks[i].ChangeDeflateChunkToNormal();
+      } else if (tgt_chunks[i] == *src) {
+        tgt_chunks[i].ChangeDeflateChunkToNormal();
+        src->ChangeDeflateChunkToNormal();
       }
     }
   }
@@ -948,14 +971,15 @@ int imgdiff(int argc, const char** argv) {
     // For zips, we only need to do this to the target:  deflated
     // chunks are matched via filename, and normal chunks are patched
     // using the entire source file as the source.
-    MergeAdjacentNormalChunks(tgt_chunks, &num_tgt_chunks);
+    MergeAdjacentNormalChunks(&tgt_chunks);
+
   } else {
     // For images, we need to maintain the parallel structure of the
     // chunk lists, so do the merging in both the source and target
     // lists.
-    MergeAdjacentNormalChunks(tgt_chunks, &num_tgt_chunks);
-    MergeAdjacentNormalChunks(src_chunks, &num_src_chunks);
-    if (num_src_chunks != num_tgt_chunks) {
+    MergeAdjacentNormalChunks(&tgt_chunks);
+    MergeAdjacentNormalChunks(&src_chunks);
+    if (src_chunks.size() != tgt_chunks.size()) {
       // This shouldn't happen.
       printf("merging normal chunks went awry\n");
       return 1;
@@ -965,117 +989,68 @@ int imgdiff(int argc, const char** argv) {
   // Compute bsdiff patches for each chunk's data (the uncompressed
   // data, in the case of deflate chunks).
 
-  DumpChunks(src_chunks, num_src_chunks);
+  DumpChunks(src_chunks);
 
-  printf("Construct patches for %d chunks...\n", num_tgt_chunks);
-  unsigned char** patch_data = static_cast<unsigned char**>(malloc(
-      num_tgt_chunks * sizeof(unsigned char*)));
-  size_t* patch_size = static_cast<size_t*>(malloc(num_tgt_chunks * sizeof(size_t)));
+  printf("Construct patches for %zu chunks...\n", tgt_chunks.size());
+  std::vector<std::vector<uint8_t>> patch_data(tgt_chunks.size());
   saidx_t* bsdiff_cache = nullptr;
-  for (i = 0; i < num_tgt_chunks; ++i) {
+  for (size_t i = 0; i < tgt_chunks.size(); ++i) {
     if (zip_mode) {
       ImageChunk* src;
-      if (tgt_chunks[i].type == CHUNK_DEFLATE &&
-          (src = FindChunkByName(tgt_chunks[i].filename, src_chunks, num_src_chunks))) {
-        patch_data[i] = MakePatch(src, tgt_chunks + i, patch_size + i, nullptr);
+      if (tgt_chunks[i].GetType() == CHUNK_DEFLATE &&
+          (src = FindChunkByName(tgt_chunks[i].GetEntryName(), src_chunks))) {
+        MakePatch(src, &tgt_chunks[i], &patch_data[i], nullptr);
       } else {
-        patch_data[i] = MakePatch(src_chunks, tgt_chunks + i, patch_size + i, &bsdiff_cache);
+        MakePatch(&src_chunks[0], &tgt_chunks[i], &patch_data[i], &bsdiff_cache);
       }
     } else {
-      if (i == 1 && bonus_data) {
-        printf("  using %zu bytes of bonus data for chunk %d\n", bonus_size, i);
-        src_chunks[i].data =
-            static_cast<unsigned char*>(realloc(src_chunks[i].data, src_chunks[i].len + bonus_size));
-        memcpy(src_chunks[i].data + src_chunks[i].len, bonus_data, bonus_size);
-        src_chunks[i].len += bonus_size;
+      if (i == 1 && !bonus_data.empty()) {
+        printf("  using %zu bytes of bonus data for chunk %zu\n", bonus_data.size(), i);
+        src_chunks[i].SetBonusData(bonus_data);
       }
 
-      patch_data[i] = MakePatch(src_chunks + i, tgt_chunks + i, patch_size + i, nullptr);
+      MakePatch(&src_chunks[i], &tgt_chunks[i], &patch_data[i], nullptr);
     }
-    printf("patch %3d is %zu bytes (of %zu)\n", i, patch_size[i], tgt_chunks[i].source_len);
+    printf("patch %3zu is %zu bytes (of %zu)\n", i, patch_data[i].size(),
+           src_chunks[i].GetRawDataLength());
+  }
+
+  if (bsdiff_cache != nullptr) {
+    free(bsdiff_cache);
   }
-  free(bsdiff_cache);
-  free(src_chunks);
 
   // Figure out how big the imgdiff file header is going to be, so
   // that we can correctly compute the offset of each bsdiff patch
   // within the file.
 
   size_t total_header_size = 12;
-  for (i = 0; i < num_tgt_chunks; ++i) {
-    total_header_size += 4;
-    switch (tgt_chunks[i].type) {
-      case CHUNK_NORMAL:
-        total_header_size += 8*3;
-        break;
-      case CHUNK_DEFLATE:
-        total_header_size += 8*5 + 4*5;
-        break;
-      case CHUNK_RAW:
-        total_header_size += 4 + patch_size[i];
-        break;
-    }
+  for (size_t i = 0; i < tgt_chunks.size(); ++i) {
+    total_header_size += tgt_chunks[i].GetHeaderSize(patch_data[i].size());
   }
 
   size_t offset = total_header_size;
 
   FILE* f = fopen(argv[3], "wb");
+  if (f == nullptr) {
+    printf("failed to open \"%s\": %s\n", argv[3], strerror(errno));
+  }
 
   // Write out the headers.
 
   fwrite("IMGDIFF2", 1, 8, f);
-  Write4(num_tgt_chunks, f);
-  for (i = 0; i < num_tgt_chunks; ++i) {
-    Write4(tgt_chunks[i].type, f);
-
-    switch (tgt_chunks[i].type) {
-      case CHUNK_NORMAL:
-        printf("chunk %3d: normal   (%10zu, %10zu)  %10zu\n", i,
-               tgt_chunks[i].start, tgt_chunks[i].len, patch_size[i]);
-        Write8(tgt_chunks[i].source_start, f);
-        Write8(tgt_chunks[i].source_len, f);
-        Write8(offset, f);
-        offset += patch_size[i];
-        break;
-
-      case CHUNK_DEFLATE:
-        printf("chunk %3d: deflate  (%10zu, %10zu)  %10zu  %s\n", i,
-               tgt_chunks[i].start, tgt_chunks[i].deflate_len, patch_size[i],
-               tgt_chunks[i].filename);
-        Write8(tgt_chunks[i].source_start, f);
-        Write8(tgt_chunks[i].source_len, f);
-        Write8(offset, f);
-        Write8(tgt_chunks[i].source_uncompressed_len, f);
-        Write8(tgt_chunks[i].len, f);
-        Write4(tgt_chunks[i].level, f);
-        Write4(tgt_chunks[i].method, f);
-        Write4(tgt_chunks[i].windowBits, f);
-        Write4(tgt_chunks[i].memLevel, f);
-        Write4(tgt_chunks[i].strategy, f);
-        offset += patch_size[i];
-        break;
-
-      case CHUNK_RAW:
-        printf("chunk %3d: raw      (%10zu, %10zu)\n", i,
-               tgt_chunks[i].start, tgt_chunks[i].len);
-        Write4(patch_size[i], f);
-        fwrite(patch_data[i], 1, patch_size[i], f);
-        break;
-    }
+  Write4(static_cast<int32_t>(tgt_chunks.size()), f);
+  for (size_t i = 0; i < tgt_chunks.size(); ++i) {
+    printf("chunk %zu: ", i);
+    offset = tgt_chunks[i].WriteHeaderToFile(f, patch_data[i], offset);
   }
 
   // Append each chunk's bsdiff patch, in order.
-
-  for (i = 0; i < num_tgt_chunks; ++i) {
-    if (tgt_chunks[i].type != CHUNK_RAW) {
-      fwrite(patch_data[i], 1, patch_size[i], f);
+  for (size_t i = 0; i < tgt_chunks.size(); ++i) {
+    if (tgt_chunks[i].GetType() != CHUNK_RAW) {
+      fwrite(patch_data[i].data(), 1, patch_data[i].size(), f);
     }
   }
 
-  free(tgt_chunks);
-  free(patch_data);
-  free(patch_size);
-
   fclose(f);
 
   return 0;
diff --git a/tests/Android.mk b/tests/Android.mk
index c7d33256a..ec971b38c 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -177,6 +177,7 @@ LOCAL_STATIC_LIBRARIES := \
     libbsdiff \
     libbspatch \
     libziparchive \
+    libutils \
     libbase \
     libcrypto \
     libbz \
diff --git a/tests/component/imgdiff_test.cpp b/tests/component/imgdiff_test.cpp
index 7ad330783..be2dd385b 100644
--- a/tests/component/imgdiff_test.cpp
+++ b/tests/component/imgdiff_test.cpp
@@ -236,6 +236,62 @@ TEST(ImgdiffTest, zip_mode_smoke_compressed) {
   ASSERT_EQ(tgt, patched);
 }
 
+TEST(ImgdiffTest, zip_mode_smoke_trailer_zeros) {
+  // Construct src and tgt zip files.
+  TemporaryFile src_file;
+  FILE* src_file_ptr = fdopen(src_file.fd, "wb");
+  ZipWriter src_writer(src_file_ptr);
+  ASSERT_EQ(0, src_writer.StartEntry("file1.txt", ZipWriter::kCompress));
+  const std::string src_content("abcdefg");
+  ASSERT_EQ(0, src_writer.WriteBytes(src_content.data(), src_content.size()));
+  ASSERT_EQ(0, src_writer.FinishEntry());
+  ASSERT_EQ(0, src_writer.Finish());
+  ASSERT_EQ(0, fclose(src_file_ptr));
+
+  TemporaryFile tgt_file;
+  FILE* tgt_file_ptr = fdopen(tgt_file.fd, "wb");
+  ZipWriter tgt_writer(tgt_file_ptr);
+  ASSERT_EQ(0, tgt_writer.StartEntry("file1.txt", ZipWriter::kCompress));
+  const std::string tgt_content("abcdefgxyz");
+  ASSERT_EQ(0, tgt_writer.WriteBytes(tgt_content.data(), tgt_content.size()));
+  ASSERT_EQ(0, tgt_writer.FinishEntry());
+  ASSERT_EQ(0, tgt_writer.Finish());
+  // Add trailing zeros to the target zip file.
+  std::vector<uint8_t> zeros(10);
+  ASSERT_EQ(zeros.size(), fwrite(zeros.data(), sizeof(uint8_t), zeros.size(), tgt_file_ptr));
+  ASSERT_EQ(0, fclose(tgt_file_ptr));
+
+  // Compute patch.
+  TemporaryFile patch_file;
+  std::vector<const char*> args = {
+    "imgdiff", "-z", src_file.path, tgt_file.path, patch_file.path,
+  };
+  ASSERT_EQ(0, imgdiff(args.size(), args.data()));
+
+  // Verify.
+  std::string tgt;
+  ASSERT_TRUE(android::base::ReadFileToString(tgt_file.path, &tgt));
+  std::string src;
+  ASSERT_TRUE(android::base::ReadFileToString(src_file.path, &src));
+  std::string patch;
+  ASSERT_TRUE(android::base::ReadFileToString(patch_file.path, &patch));
+
+  // Expect three entries: CHUNK_RAW (header) + CHUNK_DEFLATE (data) + CHUNK_RAW (footer).
+  size_t num_normal;
+  size_t num_raw;
+  size_t num_deflate;
+  verify_patch_header(patch, &num_normal, &num_raw, &num_deflate);
+  ASSERT_EQ(0U, num_normal);
+  ASSERT_EQ(1U, num_deflate);
+  ASSERT_EQ(2U, num_raw);
+
+  std::string patched;
+  ASSERT_EQ(0, ApplyImagePatch(reinterpret_cast<const unsigned char*>(src.data()), src.size(),
+                               reinterpret_cast<const unsigned char*>(patch.data()), patch.size(),
+                               MemorySink, &patched));
+  ASSERT_EQ(tgt, patched);
+}
+
 TEST(ImgdiffTest, image_mode_simple) {
   // src: "abcdefgh" + gzipped "xyz" (echo -n "xyz" | gzip -f | hd).
   const std::vector<char> src_data = { 'a',    'b',    'c',    'd',    'e',    'f',    'g',