3 files changed, 91 insertions, 15 deletions

diff --git a/applypatch/Android.mk b/applypatch/Android.mk
index 61e110617..ec3c6ee38 100644
--- a/applypatch/Android.mk
+++ b/applypatch/Android.mk
@@ -124,6 +124,7 @@ libimgdiff_cflags := \
 
 libimgdiff_static_libraries := \
     libbsdiff \
+    libbase \
     libz
 
 # libimgdiff (static library)
diff --git a/applypatch/imgdiff.cpp b/applypatch/imgdiff.cpp
index 18a15a164..62de726a4 100644
--- a/applypatch/imgdiff.cpp
+++ b/applypatch/imgdiff.cpp
@@ -124,6 +124,7 @@
 #include "applypatch/imgdiff.h"
 
 #include <errno.h>
+#include <fcntl.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -131,6 +132,9 @@
 #include <sys/types.h>
 #include <unistd.h>
 
+#include <android-base/file.h>
+#include <android-base/unique_fd.h>
+
 #include <bsdiff.h>
 #include <zlib.h>
 
@@ -382,19 +386,12 @@ unsigned char* ReadImage(const char* filename, int* num_chunks, ImageChunk** chu
   }
 
   size_t sz = static_cast<size_t>(st.st_size);
-  unsigned char* img = static_cast<unsigned char*>(malloc(sz + 4));
-  FILE* f = fopen(filename, "rb");
-  if (fread(img, 1, sz, f) != sz) {
+  unsigned char* img = static_cast<unsigned char*>(malloc(sz));
+  android::base::unique_fd fd(open(filename, O_RDONLY));
+  if (!android::base::ReadFully(fd, img, sz)) {
     printf("failed to read \"%s\" %s\n", filename, strerror(errno));
-    fclose(f);
-    return NULL;
+    return nullptr;
   }
-  fclose(f);
-
-  // append 4 zero bytes to the data so we can always search for the
-  // four-byte string 1f8b0800 starting at any point in the actual
-  // file data, without special-casing the end of the data.
-  memset(img+sz, 0, 4);
 
   size_t pos = 0;
 
@@ -518,10 +515,8 @@ unsigned char* ReadImage(const char* filename, int* num_chunks, ImageChunk** chu
       curr->data = p;
 
       for (curr->len = 0; curr->len < (sz - pos); ++curr->len) {
-        if (p[curr->len] == 0x1f &&
-            p[curr->len+1] == 0x8b &&
-            p[curr->len+2] == 0x08 &&
-            p[curr->len+3] == 0x00) {
+        if (sz - pos >= 4 && p[curr->len] == 0x1f && p[curr->len + 1] == 0x8b &&
+            p[curr->len + 2] == 0x08 && p[curr->len + 3] == 0x00) {
           break;
         }
       }
diff --git a/tests/component/imgdiff_test.cpp b/tests/component/imgdiff_test.cpp
index 3711859de..7ad330783 100644
--- a/tests/component/imgdiff_test.cpp
+++ b/tests/component/imgdiff_test.cpp
@@ -404,6 +404,86 @@ TEST(ImgdiffTest, image_mode_spurious_magic) {
   ASSERT_EQ(tgt, patched);
 }
 
+TEST(ImgdiffTest, image_mode_short_input1) {
+  // src: "abcdefgh" + '0x1f8b0b'.
+  const std::vector<char> src_data = { 'a', 'b', 'c',    'd',    'e',   'f',
+                                       'g', 'h', '\x1f', '\x8b', '\x08' };
+  const std::string src(src_data.cbegin(), src_data.cend());
+  TemporaryFile src_file;
+  ASSERT_TRUE(android::base::WriteStringToFile(src, src_file.path));
+
+  // tgt: "abcdefgxyz".
+  const std::vector<char> tgt_data = { 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'x', 'y', 'z' };
+  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));
+
+  // Expect one CHUNK_RAW (header) entry.
+  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(0U, num_deflate);
+  ASSERT_EQ(1U, 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_short_input2) {
+  // src: "abcdefgh" + '0x1f8b0b00'.
+  const std::vector<char> src_data = { 'a', 'b', 'c',    'd',    'e',    'f',
+                                       'g', 'h', '\x1f', '\x8b', '\x08', '\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".
+  const std::vector<char> tgt_data = { 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'x', 'y', 'z' };
+  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));
+
+  // Expect one CHUNK_RAW (header) entry.
+  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(0U, num_deflate);
+  ASSERT_EQ(1U, 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_single_entry_long) {
   // src: "abcdefgh" + '0x1f8b0b00' + some bytes.
   const std::vector<char> src_data = { 'a',    'b',    'c',    'd',    'e',    'f',    'g',