1 files changed, 24 insertions, 11 deletions

diff --git a/verity/hash_tree_builder.cpp b/verity/hash_tree_builder.cpp
index 79d30405..1cd2e291 100644
--- a/verity/hash_tree_builder.cpp
+++ b/verity/hash_tree_builder.cpp
@@ -17,6 +17,7 @@
 #include "verity/hash_tree_builder.h"
 
 #include <algorithm>
+#include <functional>
 #include <memory>
 
 #include <android-base/file.h>
@@ -41,6 +42,9 @@ const EVP_MD* HashTreeBuilder::HashFunction(const std::string& hash_name) {
   if (android::base::EqualsIgnoreCase(hash_name, "sha512")) {
     return EVP_sha512();
   }
+  if (android::base::EqualsIgnoreCase(hash_name, "blake2b-256")) {
+    return EVP_blake2b256();
+  }
 
   LOG(ERROR) << "Unsupported hash algorithm " << hash_name;
   return nullptr;
@@ -92,18 +96,19 @@ bool HashTreeBuilder::ParseBytesArrayFromString(
   return true;
 }
 
-uint64_t HashTreeBuilder::CalculateSize(uint64_t input_size) const {
+uint64_t HashTreeBuilder::CalculateSize(
+    uint64_t input_size, size_t block_size, size_t hash_size) {
   uint64_t verity_blocks = 0;
   size_t level_blocks;
   size_t levels = 0;
   do {
     level_blocks =
-        verity_tree_blocks(input_size, block_size_, hash_size_, levels);
+        verity_tree_blocks(input_size, block_size, hash_size, levels);
     levels++;
     verity_blocks += level_blocks;
   } while (level_blocks > 1);
 
-  return verity_blocks * block_size_;
+  return verity_blocks * block_size;
 }
 
 bool HashTreeBuilder::Initialize(int64_t expected_data_size,
@@ -297,19 +302,14 @@ bool HashTreeBuilder::WriteHashTreeToFile(const std::string& output) const {
   return WriteHashTreeToFd(output_fd, 0);
 }
 
-bool HashTreeBuilder::WriteHashTreeToFd(int fd, uint64_t offset) const {
+bool HashTreeBuilder::WriteHashTree(
+    std::function<bool(const void*, size_t)> callback) const {
   CHECK(!verity_tree_.empty());
 
-  if (lseek(fd, offset, SEEK_SET) != offset) {
-    PLOG(ERROR) << "Failed to seek the output fd, offset: " << offset;
-    return false;
-  }
-
   // 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 (!android::base::WriteFully(fd, level_blocks.data(),
-                                   level_blocks.size())) {
+    if (!callback(level_blocks.data(), level_blocks.size())) {
       PLOG(ERROR) << "Failed to write the hash tree level " << i;
       return false;
     }
@@ -318,6 +318,19 @@ bool HashTreeBuilder::WriteHashTreeToFd(int fd, uint64_t offset) const {
   return true;
 }
 
+bool HashTreeBuilder::WriteHashTreeToFd(int fd, uint64_t offset) const {
+  CHECK(!verity_tree_.empty());
+
+  if (lseek(fd, offset, SEEK_SET) != offset) {
+    PLOG(ERROR) << "Failed to seek the output fd, offset: " << offset;
+    return false;
+  }
+
+  return WriteHashTree([fd](auto data, auto size) {
+    return android::base::WriteFully(fd, data, size);
+  });
+}
+
 void HashTreeBuilder::AppendPaddings(std::vector<unsigned char>* data) {
   size_t remainder = data->size() % block_size_;
   if (remainder != 0) {