Merge "verity: Add HashTreeBuilder::CheckHashTree()."

2 files changed, 30 insertions, 0 deletions

diff --git a/verity/hash_tree_builder.cpp b/verity/hash_tree_builder.cpp
index c5d64f87..a776572f 100644
--- a/verity/hash_tree_builder.cpp
+++ b/verity/hash_tree_builder.cpp
@@ -219,6 +219,34 @@ bool HashTreeBuilder::BuildHashTree() {
   return true;
 }
 
+bool HashTreeBuilder::CheckHashTree(
+    const std::vector<unsigned char>& hash_tree) const {
+  size_t offset = 0;
+  // Reads reversely to output the verity tree top-down.
+  for (size_t i = verity_tree_.size(); i > 0; i--) {
+    const auto& level_blocks = verity_tree_[i - 1];
+    if (offset + level_blocks.size() > hash_tree.size()) {
+      LOG(ERROR) << "Hash tree too small: " << hash_tree.size();
+      return false;
+    }
+    auto iter = std::mismatch(level_blocks.begin(), level_blocks.end(),
+                              hash_tree.begin() + offset)
+                    .first;
+    if (iter != level_blocks.end()) {
+      LOG(ERROR) << "Mismatch found at the hash tree level " << i << " offset "
+                 << std::distance(level_blocks.begin(), iter);
+      return false;
+    }
+    offset += level_blocks.size();
+  }
+  if (offset != hash_tree.size()) {
+    LOG(ERROR) << "Hash tree size mismatch: " << hash_tree.size()
+               << " != " << offset;
+    return false;
+  }
+  return true;
+}
+
 bool HashTreeBuilder::WriteHashTreeToFile(const std::string& output) const {
   android::base::unique_fd output_fd(
       open(output.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666));
diff --git a/verity/include/verity/hash_tree_builder.h b/verity/include/verity/hash_tree_builder.h
index 333e5b9d..c6f00c0e 100644
--- a/verity/include/verity/hash_tree_builder.h
+++ b/verity/include/verity/hash_tree_builder.h
@@ -46,6 +46,8 @@ class HashTreeBuilder {
   bool Update(const unsigned char* data, size_t len);
   // Computes the upper levels of the hash tree based on the 0th level.
   bool BuildHashTree();
+  // Check the built hash tree against |hash_tree|, return true if they match.
+  bool CheckHashTree(const std::vector<unsigned char>& hash_tree) const;
   // Writes the computed hash tree top-down to |output|.
   bool WriteHashTreeToFile(const std::string& output) const;
   bool WriteHashTreeToFd(int fd, uint64_t offset) const;