3 files changed, 83 insertions, 0 deletions

diff --git a/tests/Android.mk b/tests/Android.mk
index 58ef3a2af..daf4853b9 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -137,6 +137,7 @@ libupdater_static_libraries := \
     libext4_utils \
     libfec \
     libfec_rs \
+    libverity_tree \
     libfs_mgr \
     libgtest_prod \
     liblog \
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index 9fcf17f13..248b469b0 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -37,6 +37,7 @@
 #include <brotli/encode.h>
 #include <bsdiff/bsdiff.h>
 #include <gtest/gtest.h>
+#include <verity/hash_tree_builder.h>
 #include <ziparchive/zip_archive.h>
 #include <ziparchive/zip_writer.h>
 
@@ -389,6 +390,86 @@ TEST_F(UpdaterTest, read_file) {
   expect("", script, kNoCause);
 }
 
+TEST_F(UpdaterTest, compute_hash_tree_smoke) {
+  std::string data;
+  for (unsigned char i = 0; i < 128; i++) {
+    data += std::string(4096, i);
+  }
+  // Appends an additional block for verity data.
+  data += std::string(4096, 0);
+  ASSERT_EQ(129 * 4096, data.size());
+  ASSERT_TRUE(android::base::WriteStringToFile(data, image_file_));
+
+  std::string salt = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7";
+  std::string expected_root_hash =
+      "7e0a8d8747f54384014ab996f5b2dc4eb7ff00c630eede7134c9e3f05c0dd8ca";
+  // hash_tree_ranges, source_ranges, hash_algorithm, salt_hex, root_hash
+  std::vector<std::string> tokens{ "compute_hash_tree", "2,128,129", "2,0,128", "sha256", salt,
+                                   expected_root_hash };
+  std::string hash_tree_command = android::base::Join(tokens, " ");
+
+  std::vector<std::string> transfer_list{
+    "4", "2", "0", "2", hash_tree_command,
+  };
+
+  PackageEntries entries{
+    { "new_data", "" },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, "\n") },
+  };
+
+  RunBlockImageUpdate(false, entries, image_file_, "t");
+
+  std::string updated;
+  ASSERT_TRUE(android::base::ReadFileToString(image_file_, &updated));
+  ASSERT_EQ(129 * 4096, updated.size());
+  ASSERT_EQ(data.substr(0, 128 * 4096), updated.substr(0, 128 * 4096));
+
+  // Computes the SHA256 of the salt + hash_tree_data and expects the result to match with the
+  // root_hash.
+  std::vector<unsigned char> salt_bytes;
+  ASSERT_TRUE(HashTreeBuilder::ParseBytesArrayFromString(salt, &salt_bytes));
+  std::vector<unsigned char> hash_tree = std::move(salt_bytes);
+  hash_tree.insert(hash_tree.end(), updated.begin() + 128 * 4096, updated.end());
+
+  std::vector<unsigned char> digest(SHA256_DIGEST_LENGTH);
+  SHA256(hash_tree.data(), hash_tree.size(), digest.data());
+  ASSERT_EQ(expected_root_hash, HashTreeBuilder::BytesArrayToString(digest));
+}
+
+TEST_F(UpdaterTest, compute_hash_tree_root_mismatch) {
+  std::string data;
+  for (size_t i = 0; i < 128; i++) {
+    data += std::string(4096, i);
+  }
+  // Appends an additional block for verity data.
+  data += std::string(4096, 0);
+  ASSERT_EQ(129 * 4096, data.size());
+  // Corrupts one bit
+  data[4096] = 'A';
+  ASSERT_TRUE(android::base::WriteStringToFile(data, image_file_));
+
+  std::string salt = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7";
+  std::string expected_root_hash =
+      "7e0a8d8747f54384014ab996f5b2dc4eb7ff00c630eede7134c9e3f05c0dd8ca";
+  // hash_tree_ranges, source_ranges, hash_algorithm, salt_hex, root_hash
+  std::vector<std::string> tokens{ "compute_hash_tree", "2,128,129", "2,0,128", "sha256", salt,
+                                   expected_root_hash };
+  std::string hash_tree_command = android::base::Join(tokens, " ");
+
+  std::vector<std::string> transfer_list{
+    "4", "2", "0", "2", hash_tree_command,
+  };
+
+  PackageEntries entries{
+    { "new_data", "" },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, "\n") },
+  };
+
+  RunBlockImageUpdate(false, entries, image_file_, "", kHashTreeComputationFailure);
+}
+
 TEST_F(UpdaterTest, write_value) {
   // write_value() expects two arguments.
   expect(nullptr, "write_value()", kArgsParsingFailure);
diff --git a/tests/unit/commands_test.cpp b/tests/unit/commands_test.cpp
index 3daa58f33..9679a9e73 100644
--- a/tests/unit/commands_test.cpp
+++ b/tests/unit/commands_test.cpp
@@ -30,6 +30,7 @@ TEST(CommandsTest, ParseType) {
   ASSERT_EQ(Command::Type::IMGDIFF, Command::ParseType("imgdiff"));
   ASSERT_EQ(Command::Type::STASH, Command::ParseType("stash"));
   ASSERT_EQ(Command::Type::FREE, Command::ParseType("free"));
+  ASSERT_EQ(Command::Type::COMPUTE_HASH_TREE, Command::ParseType("compute_hash_tree"));
 }
 
 TEST(CommandsTest, ParseType_InvalidCommand) {