Add WriteHashTree() to allow users process hashtree data themselves

Sometimes we may want to write hash tree data somewhere else other than
fd. For example, we may want to send hash tree data to CowWriter.
Currently there's no interface to do that, so we added a generate
WriteHashTree() so users can process the hashtree data themselves

Test: treehugger
Change-Id: I56258ee8497336abcb804facef017ed71b5c9415

2 files changed, 19 insertions, 8 deletions

diff --git a/verity/hash_tree_builder.cpp b/verity/hash_tree_builder.cpp
index 79d30405..df6f7864 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>
@@ -297,19 +298,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 +314,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) {
diff --git a/verity/include/verity/hash_tree_builder.h b/verity/include/verity/hash_tree_builder.h
index 7e870bf4..f4ceab8c 100644
--- a/verity/include/verity/hash_tree_builder.h
+++ b/verity/include/verity/hash_tree_builder.h
@@ -20,6 +20,7 @@
 #include <inttypes.h>
 #include <stddef.h>
 
+#include <functional>
 #include <string>
 #include <vector>
 
@@ -51,6 +52,7 @@ class HashTreeBuilder {
   // Writes the computed hash tree top-down to |output|.
   bool WriteHashTreeToFile(const std::string& output) const;
   bool WriteHashTreeToFd(int fd, uint64_t offset) const;
+  bool WriteHashTree(std::function<bool(const void*, size_t)> callback) const;
 
   size_t hash_size() const { return hash_size_; }
   const std::vector<unsigned char>& root_hash() const { return root_hash_; }