payload_generator: Verify hash tree.
am: 5cc2a3f01a

Change-Id: I334fa808517751f30eea1dec39f9df0d007aae86
diff --git a/Android.mk b/Android.mk
index 4d306c4..eb31771 100644
--- a/Android.mk
+++ b/Android.mk
@@ -634,10 +634,12 @@
     liblzma \
     libpayload_consumer \
     libpuffdiff \
+    libverity_tree \
     update_metadata-protos \
     $(ue_libpayload_consumer_exported_static_libraries) \
     $(ue_update_metadata_protos_exported_static_libraries)
 ue_libpayload_generator_exported_shared_libraries := \
+    libbase \
     libext2fs \
     $(ue_libpayload_consumer_exported_shared_libraries) \
     $(ue_update_metadata_protos_exported_shared_libraries)
diff --git a/payload_generator/payload_generation_config_android.cc b/payload_generator/payload_generation_config_android.cc
index 5ffd11e..7d9988f 100644
--- a/payload_generator/payload_generation_config_android.cc
+++ b/payload_generator/payload_generation_config_android.cc
@@ -19,6 +19,7 @@
 #include <base/logging.h>
 #include <brillo/secure_blob.h>
 #include <libavb/libavb.h>
+#include <verity/hash_tree_builder.h>
 
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_generator/extent_ranges.h"
@@ -55,9 +56,38 @@
 
   TEST_AND_RETURN_FALSE(hashtree.hash_block_size ==
                         part->fs_interface->GetBlockSize());
+
+  // Generate hash tree based on the descriptor and verify that it matches
+  // the hash tree stored in the image.
+  auto hash_function =
+      HashTreeBuilder::HashFunction(part->verity.hash_tree_algorithm);
+  TEST_AND_RETURN_FALSE(hash_function != nullptr);
+  HashTreeBuilder hash_tree_builder(hashtree.data_block_size, hash_function);
+  TEST_AND_RETURN_FALSE(hash_tree_builder.Initialize(
+      hashtree.image_size, part->verity.hash_tree_salt));
+  TEST_AND_RETURN_FALSE(hash_tree_builder.CalculateSize(hashtree.image_size) ==
+                        hashtree.tree_size);
+
+  brillo::Blob buffer;
+  for (uint64_t offset = 0; offset < hashtree.image_size;) {
+    constexpr uint64_t kBufferSize = 1024 * 1024;
+    size_t bytes_to_read = std::min(kBufferSize, hashtree.image_size - offset);
+    TEST_AND_RETURN_FALSE(
+        utils::ReadFileChunk(part->path, offset, bytes_to_read, &buffer));
+    TEST_AND_RETURN_FALSE(
+        hash_tree_builder.Update(buffer.data(), buffer.size()));
+    offset += buffer.size();
+    buffer.clear();
+  }
+  TEST_AND_RETURN_FALSE(hash_tree_builder.BuildHashTree());
+  TEST_AND_RETURN_FALSE(utils::ReadFileChunk(
+      part->path, hashtree.tree_offset, hashtree.tree_size, &buffer));
+  TEST_AND_RETURN_FALSE(hash_tree_builder.CheckHashTree(buffer));
+
   part->verity.hash_tree_extent = ExtentForBytes(
       hashtree.hash_block_size, hashtree.tree_offset, hashtree.tree_size);
 
+  // TODO(senj): Verify FEC data.
   part->verity.fec_data_extent =
       ExtentForBytes(hashtree.data_block_size, 0, hashtree.fec_offset);
   part->verity.fec_extent = ExtentForBytes(