2 files changed, 101 insertions, 99 deletions

diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index b253abccf..24c63e776 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -41,6 +41,7 @@
 #include <ziparchive/zip_archive.h>
 #include <ziparchive/zip_writer.h>
 
+#include "applypatch/applypatch.h"
 #include "common/test_constants.h"
 #include "edify/expr.h"
 #include "otautil/error_code.h"
@@ -214,55 +215,47 @@ TEST_F(UpdaterTest, getprop) {
     expect(nullptr, "getprop(\"arg1\", \"arg2\")", kArgsParsingFailure);
 }
 
-TEST_F(UpdaterTest, apply_patch_check) {
-  // Zero-argument is not valid.
-  expect(nullptr, "apply_patch_check()", kArgsParsingFailure);
-
-  // File not found.
-  expect("", "apply_patch_check(\"/doesntexist\")", kNoCause);
-
-  std::string src_file = from_testdata_base("boot.img");
-  std::string src_content;
-  ASSERT_TRUE(android::base::ReadFileToString(src_file, &src_content));
-  size_t src_size = src_content.size();
-  std::string src_hash = get_sha1(src_content);
-
-  // One-argument with EMMC:file:size:sha1 should pass the check.
-  std::string filename = android::base::Join(
-      std::vector<std::string>{ "EMMC", src_file, std::to_string(src_size), src_hash }, ":");
-  std::string cmd = "apply_patch_check(\"" + filename + "\")";
-  expect("t", cmd, kNoCause);
-
-  // EMMC:file:(size-1):sha1:(size+1):sha1 should fail the check.
-  std::string filename_bad = android::base::Join(
-      std::vector<std::string>{ "EMMC", src_file, std::to_string(src_size - 1), src_hash,
-                                std::to_string(src_size + 1), src_hash },
-      ":");
-  cmd = "apply_patch_check(\"" + filename_bad + "\")";
-  expect("", cmd, kNoCause);
-
-  // EMMC:file:(size-1):sha1:size:sha1:(size+1):sha1 should pass the check.
-  filename_bad =
-      android::base::Join(std::vector<std::string>{ "EMMC", src_file, std::to_string(src_size - 1),
-                                                    src_hash, std::to_string(src_size), src_hash,
-                                                    std::to_string(src_size + 1), src_hash },
-                          ":");
-  cmd = "apply_patch_check(\"" + filename_bad + "\")";
+TEST_F(UpdaterTest, patch_partition_check) {
+  // Zero argument is not valid.
+  expect(nullptr, "patch_partition_check()", kArgsParsingFailure);
+
+  std::string source_file = from_testdata_base("boot.img");
+  std::string source_content;
+  ASSERT_TRUE(android::base::ReadFileToString(source_file, &source_content));
+  size_t source_size = source_content.size();
+  std::string source_hash = get_sha1(source_content);
+  Partition source(source_file, source_size, source_hash);
+
+  std::string target_file = from_testdata_base("recovery.img");
+  std::string target_content;
+  ASSERT_TRUE(android::base::ReadFileToString(target_file, &target_content));
+  size_t target_size = target_content.size();
+  std::string target_hash = get_sha1(target_content);
+  Partition target(target_file, target_size, target_hash);
+
+  // One argument is not valid.
+  expect(nullptr, "patch_partition_check(\"" + source.ToString() + "\")", kArgsParsingFailure);
+  expect(nullptr, "patch_partition_check(\"" + target.ToString() + "\")", kArgsParsingFailure);
+
+  // Both of the source and target have the desired checksum.
+  std::string cmd =
+      "patch_partition_check(\"" + source.ToString() + "\", \"" + target.ToString() + "\")";
   expect("t", cmd, kNoCause);
 
-  // Multiple arguments.
-  // As long as it successfully loads the partition specified in filename, it won't check against
-  // any given SHAs.
-  cmd = "apply_patch_check(\"" + filename + "\", \"wrong_sha1\", \"wrong_sha2\")";
+  // Only source partition has the desired checksum.
+  Partition bad_target(target_file, target_size - 1, target_hash);
+  cmd = "patch_partition_check(\"" + source.ToString() + "\", \"" + bad_target.ToString() + "\")";
   expect("t", cmd, kNoCause);
 
-  cmd = "apply_patch_check(\"" + filename + "\", \"wrong_sha1\", \"" + src_hash +
-        "\", \"wrong_sha2\")";
+  // Only target partition has the desired checksum.
+  Partition bad_source(source_file, source_size + 1, source_hash);
+  cmd = "patch_partition_check(\"" + bad_source.ToString() + "\", \"" + target.ToString() + "\")";
   expect("t", cmd, kNoCause);
 
-  cmd = "apply_patch_check(\"" + filename_bad + "\", \"wrong_sha1\", \"" + src_hash +
-        "\", \"wrong_sha2\")";
-  expect("t", cmd, kNoCause);
+  // Neither of the source or target has the desired checksum.
+  cmd =
+      "patch_partition_check(\"" + bad_source.ToString() + "\", \"" + bad_target.ToString() + "\")";
+  expect("", cmd, kNoCause);
 }
 
 TEST_F(UpdaterTest, file_getprop) {
diff --git a/tests/unit/applypatch_test.cpp b/tests/unit/applypatch_test.cpp
index fe00e1e66..4204074e7 100644
--- a/tests/unit/applypatch_test.cpp
+++ b/tests/unit/applypatch_test.cpp
@@ -38,6 +38,7 @@
 
 #include "applypatch/applypatch.h"
 #include "common/test_constants.h"
+#include "edify/expr.h"
 #include "otautil/paths.h"
 #include "otautil/print_sha1.h"
 
@@ -58,12 +59,19 @@ class ApplyPatchTest : public ::testing::Test {
     target_size = recovery_fc.data.size();
     target_sha1 = print_sha1(recovery_fc.sha1);
 
+    source_partition = Partition(source_file, source_size, source_sha1);
+    target_partition = Partition(partition_file.path, target_size, target_sha1);
+
     srand(time(nullptr));
     bad_sha1_a = android::base::StringPrintf("%040x", rand());
     bad_sha1_b = android::base::StringPrintf("%040x", rand());
 
     // Reset the cache backup file.
-    Paths::Get().set_cache_temp_source("/cache/saved.file");
+    Paths::Get().set_cache_temp_source(cache_temp_source.path);
+  }
+
+  void TearDown() override {
+    ASSERT_TRUE(android::base::RemoveFileIfExists(cache_temp_source.path));
   }
 
   std::string source_file;
@@ -76,74 +84,75 @@ class ApplyPatchTest : public ::testing::Test {
 
   std::string bad_sha1_a;
   std::string bad_sha1_b;
+
+  Partition source_partition;
+  Partition target_partition;
+
+ private:
+  TemporaryFile partition_file;
+  TemporaryFile cache_temp_source;
 };
 
-TEST_F(ApplyPatchTest, CheckMode) {
-  std::string partition =
-      "EMMC:" + source_file + ":" + std::to_string(source_size) + ":" + source_sha1;
-  ASSERT_EQ(0, applypatch_check(partition, {}));
-  ASSERT_EQ(0, applypatch_check(partition, { source_sha1 }));
-  ASSERT_EQ(0, applypatch_check(partition, { bad_sha1_a, bad_sha1_b }));
-  ASSERT_EQ(0, applypatch_check(partition, { bad_sha1_a, source_sha1, bad_sha1_b }));
+TEST_F(ApplyPatchTest, CheckPartition) {
+  ASSERT_TRUE(CheckPartition(source_partition));
 }
 
-TEST_F(ApplyPatchTest, CheckMode_NonEmmcTarget) {
-  ASSERT_NE(0, applypatch_check(source_file, {}));
-  ASSERT_NE(0, applypatch_check(source_file, { source_sha1 }));
-  ASSERT_NE(0, applypatch_check(source_file, { bad_sha1_a, bad_sha1_b }));
-  ASSERT_NE(0, applypatch_check(source_file, { bad_sha1_a, source_sha1, bad_sha1_b }));
+TEST_F(ApplyPatchTest, CheckPartition_Mismatching) {
+  ASSERT_FALSE(CheckPartition(Partition(source_file, target_size, target_sha1)));
+  ASSERT_FALSE(CheckPartition(Partition(source_file, source_size, bad_sha1_a)));
+
+  ASSERT_FALSE(CheckPartition(Partition(source_file, source_size - 1, source_sha1)));
+  ASSERT_FALSE(CheckPartition(Partition(source_file, source_size + 1, source_sha1)));
 }
 
-TEST_F(ApplyPatchTest, CheckMode_EmmcTarget) {
-  // EMMC:source_file:size:sha1 should pass the check.
-  std::string src_file =
-      "EMMC:" + source_file + ":" + std::to_string(source_size) + ":" + source_sha1;
-  ASSERT_EQ(0, applypatch_check(src_file, {}));
-
-  // EMMC:source_file:(size-1):sha1:(size+1):sha1 should fail the check.
-  src_file = "EMMC:" + source_file + ":" + std::to_string(source_size - 1) + ":" + source_sha1 +
-             ":" + std::to_string(source_size + 1) + ":" + source_sha1;
-  ASSERT_NE(0, applypatch_check(src_file, {}));
-
-  // EMMC:source_file:(size-1):sha1:size:sha1:(size+1):sha1 should pass the check.
-  src_file = "EMMC:" + source_file + ":" + std::to_string(source_size - 1) + ":" + source_sha1 +
-             ":" + std::to_string(source_size) + ":" + source_sha1 + ":" +
-             std::to_string(source_size + 1) + ":" + source_sha1;
-  ASSERT_EQ(0, applypatch_check(src_file, {}));
-
-  // EMMC:source_file:(size+1):sha1:(size-1):sha1:size:sha1 should pass the check.
-  src_file = "EMMC:" + source_file + ":" + std::to_string(source_size + 1) + ":" + source_sha1 +
-             ":" + std::to_string(source_size - 1) + ":" + source_sha1 + ":" +
-             std::to_string(source_size) + ":" + source_sha1;
-  ASSERT_EQ(0, applypatch_check(src_file, {}));
-
-  // EMMC:target_file:(size+1):source_sha1:(size-1):source_sha1:size:source_sha1:size:target_sha1
-  // should pass the check.
-  src_file = "EMMC:" + target_file + ":" + std::to_string(source_size + 1) + ":" + source_sha1 +
-             ":" + std::to_string(source_size - 1) + ":" + source_sha1 + ":" +
-             std::to_string(source_size) + ":" + source_sha1 + ":" + std::to_string(target_size) +
-             ":" + target_sha1;
-  ASSERT_EQ(0, applypatch_check(src_file, {}));
+TEST_F(ApplyPatchTest, PatchPartitionCheck) {
+  ASSERT_TRUE(PatchPartitionCheck(target_partition, source_partition));
+
+  ASSERT_TRUE(
+      PatchPartitionCheck(Partition(source_file, source_size - 1, source_sha1), source_partition));
+
+  ASSERT_TRUE(
+      PatchPartitionCheck(Partition(source_file, source_size + 1, source_sha1), source_partition));
 }
 
-TEST_F(ApplyPatchTest, CheckMode_UseBackup) {
-  std::string corrupted =
-      "EMMC:" + source_file + ":" + std::to_string(source_size) + ":" + bad_sha1_a;
-  ASSERT_NE(0, applypatch_check(corrupted, { source_sha1 }));
+TEST_F(ApplyPatchTest, PatchPartitionCheck_UseBackup) {
+  ASSERT_FALSE(
+      PatchPartitionCheck(target_partition, Partition(target_file, source_size, source_sha1)));
 
   Paths::Get().set_cache_temp_source(source_file);
-  ASSERT_EQ(0, applypatch_check(corrupted, { source_sha1 }));
-  ASSERT_EQ(0, applypatch_check(corrupted, { bad_sha1_a, source_sha1, bad_sha1_b }));
+  ASSERT_TRUE(
+      PatchPartitionCheck(target_partition, Partition(target_file, source_size, source_sha1)));
 }
 
-TEST_F(ApplyPatchTest, CheckMode_UseBackup_BothCorrupted) {
-  std::string corrupted =
-      "EMMC:" + source_file + ":" + std::to_string(source_size) + ":" + bad_sha1_a;
-  ASSERT_NE(0, applypatch_check(corrupted, {}));
-  ASSERT_NE(0, applypatch_check(corrupted, { source_sha1 }));
+TEST_F(ApplyPatchTest, PatchPartitionCheck_UseBackup_BothCorrupted) {
+  ASSERT_FALSE(
+      PatchPartitionCheck(target_partition, Partition(target_file, source_size, source_sha1)));
 
-  Paths::Get().set_cache_temp_source(source_file);
-  ASSERT_NE(0, applypatch_check(corrupted, { bad_sha1_a, bad_sha1_b }));
+  Paths::Get().set_cache_temp_source(target_file);
+  ASSERT_FALSE(
+      PatchPartitionCheck(target_partition, Partition(target_file, source_size, source_sha1)));
+}
+
+TEST_F(ApplyPatchTest, PatchPartition) {
+  FileContents patch_fc;
+  ASSERT_EQ(0, LoadFileContents(from_testdata_base("recovery-from-boot.p"), &patch_fc));
+  Value patch(Value::Type::BLOB, std::string(patch_fc.data.cbegin(), patch_fc.data.cend()));
+
+  FileContents bonus_fc;
+  ASSERT_EQ(0, LoadFileContents(from_testdata_base("bonus.file"), &bonus_fc));
+  Value bonus(Value::Type::BLOB, std::string(bonus_fc.data.cbegin(), bonus_fc.data.cend()));
+
+  ASSERT_TRUE(PatchPartition(target_partition, source_partition, patch, &bonus));
+}
+
+// Tests patching an eMMC target without a separate bonus file (i.e. recovery-from-boot patch has
+// everything).
+TEST_F(ApplyPatchTest, PatchPartitionWithoutBonusFile) {
+  FileContents patch_fc;
+  ASSERT_EQ(0, LoadFileContents(from_testdata_base("recovery-from-boot-with-bonus.p"), &patch_fc));
+  Value patch(Value::Type::BLOB, std::string(patch_fc.data.cbegin(), patch_fc.data.cend()));
+
+  ASSERT_TRUE(PatchPartition(target_partition, source_partition, patch, nullptr));
 }
 
 class FreeCacheTest : public ::testing::Test {