4 files changed, 33 insertions, 25 deletions

diff --git a/applypatch/imgdiff.cpp b/applypatch/imgdiff.cpp
index 69ad75f37..f57e7942c 100644
--- a/applypatch/imgdiff.cpp
+++ b/applypatch/imgdiff.cpp
@@ -163,7 +163,7 @@
 #include <android-base/stringprintf.h>
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
-#include <bsdiff.h>
+#include <bsdiff/bsdiff.h>
 #include <ziparchive/zip_archive.h>
 #include <zlib.h>
 
@@ -324,7 +324,8 @@ void ImageChunk::MergeAdjacentNormal(const ImageChunk& other) {
 }
 
 bool ImageChunk::MakePatch(const ImageChunk& tgt, const ImageChunk& src,
-                           std::vector<uint8_t>* patch_data, saidx_t** bsdiff_cache) {
+                           std::vector<uint8_t>* patch_data,
+                           bsdiff::SuffixArrayIndexInterface** bsdiff_cache) {
 #if defined(__ANDROID__)
   char ptemp[] = "/data/local/tmp/imgdiff-patch-XXXXXX";
 #else
@@ -1083,7 +1084,7 @@ bool ZipModeImage::GeneratePatchesInternal(const ZipModeImage& tgt_image,
   printf("Construct patches for %zu chunks...\n", tgt_image.NumOfChunks());
   patch_chunks->clear();
 
-  saidx_t* bsdiff_cache = nullptr;
+  bsdiff::SuffixArrayIndexInterface* bsdiff_cache = nullptr;
   for (size_t i = 0; i < tgt_image.NumOfChunks(); i++) {
     const auto& tgt_chunk = tgt_image[i];
 
@@ -1097,7 +1098,8 @@ bool ZipModeImage::GeneratePatchesInternal(const ZipModeImage& tgt_image,
                                       : src_image.FindChunkByName(tgt_chunk.GetEntryName());
 
     const auto& src_ref = (src_chunk == nullptr) ? src_image.PseudoSource() : *src_chunk;
-    saidx_t** bsdiff_cache_ptr = (src_chunk == nullptr) ? &bsdiff_cache : nullptr;
+    bsdiff::SuffixArrayIndexInterface** bsdiff_cache_ptr =
+        (src_chunk == nullptr) ? &bsdiff_cache : nullptr;
 
     std::vector<uint8_t> patch_data;
     if (!ImageChunk::MakePatch(tgt_chunk, src_ref, &patch_data, bsdiff_cache_ptr)) {
@@ -1114,7 +1116,7 @@ bool ZipModeImage::GeneratePatchesInternal(const ZipModeImage& tgt_image,
       patch_chunks->emplace_back(tgt_chunk, src_ref, std::move(patch_data));
     }
   }
-  free(bsdiff_cache);
+  delete bsdiff_cache;
 
   CHECK_EQ(patch_chunks->size(), tgt_image.NumOfChunks());
   return true;
diff --git a/applypatch/imgpatch.cpp b/applypatch/imgpatch.cpp
index 7a43ddbef..25ba0a182 100644
--- a/applypatch/imgpatch.cpp
+++ b/applypatch/imgpatch.cpp
@@ -59,13 +59,13 @@ static bool ApplyBSDiffPatchAndStreamOutput(const uint8_t* src_data, size_t src_
   int mem_level = Read4(deflate_header + 52);
   int strategy = Read4(deflate_header + 56);
 
-  std::unique_ptr<z_stream, decltype(&deflateEnd)> strm(new z_stream(), deflateEnd);
-  strm->zalloc = Z_NULL;
-  strm->zfree = Z_NULL;
-  strm->opaque = Z_NULL;
-  strm->avail_in = 0;
-  strm->next_in = nullptr;
-  int ret = deflateInit2(strm.get(), level, method, window_bits, mem_level, strategy);
+  z_stream strm;
+  strm.zalloc = Z_NULL;
+  strm.zfree = Z_NULL;
+  strm.opaque = Z_NULL;
+  strm.avail_in = 0;
+  strm.next_in = nullptr;
+  int ret = deflateInit2(&strm, level, method, window_bits, mem_level, strategy);
   if (ret != Z_OK) {
     LOG(ERROR) << "Failed to init uncompressed data deflation: " << ret;
     return false;
@@ -76,18 +76,19 @@ static bool ApplyBSDiffPatchAndStreamOutput(const uint8_t* src_data, size_t src_
   size_t actual_target_length = 0;
   size_t total_written = 0;
   static constexpr size_t buffer_size = 32768;
-  auto compression_sink = [&](const uint8_t* data, size_t len) -> size_t {
+  auto compression_sink = [&strm, &actual_target_length, &expected_target_length, &total_written,
+                           &ret, &ctx, &sink](const uint8_t* data, size_t len) -> size_t {
     // The input patch length for an update never exceeds INT_MAX.
-    strm->avail_in = len;
-    strm->next_in = data;
+    strm.avail_in = len;
+    strm.next_in = data;
     do {
       std::vector<uint8_t> buffer(buffer_size);
-      strm->avail_out = buffer_size;
-      strm->next_out = buffer.data();
+      strm.avail_out = buffer_size;
+      strm.next_out = buffer.data();
       if (actual_target_length + len < expected_target_length) {
-        ret = deflate(strm.get(), Z_NO_FLUSH);
+        ret = deflate(&strm, Z_NO_FLUSH);
       } else {
-        ret = deflate(strm.get(), Z_FINISH);
+        ret = deflate(&strm, Z_FINISH);
       }
       if (ret != Z_OK && ret != Z_STREAM_END) {
         LOG(ERROR) << "Failed to deflate stream: " << ret;
@@ -95,20 +96,24 @@ static bool ApplyBSDiffPatchAndStreamOutput(const uint8_t* src_data, size_t src_
         return 0;
       }
 
-      size_t have = buffer_size - strm->avail_out;
+      size_t have = buffer_size - strm.avail_out;
       total_written += have;
       if (sink(buffer.data(), have) != have) {
         LOG(ERROR) << "Failed to write " << have << " compressed bytes to output.";
         return 0;
       }
       if (ctx) SHA1_Update(ctx, buffer.data(), have);
-    } while ((strm->avail_in != 0 || strm->avail_out == 0) && ret != Z_STREAM_END);
+    } while ((strm.avail_in != 0 || strm.avail_out == 0) && ret != Z_STREAM_END);
 
     actual_target_length += len;
     return len;
   };
 
-  if (ApplyBSDiffPatch(src_data, src_len, patch, patch_offset, compression_sink, nullptr) != 0) {
+  int bspatch_result =
+      ApplyBSDiffPatch(src_data, src_len, patch, patch_offset, compression_sink, nullptr);
+  deflateEnd(&strm);
+
+  if (bspatch_result != 0) {
     return false;
   }
 
diff --git a/applypatch/include/applypatch/imgdiff_image.h b/applypatch/include/applypatch/imgdiff_image.h
index 3d29547cb..00a84f3a9 100644
--- a/applypatch/include/applypatch/imgdiff_image.h
+++ b/applypatch/include/applypatch/imgdiff_image.h
@@ -24,7 +24,7 @@
 #include <string>
 #include <vector>
 
-#include <bsdiff.h>
+#include <bsdiff/bsdiff.h>
 #include <ziparchive/zip_archive.h>
 #include <zlib.h>
 
@@ -98,7 +98,8 @@ class ImageChunk {
    * repeatedly, pass nullptr if not needed.
    */
   static bool MakePatch(const ImageChunk& tgt, const ImageChunk& src,
-                        std::vector<uint8_t>* patch_data, saidx_t** bsdiff_cache);
+                        std::vector<uint8_t>* patch_data,
+                        bsdiff::SuffixArrayIndexInterface** bsdiff_cache);
 
  private:
   const uint8_t* GetRawData() const;
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index 266657d31..d9d01d427 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -33,7 +33,7 @@
 #include <android-base/test_utils.h>
 #include <bootloader_message/bootloader_message.h>
 #include <brotli/encode.h>
-#include <bsdiff.h>
+#include <bsdiff/bsdiff.h>
 #include <gtest/gtest.h>
 #include <ziparchive/zip_archive.h>
 #include <ziparchive/zip_writer.h>