8 files changed, 98 insertions, 48 deletions

diff --git a/Android.bp b/Android.bp
index f919ebc83..49438ad9e 100644
--- a/Android.bp
+++ b/Android.bp
@@ -1,3 +1,4 @@
 subdirs = [
+    "bootloader_message",
     "otautil",
 ]
diff --git a/Android.mk b/Android.mk
index 1f69d5d5a..e619db031 100644
--- a/Android.mk
+++ b/Android.mk
@@ -185,7 +185,6 @@ include $(BUILD_STATIC_LIBRARY)
 include \
     $(LOCAL_PATH)/applypatch/Android.mk \
     $(LOCAL_PATH)/boot_control/Android.mk \
-    $(LOCAL_PATH)/bootloader_message/Android.mk \
     $(LOCAL_PATH)/edify/Android.mk \
     $(LOCAL_PATH)/minadbd/Android.mk \
     $(LOCAL_PATH)/minui/Android.mk \
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/bootloader_message/Android.bp b/bootloader_message/Android.bp
new file mode 100644
index 000000000..f0d76e718
--- /dev/null
+++ b/bootloader_message/Android.bp
@@ -0,0 +1,26 @@
+//
+// Copyright (C) 2017 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+cc_library_static {
+    name: "libbootloader_message",
+    srcs: ["bootloader_message.cpp"],
+    cppflags: ["-Werror"],
+    static_libs: [
+        "libbase",
+        "libfs_mgr",
+    ],
+    export_include_dirs: ["include"],
+}
diff --git a/bootloader_message/Android.mk b/bootloader_message/Android.mk
deleted file mode 100644
index a8c50819b..000000000
--- a/bootloader_message/Android.mk
+++ /dev/null
@@ -1,25 +0,0 @@
-# Copyright (C) 2016 The Android Open Source Project
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-LOCAL_PATH := $(call my-dir)
-
-include $(CLEAR_VARS)
-LOCAL_CLANG := true
-LOCAL_SRC_FILES := bootloader_message.cpp
-LOCAL_MODULE := libbootloader_message
-LOCAL_STATIC_LIBRARIES := libbase libfs_mgr
-LOCAL_CFLAGS := -Werror
-LOCAL_C_INCLUDES := $(LOCAL_PATH)/include
-LOCAL_EXPORT_C_INCLUDE_DIRS := $(LOCAL_PATH)/include
-include $(BUILD_STATIC_LIBRARY)
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; }));
+}