4 files changed, 71 insertions, 22 deletions

diff --git a/applypatch/bspatch.cpp b/applypatch/bspatch.cpp
index f75a2c680..65ee614ef 100644
--- a/applypatch/bspatch.cpp
+++ b/applypatch/bspatch.cpp
@@ -23,10 +23,14 @@
 #include <stdio.h>
 #include <sys/types.h>
 
+#include <string>
+
+#include <android-base/logging.h>
 #include <bspatch.h>
 #include <openssl/sha.h>
 
 #include "applypatch/applypatch.h"
+#include "print_sha1.h"
 
 void ShowBSDiffLicense() {
     puts("The bsdiff library used herein is:\n"
@@ -67,18 +71,24 @@ int ApplyBSDiffPatch(const unsigned char* old_data, size_t old_size, const Value
     if (ctx) SHA1_Update(ctx, data, len);
     return len;
   };
-  return bsdiff::bspatch(old_data, old_size,
-                         reinterpret_cast<const uint8_t*>(&patch->data[patch_offset]),
-                         patch->data.size(), sha_sink);
-}
 
-int ApplyBSDiffPatchMem(const unsigned char* old_data, size_t old_size, const Value* patch,
-                        size_t patch_offset, std::vector<unsigned char>* new_data) {
-  auto vector_sink = [new_data](const uint8_t* data, size_t len) {
-    new_data->insert(new_data->end(), data, data + len);
-    return len;
-  };
-  return bsdiff::bspatch(old_data, old_size,
-                         reinterpret_cast<const uint8_t*>(&patch->data[patch_offset]),
-                         patch->data.size(), vector_sink);
-}
+  CHECK(patch != nullptr);
+  CHECK_LE(patch_offset, patch->data.size());
+
+  int result = bsdiff::bspatch(old_data, old_size,
+                               reinterpret_cast<const uint8_t*>(&patch->data[patch_offset]),
+                               patch->data.size() - patch_offset, sha_sink);
+  if (result != 0) {
+    LOG(ERROR) << "bspatch failed, result: " << result;
+    // print SHA1 of the patch in the case of a data error.
+    if (result == 2) {
+      uint8_t digest[SHA_DIGEST_LENGTH];
+      SHA1(reinterpret_cast<const uint8_t*>(patch->data.data() + patch_offset),
+           patch->data.size() - patch_offset, digest);
+      std::string patch_sha1 = print_sha1(digest);
+      LOG(ERROR) << "Patch may be corrupted, offset: " << patch_offset << ", SHA1: "
+                 << patch_sha1;
+    }
+  }
+  return result;
+}
\ No newline at end of file
diff --git a/applypatch/imgpatch.cpp b/applypatch/imgpatch.cpp
index 3d905f7b4..702a624ae 100644
--- a/applypatch/imgpatch.cpp
+++ b/applypatch/imgpatch.cpp
@@ -200,12 +200,14 @@ int ApplyImagePatch(const unsigned char* old_data, size_t old_size, const Value*
       }
 
       // Next, apply the bsdiff patch (in memory) to the uncompressed data.
-      std::vector<unsigned char> uncompressed_target_data;
-      // TODO(senj): Remove the only usage of ApplyBSDiffPatchMem here,
-      // replace it with ApplyBSDiffPatch with a custom sink function that
-      // wraps the given sink function to stream output to save memory.
-      if (ApplyBSDiffPatchMem(expanded_source.data(), expanded_len, patch, patch_offset,
-                              &uncompressed_target_data) != 0) {
+      std::vector<uint8_t> uncompressed_target_data;
+      // TODO: replace the custom sink function passed into ApplyBSDiffPatch so that it wraps the
+      // given sink function to stream output to save memory.
+      if (ApplyBSDiffPatch(expanded_source.data(), expanded_len, patch, patch_offset,
+        [&uncompressed_target_data](const uint8_t* data, size_t len) {
+          uncompressed_target_data.insert(uncompressed_target_data.end(), data, data + len);
+          return len;
+        }, nullptr) != 0) {
         return -1;
       }
       if (uncompressed_target_data.size() != target_len) {
diff --git a/applypatch/include/applypatch/applypatch.h b/applypatch/include/applypatch/applypatch.h
index da55432d5..581360ef1 100644
--- a/applypatch/include/applypatch/applypatch.h
+++ b/applypatch/include/applypatch/applypatch.h
@@ -69,8 +69,6 @@ int SaveFileContents(const char* filename, const FileContents* file);
 void ShowBSDiffLicense();
 int ApplyBSDiffPatch(const unsigned char* old_data, size_t old_size, const Value* patch,
                      size_t patch_offset, SinkFn sink, SHA_CTX* ctx);
-int ApplyBSDiffPatchMem(const unsigned char* old_data, size_t old_size, const Value* patch,
-                        size_t patch_offset, std::vector<unsigned char>* new_data);
 
 // imgpatch.cpp
 int ApplyImagePatch(const unsigned char* old_data, size_t old_size, const Value* patch, SinkFn sink,
diff --git a/tests/component/imgdiff_test.cpp b/tests/component/imgdiff_test.cpp
index 7d00a3d53..6f5960bbd 100644
--- a/tests/component/imgdiff_test.cpp
+++ b/tests/component/imgdiff_test.cpp
@@ -551,3 +551,42 @@ TEST(ImgdiffTest, image_mode_single_entry_long) {
 
   verify_patched_image(src, patch, tgt);
 }
+
+TEST(ImgpatchTest, image_mode_patch_corruption) {
+  // src: "abcdefgh" + gzipped "xyz" (echo -n "xyz" | gzip -f | hd).
+  const std::vector<char> src_data = { 'a',    'b',    'c',    'd',    'e',    'f',    'g',
+                                       'h',    '\x1f', '\x8b', '\x08', '\x00', '\xc4', '\x1e',
+                                       '\x53', '\x58', '\x00', '\x03', '\xab', '\xa8', '\xac',
+                                       '\x02', '\x00', '\x67', '\xba', '\x8e', '\xeb', '\x03',
+                                       '\x00', '\x00', '\x00' };
+  const std::string src(src_data.cbegin(), src_data.cend());
+  TemporaryFile src_file;
+  ASSERT_TRUE(android::base::WriteStringToFile(src, src_file.path));
+
+  // tgt: "abcdefgxyz" + gzipped "xxyyzz".
+  const std::vector<char> tgt_data = {
+    'a',    'b',    'c',    'd',    'e',    'f',    'g',    'x',    'y',    'z',    '\x1f', '\x8b',
+    '\x08', '\x00', '\x62', '\x1f', '\x53', '\x58', '\x00', '\x03', '\xab', '\xa8', '\xa8', '\xac',
+    '\xac', '\xaa', '\x02', '\x00', '\x96', '\x30', '\x06', '\xb7', '\x06', '\x00', '\x00', '\x00'
+  };
+  const std::string tgt(tgt_data.cbegin(), tgt_data.cend());
+  TemporaryFile tgt_file;
+  ASSERT_TRUE(android::base::WriteStringToFile(tgt, tgt_file.path));
+
+  TemporaryFile patch_file;
+  std::vector<const char*> args = {
+    "imgdiff", src_file.path, tgt_file.path, patch_file.path,
+  };
+  ASSERT_EQ(0, imgdiff(args.size(), args.data()));
+
+  // Verify.
+  std::string patch;
+  ASSERT_TRUE(android::base::ReadFileToString(patch_file.path, &patch));
+  verify_patched_image(src, patch, tgt);
+
+  // Corrupt the end of the patch and expect the ApplyImagePatch to fail.
+  patch.insert(patch.end() - 10, 10, '0');
+  ASSERT_EQ(-1, ApplyImagePatch(reinterpret_cast<const unsigned char*>(src.data()), src.size(),
+                                reinterpret_cast<const unsigned char*>(patch.data()), patch.size(),
+                                [](const unsigned char* /*data*/, size_t len) { return len; }));
+}