Add support of more hash algorithms for verity tree builder

Currently, verified boot 1.0 is using SHA256 to compute the hash tree
while AVB is using SHA1. We should support at least these two hash
functionss in the HashTreeBuilder and command line parser. And we can
potentially add more algorithms in the future.

Bug: 25170618
Test: unit tests pass
Change-Id: I5f6dc8a545c0cef75acbbc2044959e8100f9f842

4 files changed, 72 insertions, 9 deletions

diff --git a/verity/build_verity_tree_main.cpp b/verity/build_verity_tree_main.cpp
index 9eaaa79e..76164cc5 100644
--- a/verity/build_verity_tree_main.cpp
+++ b/verity/build_verity_tree_main.cpp
@@ -52,6 +52,7 @@ int main(int argc, char** argv) {
   bool sparse = false;
   uint64_t calculate_size = 0;
   bool verbose = false;
+  std::string hash_algorithm;
 
   while (1) {
     constexpr struct option long_options[] = {
@@ -61,8 +62,10 @@ int main(int argc, char** argv) {
         {"sparse", no_argument, nullptr, 'S'},
         {"verity-size", required_argument, nullptr, 's'},
         {"verbose", no_argument, nullptr, 'v'},
+        {"hash-algorithm", required_argument, nullptr, 0},
         {nullptr, 0, nullptr, 0}};
-    int c = getopt_long(argc, argv, "a:A:hSs:v", long_options, nullptr);
+    int option_index;
+    int c = getopt_long(argc, argv, "a:A:hSs:v", long_options, &option_index);
     if (c < 0) {
       break;
     }
@@ -102,6 +105,12 @@ int main(int argc, char** argv) {
       case 'v':
         verbose = true;
         break;
+      case 0: {
+        std::string option = long_options[option_index].name;
+        if (option == "hash-algorithm") {
+          hash_algorithm = optarg;
+        }
+      } break;
       case '?':
         usage();
         return 1;
@@ -113,7 +122,13 @@ int main(int argc, char** argv) {
   argc -= optind;
   argv += optind;
 
-  HashTreeBuilder builder(kBlockSize);
+  auto hash_function = hash_algorithm.empty()
+                           ? EVP_sha256()
+                           : HashTreeBuilder::HashFunction(hash_algorithm);
+  if (hash_function == nullptr) {
+    return 1;
+  }
+  HashTreeBuilder builder(kBlockSize, hash_function);
 
   if (calculate_size) {
     if (argc != 0) {
diff --git a/verity/build_verity_tree_test.cpp b/verity/build_verity_tree_test.cpp
index 36aa34cb..50f6e397 100644
--- a/verity/build_verity_tree_test.cpp
+++ b/verity/build_verity_tree_test.cpp
@@ -22,6 +22,7 @@
 #include <vector>
 
 #include <gtest/gtest.h>
+#include <openssl/evp.h>
 
 #include "build_verity_tree_utils.h"
 #include "hash_tree_builder.h"
@@ -38,7 +39,7 @@ class BuildVerityTreeTest : public ::testing::Test {
   void SetUp() override {
     salt_hex =
         std::vector<unsigned char>(kSaltHex, kSaltHex + sizeof(kSaltHex));
-    builder.reset(new HashTreeBuilder(4096));
+    builder.reset(new HashTreeBuilder(4096, EVP_sha256()));
   }
 
   const std::vector<unsigned char>& zero_block_hash() const {
@@ -133,3 +134,15 @@ TEST_F(BuildVerityTreeTest, StreamingDataMultipleBlocks) {
   ASSERT_EQ("6e73d59b0b6baf026e921814979a7db02244c95a46b869a17aa1310dad066deb",
             HashTreeBuilder::BytesArrayToString(builder->root_hash()));
 }
+
+TEST_F(BuildVerityTreeTest, SHA1MultipleBlocks) {
+  std::vector<unsigned char> data(128 * 4096, 0xff);
+
+  builder.reset(
+      new HashTreeBuilder(4096, HashTreeBuilder::HashFunction("SHA1")));
+
+  GenerateHashTree(data, salt_hex);
+  ASSERT_EQ(1u, verity_tree().size());
+  ASSERT_EQ("7ea287e6167929988810077abaafbc313b2b8593000000000000000000000000",
+            HashTreeBuilder::BytesArrayToString(builder->root_hash()));
+}
diff --git a/verity/hash_tree_builder.cpp b/verity/hash_tree_builder.cpp
index 05c06ada..41649900 100644
--- a/verity/hash_tree_builder.cpp
+++ b/verity/hash_tree_builder.cpp
@@ -16,18 +16,45 @@
 
 #include "hash_tree_builder.h"
 
+#include <algorithm>
+
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/stringprintf.h>
+#include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 
 #include "build_verity_tree_utils.h"
 
-HashTreeBuilder::HashTreeBuilder(size_t block_size)
-    : block_size_(block_size), data_size_(0), md_(EVP_sha256()) {
+const EVP_MD* HashTreeBuilder::HashFunction(const std::string& hash_name) {
+  if (android::base::EqualsIgnoreCase(hash_name, "sha1")) {
+    return EVP_sha1();
+  }
+  if (android::base::EqualsIgnoreCase(hash_name, "sha256")) {
+    return EVP_sha256();
+  }
+  if (android::base::EqualsIgnoreCase(hash_name, "sha384")) {
+    return EVP_sha384();
+  }
+  if (android::base::EqualsIgnoreCase(hash_name, "sha512")) {
+    return EVP_sha512();
+  }
+
+  LOG(ERROR) << "Unsupported hash algorithm " << hash_name;
+  return nullptr;
+}
+
+HashTreeBuilder::HashTreeBuilder(size_t block_size, const EVP_MD* md)
+    : block_size_(block_size), data_size_(0), md_(md) {
   CHECK(md_ != nullptr) << "Failed to initialize md";
 
-  hash_size_ = EVP_MD_size(md_);
+  hash_size_raw_ = EVP_MD_size(md_);
+
+  // Round up the hash size to the next power of 2.
+  hash_size_ = 1;
+  while (hash_size_ < hash_size_raw_) {
+    hash_size_ = hash_size_ << 1;
+  }
   CHECK_LT(hash_size_ * 2, block_size_);
 }
 
@@ -94,7 +121,9 @@ bool HashTreeBuilder::HashBlock(const unsigned char* block,
   EVP_MD_CTX_destroy(mdctx);
 
   CHECK_EQ(1, ret);
-  CHECK_EQ(hash_size_, s);
+  CHECK_EQ(hash_size_raw_, s);
+  std::fill(out + s, out + hash_size_, 0);
+
   return true;
 }
 
diff --git a/verity/hash_tree_builder.h b/verity/hash_tree_builder.h
index 8f180f4b..2a1ee56c 100644
--- a/verity/hash_tree_builder.h
+++ b/verity/hash_tree_builder.h
@@ -30,10 +30,9 @@
 // the total data size should be know in advance. Once all the data is ready,
 // appropriate functions can be called to build the upper levels of the hash
 // tree and output the tree to a file.
-// TODO(xunchang) add support of various hash algorithms.
 class HashTreeBuilder {
  public:
-  explicit HashTreeBuilder(size_t block_size);
+  HashTreeBuilder(size_t block_size, const EVP_MD* md);
   // Returns the size of the verity tree in bytes given the input data size.
   uint64_t CalculateSize(uint64_t input_size) const;
   // Gets ready for the hash tree computation. We expect |expected_data_size|
@@ -57,6 +56,10 @@ class HashTreeBuilder {
   static std::string BytesArrayToString(
       const std::vector<unsigned char>& bytes);
 
+  // Returns the hash function given the name of the hash algorithm. Returns
+  // nullptr if the algorithm is unrecongnized or not supported.
+  static const EVP_MD* HashFunction(const std::string& hash_name);
+
  private:
   friend class BuildVerityTreeTest;
   // Calculates the hash of one single block. Write the result to |out|, a
@@ -75,6 +78,9 @@ class HashTreeBuilder {
   uint64_t data_size_;
   std::vector<unsigned char> salt_;
   const EVP_MD* md_;
+  // The raw hash size of the hash algorithm specified by md_.
+  size_t hash_size_raw_;
+  // Hash size rounded up to the next power of 2. (e.g. 20 -> 32)
   size_t hash_size_;
 
   // Pre-calculated hash of a zero block.