5 files changed, 140 insertions, 92 deletions

diff --git a/tests/Android.mk b/tests/Android.mk
index 31c7de17..b0f71a83 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -201,6 +201,7 @@ LOCAL_SRC_FILES := \
 LOCAL_STATIC_LIBRARIES := \
     libimgdiff \
     libimgpatch \
+    libotautil \
     libbsdiff \
     libbspatch \
     libziparchive \
diff --git a/tests/component/applypatch_test.cpp b/tests/component/applypatch_test.cpp
index 42542898..15ec08fe 100644
--- a/tests/component/applypatch_test.cpp
+++ b/tests/component/applypatch_test.cpp
@@ -35,7 +35,7 @@
 #include "applypatch/applypatch.h"
 #include "applypatch/applypatch_modes.h"
 #include "common/test_constants.h"
-#include "print_sha1.h"
+#include "otautil/print_sha1.h"
 
 static void sha1sum(const std::string& fname, std::string* sha1, size_t* fsize = nullptr) {
   ASSERT_NE(nullptr, sha1);
diff --git a/tests/component/imgdiff_test.cpp b/tests/component/imgdiff_test.cpp
index 161d58d4..bf591dae 100644
--- a/tests/component/imgdiff_test.cpp
+++ b/tests/component/imgdiff_test.cpp
@@ -24,6 +24,7 @@
 #include <android-base/file.h>
 #include <android-base/memory.h>
 #include <android-base/stringprintf.h>
+#include <android-base/strings.h>
 #include <android-base/test_utils.h>
 #include <applypatch/imgdiff.h>
 #include <applypatch/imgdiff_image.h>
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index 19653ec5..d9d01d42 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -23,6 +23,7 @@
 #include <algorithm>
 #include <memory>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include <android-base/file.h>
@@ -39,9 +40,9 @@
 
 #include "common/test_constants.h"
 #include "edify/expr.h"
-#include "error_code.h"
 #include "otautil/SysUtil.h"
-#include "print_sha1.h"
+#include "otautil/error_code.h"
+#include "otautil/print_sha1.h"
 #include "updater/blockimg.h"
 #include "updater/install.h"
 #include "updater/updater.h"
@@ -74,6 +75,23 @@ static void expect(const char* expected, const char* expr_str, CauseCode cause_c
   ASSERT_EQ(cause_code, state.cause_code);
 }
 
+static void BuildUpdatePackage(const std::unordered_map<std::string, std::string>& entries,
+                               int fd) {
+  FILE* zip_file_ptr = fdopen(fd, "wb");
+  ZipWriter zip_writer(zip_file_ptr);
+
+  for (const auto& entry : entries) {
+    ASSERT_EQ(0, zip_writer.StartEntry(entry.first.c_str(), 0));
+    if (!entry.second.empty()) {
+      ASSERT_EQ(0, zip_writer.WriteBytes(entry.second.data(), entry.second.size()));
+    }
+    ASSERT_EQ(0, zip_writer.FinishEntry());
+  }
+
+  ASSERT_EQ(0, zip_writer.Finish());
+  ASSERT_EQ(0, fclose(zip_file_ptr));
+}
+
 static std::string get_sha1(const std::string& content) {
   uint8_t digest[SHA_DIGEST_LENGTH];
   SHA1(reinterpret_cast<const uint8_t*>(content.c_str()), content.size(), digest);
@@ -420,30 +438,19 @@ TEST_F(UpdaterTest, show_progress) {
   ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
 }
 
-TEST_F(UpdaterTest, block_image_update) {
-  // Create a zip file with new_data and patch_data.
-  TemporaryFile zip_file;
-  FILE* zip_file_ptr = fdopen(zip_file.release(), "wb");
-  ZipWriter zip_writer(zip_file_ptr);
-
-  // Add a dummy new data.
-  ASSERT_EQ(0, zip_writer.StartEntry("new_data", 0));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-
-  // Generate and add the patch data.
+TEST_F(UpdaterTest, block_image_update_patch_data) {
   std::string src_content = std::string(4096, 'a') + std::string(4096, 'c');
   std::string tgt_content = std::string(4096, 'b') + std::string(4096, 'd');
+
+  // Generate the patch data.
   TemporaryFile patch_file;
   ASSERT_EQ(0, bsdiff::bsdiff(reinterpret_cast<const uint8_t*>(src_content.data()),
       src_content.size(), reinterpret_cast<const uint8_t*>(tgt_content.data()),
       tgt_content.size(), patch_file.path, nullptr));
   std::string patch_content;
   ASSERT_TRUE(android::base::ReadFileToString(patch_file.path, &patch_content));
-  ASSERT_EQ(0, zip_writer.StartEntry("patch_data", 0));
-  ASSERT_EQ(0, zip_writer.WriteBytes(patch_content.data(), patch_content.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
 
-  // Add two transfer lists. The first one contains a bsdiff; and we expect the update to succeed.
+  // Create the transfer list that contains a bsdiff.
   std::string src_hash = get_sha1(src_content);
   std::string tgt_hash = get_sha1(tgt_content);
   std::vector<std::string> transfer_list = {
@@ -456,27 +463,16 @@ TEST_F(UpdaterTest, block_image_update) {
                                 src_hash.c_str(), tgt_hash.c_str(), src_hash.c_str()),
     "free " + src_hash,
   };
-  ASSERT_EQ(0, zip_writer.StartEntry("transfer_list", 0));
-  std::string commands = android::base::Join(transfer_list, '\n');
-  ASSERT_EQ(0, zip_writer.WriteBytes(commands.data(), commands.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
 
-  // Stash and free some blocks, then fail the 2nd update intentionally.
-  std::vector<std::string> fail_transfer_list = {
-    "4",
-    "2",
-    "0",
-    "2",
-    "stash " + tgt_hash + " 2,0,2",
-    "free " + tgt_hash,
-    "fail",
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", "" },
+    { "patch_data", patch_content },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
   };
-  ASSERT_EQ(0, zip_writer.StartEntry("fail_transfer_list", 0));
-  std::string fail_commands = android::base::Join(fail_transfer_list, '\n');
-  ASSERT_EQ(0, zip_writer.WriteBytes(fail_commands.data(), fail_commands.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  ASSERT_EQ(0, zip_writer.Finish());
-  ASSERT_EQ(0, fclose(zip_file_ptr));
+
+  // Build the update package.
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
 
   MemMapping map;
   ASSERT_TRUE(map.MapFile(zip_file.path));
@@ -491,7 +487,7 @@ TEST_F(UpdaterTest, block_image_update) {
   updater_info.package_zip_addr = map.addr;
   updater_info.package_zip_len = map.length;
 
-  // Execute the commands in the 1st transfer list.
+  // Execute the commands in the transfer list.
   TemporaryFile update_file;
   ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
   std::string script = "block_image_update(\"" + std::string(update_file.path) +
@@ -502,44 +498,98 @@ TEST_F(UpdaterTest, block_image_update) {
   ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_content));
   ASSERT_EQ(tgt_hash, get_sha1(updated_content));
 
-  // Expect the 2nd update to fail, but expect the stashed blocks to be freed.
-  script = "block_image_update(\"" + std::string(update_file.path) +
-      R"(", package_extract_file("fail_transfer_list"), "new_data", "patch_data"))";
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
+  CloseArchive(handle);
+}
+
+TEST_F(UpdaterTest, block_image_update_fail) {
+  std::string src_content(4096 * 2, 'e');
+  std::string src_hash = get_sha1(src_content);
+  // Stash and free some blocks, then fail the update intentionally.
+  std::vector<std::string> transfer_list = {
+    "4", "2", "0", "2", "stash " + src_hash + " 2,0,2", "free " + src_hash, "fail",
+  };
+
+  // Add a new data of 10 bytes to test the deadlock.
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", std::string(10, 0) },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
+
+  // Build the update package.
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
+
+  MemMapping map;
+  ASSERT_TRUE(map.MapFile(zip_file.path));
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
+
+  // Set up the handler, command_pipe, patch offset & length.
+  UpdaterInfo updater_info;
+  updater_info.package_zip = handle;
+  TemporaryFile temp_pipe;
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
+  updater_info.package_zip_addr = map.addr;
+  updater_info.package_zip_len = map.length;
+
+  TemporaryFile update_file;
+  ASSERT_TRUE(android::base::WriteStringToFile(src_content, update_file.path));
+  // Expect the stashed blocks to be freed.
+  std::string script = "block_image_update(\"" + std::string(update_file.path) +
+                       R"(", package_extract_file("transfer_list"), "new_data", "patch_data"))";
   expect("", script.c_str(), kNoCause, &updater_info);
   // Updater generates the stash name based on the input file name.
   std::string name_digest = get_sha1(update_file.path);
   std::string stash_base = "/cache/recovery/" + name_digest;
   ASSERT_EQ(0, access(stash_base.c_str(), F_OK));
-  ASSERT_EQ(-1, access((stash_base + tgt_hash).c_str(), F_OK));
+  ASSERT_EQ(-1, access((stash_base + src_hash).c_str(), F_OK));
   ASSERT_EQ(0, rmdir(stash_base.c_str()));
 
   ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
   CloseArchive(handle);
 }
 
-TEST_F(UpdaterTest, new_data_short_write) {
-  // Create a zip file with new_data.
+TEST_F(UpdaterTest, new_data_over_write) {
+  std::vector<std::string> transfer_list = {
+    "4", "1", "0", "0", "new 2,0,1",
+  };
+
+  // Write 4096 + 100 bytes of new data.
+  std::unordered_map<std::string, std::string> entries = {
+    { "new_data", std::string(4196, 0) },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
+
+  // Build the update package.
   TemporaryFile zip_file;
-  FILE* zip_file_ptr = fdopen(zip_file.release(), "wb");
-  ZipWriter zip_writer(zip_file_ptr);
+  BuildUpdatePackage(entries, zip_file.release());
+
+  MemMapping map;
+  ASSERT_TRUE(map.MapFile(zip_file.path));
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
 
-  // Add the empty new data.
-  ASSERT_EQ(0, zip_writer.StartEntry("empty_new_data", 0));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  // Add the short written new data.
-  ASSERT_EQ(0, zip_writer.StartEntry("short_new_data", 0));
-  std::string new_data_short = std::string(10, 'a');
-  ASSERT_EQ(0, zip_writer.WriteBytes(new_data_short.data(), new_data_short.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  // Add the data of exactly one block.
-  ASSERT_EQ(0, zip_writer.StartEntry("exact_new_data", 0));
-  std::string new_data_exact = std::string(4096, 'a');
-  ASSERT_EQ(0, zip_writer.WriteBytes(new_data_exact.data(), new_data_exact.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  // Add a dummy patch data.
-  ASSERT_EQ(0, zip_writer.StartEntry("patch_data", 0));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
+  // Set up the handler, command_pipe, patch offset & length.
+  UpdaterInfo updater_info;
+  updater_info.package_zip = handle;
+  TemporaryFile temp_pipe;
+  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
+  updater_info.package_zip_addr = map.addr;
+  updater_info.package_zip_len = map.length;
 
+  TemporaryFile update_file;
+  std::string script = "block_image_update(\"" + std::string(update_file.path) +
+                       R"(", package_extract_file("transfer_list"), "new_data", "patch_data"))";
+  expect("t", script.c_str(), kNoCause, &updater_info);
+
+  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
+  CloseArchive(handle);
+}
+
+TEST_F(UpdaterTest, new_data_short_write) {
   std::vector<std::string> transfer_list = {
     "4",
     "1",
@@ -547,12 +597,17 @@ TEST_F(UpdaterTest, new_data_short_write) {
     "0",
     "new 2,0,1",
   };
-  ASSERT_EQ(0, zip_writer.StartEntry("transfer_list", 0));
-  std::string commands = android::base::Join(transfer_list, '\n');
-  ASSERT_EQ(0, zip_writer.WriteBytes(commands.data(), commands.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  ASSERT_EQ(0, zip_writer.Finish());
-  ASSERT_EQ(0, fclose(zip_file_ptr));
+
+  std::unordered_map<std::string, std::string> entries = {
+    { "empty_new_data", "" },
+    { "short_new_data", std::string(10, 'a') },
+    { "exact_new_data", std::string(4096, 'a') },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
+
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
 
   MemMapping map;
   ASSERT_TRUE(map.MapFile(zip_file.path));
@@ -587,14 +642,6 @@ TEST_F(UpdaterTest, new_data_short_write) {
 }
 
 TEST_F(UpdaterTest, brotli_new_data) {
-  // Create a zip file with new_data.
-  TemporaryFile zip_file;
-  FILE* zip_file_ptr = fdopen(zip_file.release(), "wb");
-  ZipWriter zip_writer(zip_file_ptr);
-
-  // Add a brotli compressed new data entry.
-  ASSERT_EQ(0, zip_writer.StartEntry("new.dat.br", 0));
-
   auto generator = []() { return rand() % 128; };
   // Generate 100 blocks of random data.
   std::string brotli_new_data;
@@ -602,16 +649,12 @@ TEST_F(UpdaterTest, brotli_new_data) {
   generate_n(back_inserter(brotli_new_data), 4096 * 100, generator);
 
   size_t encoded_size = BrotliEncoderMaxCompressedSize(brotli_new_data.size());
-  std::vector<uint8_t> encoded_data(encoded_size);
+  std::string encoded_data(encoded_size, 0);
   ASSERT_TRUE(BrotliEncoderCompress(
       BROTLI_DEFAULT_QUALITY, BROTLI_DEFAULT_WINDOW, BROTLI_DEFAULT_MODE, brotli_new_data.size(),
-      reinterpret_cast<const uint8_t*>(brotli_new_data.data()), &encoded_size, encoded_data.data()));
-
-  ASSERT_EQ(0, zip_writer.WriteBytes(encoded_data.data(), encoded_size));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  // Add a dummy patch data.
-  ASSERT_EQ(0, zip_writer.StartEntry("patch_data", 0));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
+      reinterpret_cast<const uint8_t*>(brotli_new_data.data()), &encoded_size,
+      reinterpret_cast<uint8_t*>(const_cast<char*>(encoded_data.data()))));
+  encoded_data.resize(encoded_size);
 
   // Write a few small chunks of new data, then a large chunk, and finally a few small chunks.
   // This helps us to catch potential short writes.
@@ -627,12 +670,15 @@ TEST_F(UpdaterTest, brotli_new_data) {
     "new 2,98,99",
     "new 2,99,100",
   };
-  ASSERT_EQ(0, zip_writer.StartEntry("transfer_list", 0));
-  std::string commands = android::base::Join(transfer_list, '\n');
-  ASSERT_EQ(0, zip_writer.WriteBytes(commands.data(), commands.size()));
-  ASSERT_EQ(0, zip_writer.FinishEntry());
-  ASSERT_EQ(0, zip_writer.Finish());
-  ASSERT_EQ(0, fclose(zip_file_ptr));
+
+  std::unordered_map<std::string, std::string> entries = {
+    { "new.dat.br", std::move(encoded_data) },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, '\n') },
+  };
+
+  TemporaryFile zip_file;
+  BuildUpdatePackage(entries, zip_file.release());
 
   MemMapping map;
   ASSERT_TRUE(map.MapFile(zip_file.path));
diff --git a/tests/unit/rangeset_test.cpp b/tests/unit/rangeset_test.cpp
index 15bcec85..b3ed9921 100644
--- a/tests/unit/rangeset_test.cpp
+++ b/tests/unit/rangeset_test.cpp
@@ -21,7 +21,7 @@
 
 #include <gtest/gtest.h>
 
-#include "rangeset.h"
+#include "otautil/rangeset.h"
 
 TEST(RangeSetTest, Parse_smoke) {
   RangeSet rs = RangeSet::Parse("2,1,10");
@@ -156,4 +156,4 @@ TEST(SortedRangeSetTest, file_range) {
   ASSERT_EQ(static_cast<size_t>(40970), rs.GetOffsetInRangeSet(4096 * 16 + 10));
   // block#10 not in range.
   ASSERT_EXIT(rs.GetOffsetInRangeSet(40970), ::testing::KilledBySignal(SIGABRT), "");
-}
\ No newline at end of file
+}