Convert the sparse_hash_ctx to the HashTreeBuilder class

Also convert the verity tree storage to std::vector. Check that the run
time doesn't increase much.

Bug: 25170618
Test: run build_verity_tree on a system image and check the root hash
Change-Id: I911c40fe9c540a35bd0397c16e09be81f28b3642

5 files changed, 475 insertions, 293 deletions

diff --git a/verity/Android.bp b/verity/Android.bp
index 2c7cda73..3bc6e754 100644
--- a/verity/Android.bp
+++ b/verity/Android.bp
@@ -68,7 +68,10 @@ cc_binary_host {
 
 cc_binary_host {
     name: "build_verity_tree",
-    srcs: ["build_verity_tree.cpp"],
+    srcs: [
+        "hash_tree_builder.cpp",
+        "build_verity_tree.cpp"
+    ],
 
     static_libs: [
         "libsparse",
diff --git a/verity/build_verity_tree.cpp b/verity/build_verity_tree.cpp
index ad812ed8..e9978cd8 100644
--- a/verity/build_verity_tree.cpp
+++ b/verity/build_verity_tree.cpp
@@ -1,112 +1,58 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 #include <openssl/bn.h>
-#include <openssl/evp.h>
 #include <sparse/sparse.h>
 
 #undef NDEBUG
 
-#include <assert.h>
 #include <errno.h>
 #include <getopt.h>
 #include <fcntl.h>
 #include <inttypes.h>
-#include <limits.h>
 #include <stdio.h>
-#include <stdlib.h>
 #include <string.h>
 #include <unistd.h>
+
+#include <limits>
 #include <vector>
 
 #include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/unique_fd.h>
 
-struct sparse_hash_ctx {
-    unsigned char *hashes;
-    const unsigned char *salt;
-    uint64_t salt_size;
-    uint64_t hash_size;
-    uint64_t block_size;
-    const unsigned char *zero_block_hash;
-    const EVP_MD *md;
-};
-
-#define div_round_up(x,y) (((x) + (y) - 1)/(y))
-
-#define round_up(x,y) (div_round_up(x,y)*(y))
+#include "build_verity_tree_utils.h"
+#include "hash_tree_builder.h"
 
 #define FATAL(x...) { \
     fprintf(stderr, x); \
     exit(1); \
 }
 
-size_t verity_tree_blocks(uint64_t data_size, size_t block_size, size_t hash_size,
-                          int level)
-{
-    size_t level_blocks = div_round_up(data_size, block_size);
-    int hashes_per_block = div_round_up(block_size, hash_size);
-
-    do {
-        level_blocks = div_round_up(level_blocks, hashes_per_block);
-    } while (level--);
-
-    return level_blocks;
-}
-
-int hash_block(const EVP_MD *md,
-               const unsigned char *block, size_t len,
-               const unsigned char *salt, size_t salt_len,
-               unsigned char *out, size_t *out_size)
-{
-    EVP_MD_CTX *mdctx;
-    unsigned int s;
-    int ret = 1;
-
-    mdctx = EVP_MD_CTX_create();
-    assert(mdctx);
-    ret &= EVP_DigestInit_ex(mdctx, md, NULL);
-    ret &= EVP_DigestUpdate(mdctx, salt, salt_len);
-    ret &= EVP_DigestUpdate(mdctx, block, len);
-    ret &= EVP_DigestFinal_ex(mdctx, out, &s);
-    EVP_MD_CTX_destroy(mdctx);
-    assert(ret == 1);
-    if (out_size) {
-        *out_size = s;
-    }
-    return 0;
-}
-
-int hash_blocks(const EVP_MD *md,
-                const unsigned char *in, size_t in_size,
-                unsigned char *out, size_t *out_size,
-                const unsigned char *salt, size_t salt_size,
-                size_t block_size)
-{
-    size_t s;
-    *out_size = 0;
-    for (size_t i = 0; i < in_size; i += block_size) {
-        hash_block(md, in + i, block_size, salt, salt_size, out, &s);
-        out += s;
-        *out_size += s;
-    }
+size_t verity_tree_blocks(uint64_t data_size, size_t block_size,
+                          size_t hash_size, size_t level) {
+  uint64_t level_blocks = div_round_up(data_size, block_size);
+  uint64_t hashes_per_block = div_round_up(block_size, hash_size);
 
-    return 0;
-}
+  do {
+    level_blocks = div_round_up(level_blocks, hashes_per_block);
+  } while (level--);
 
-int hash_chunk(void *priv, const void *data, size_t len)
-{
-    struct sparse_hash_ctx *ctx = (struct sparse_hash_ctx *)priv;
-    assert(len % ctx->block_size == 0);
-    if (data) {
-        size_t s;
-        hash_blocks(ctx->md, (const unsigned char *)data, len,
-                    ctx->hashes, &s,
-                    ctx->salt, ctx->salt_size, ctx->block_size);
-        ctx->hashes += s;
-    } else {
-        for (size_t i = 0; i < len; i += ctx->block_size) {
-            memcpy(ctx->hashes, ctx->zero_block_hash, ctx->hash_size);
-            ctx->hashes += ctx->hash_size;
-        }
-    }
-    return 0;
+  CHECK_LE(level_blocks, std::numeric_limits<size_t>::max());
+  return level_blocks;
 }
 
 void usage(void)
@@ -124,231 +70,178 @@ void usage(void)
 
 int main(int argc, char **argv)
 {
-    char *data_filename;
-    char *verity_filename;
-    std::vector<unsigned char> salt;
-    bool sparse = false;
-    size_t block_size = 4096;
-    uint64_t calculate_size = 0;
-    bool verbose = false;
-
-    while (1) {
-        const static struct option long_options[] = {
-            {"salt-str", required_argument, 0, 'a'},
-            {"salt-hex", required_argument, 0, 'A'},
-            {"help", no_argument, 0, 'h'},
-            {"sparse", no_argument, 0, 'S'},
-            {"verity-size", required_argument, 0, 's'},
-            {"verbose", no_argument, 0, 'v'},
-            {NULL, 0, 0, 0}
-        };
-        int c = getopt_long(argc, argv, "a:A:hSs:v", long_options, NULL);
-        if (c < 0) {
-            break;
+  constexpr size_t kBlockSize = 4096;
+
+  char* data_filename;
+  char* verity_filename;
+  std::vector<unsigned char> salt;
+  bool sparse = false;
+  uint64_t calculate_size = 0;
+  bool verbose = false;
+
+  while (1) {
+    const static struct option long_options[] = {
+        {"salt-str", required_argument, 0, 'a'},
+        {"salt-hex", required_argument, 0, 'A'},
+        {"help", no_argument, 0, 'h'},
+        {"sparse", no_argument, 0, 'S'},
+        {"verity-size", required_argument, 0, 's'},
+        {"verbose", no_argument, 0, 'v'},
+        {NULL, 0, 0, 0}};
+    int c = getopt_long(argc, argv, "a:A:hSs:v", long_options, NULL);
+    if (c < 0) {
+      break;
+    }
+
+    switch (c) {
+      case 'a':
+        salt.clear();
+        salt.insert(salt.end(), optarg, &optarg[strlen(optarg)]);
+        break;
+      case 'A': {
+        BIGNUM* bn = NULL;
+        if (!BN_hex2bn(&bn, optarg)) {
+          FATAL("failed to convert salt from hex\n");
         }
-
-        switch (c) {
-        case 'a':
-            salt.clear();
-            salt.insert(salt.end(), optarg, &optarg[strlen(optarg)]);
-            break;
-        case 'A': {
-                BIGNUM *bn = NULL;
-                if(!BN_hex2bn(&bn, optarg)) {
-                    FATAL("failed to convert salt from hex\n");
-                }
-                size_t salt_size = BN_num_bytes(bn);
-                salt.resize(salt_size);
-                if (BN_bn2bin(bn, salt.data()) != salt_size) {
-                    FATAL("failed to convert salt to bytes\n");
-                }
-            }
-            break;
-        case 'h':
-            usage();
-            return 1;
-        case 'S':
-            sparse = true;
-            break;
-        case 's': {
-                char* endptr;
-                errno = 0;
-                unsigned long long int inSize = strtoull(optarg, &endptr, 0);
-                if (optarg[0] == '\0' || *endptr != '\0' ||
-                        (errno == ERANGE && inSize == ULLONG_MAX)) {
-                    FATAL("invalid value of verity-size\n");
-                }
-                if (inSize > UINT64_MAX) {
-                    FATAL("invalid value of verity-size\n");
-                }
-                calculate_size = (uint64_t)inSize;
-            }
-            break;
-        case 'v':
-            verbose = true;
-            break;
-        case '?':
-            usage();
-            return 1;
-        default:
-            abort();
+        size_t salt_size = BN_num_bytes(bn);
+        salt.resize(salt_size);
+        if (BN_bn2bin(bn, salt.data()) != salt_size) {
+          FATAL("failed to convert salt to bytes\n");
         }
-    }
-
-    argc -= optind;
-    argv += optind;
-
-    const EVP_MD *md = EVP_sha256();
-    if (!md) {
-        FATAL("failed to get digest\n");
-    }
-
-    size_t hash_size = EVP_MD_size(md);
-    assert(hash_size * 2 < block_size);
-
-    if (salt.empty()) {
-        salt.resize(hash_size);
-
-        int random_fd = open("/dev/urandom", O_RDONLY);
-        if (random_fd < 0) {
-            FATAL("failed to open /dev/urandom\n");
-        }
-
-        ssize_t ret = read(random_fd, salt.data(), salt.size());
-        if (ret != static_cast<ssize_t>(salt.size())) {
-            FATAL("failed to read %zu bytes from /dev/urandom: %zd %d\n", salt.size(), ret, errno);
+      } break;
+      case 'h':
+        usage();
+        return 1;
+      case 'S':
+        sparse = true;
+        break;
+      case 's': {
+        char* endptr;
+        errno = 0;
+        unsigned long long int inSize = strtoull(optarg, &endptr, 0);
+        if (optarg[0] == '\0' || *endptr != '\0' ||
+            (errno == ERANGE && inSize == ULLONG_MAX)) {
+          FATAL("invalid value of verity-size\n");
         }
-        close(random_fd);
-    }
-
-    if (calculate_size) {
-        if (argc != 0) {
-            usage();
-            return 1;
+        if (inSize > UINT64_MAX) {
+          FATAL("invalid value of verity-size\n");
         }
-        size_t verity_blocks = 0;
-        size_t level_blocks;
-        int levels = 0;
-        do {
-            level_blocks = verity_tree_blocks(calculate_size, block_size, hash_size, levels);
-            levels++;
-            verity_blocks += level_blocks;
-        } while (level_blocks > 1);
-
-        printf("%" PRIu64 "\n", (uint64_t)verity_blocks * block_size);
-        return 0;
-    }
-
-    if (argc != 2) {
+        calculate_size = (uint64_t)inSize;
+      } break;
+      case 'v':
+        verbose = true;
+        break;
+      case '?':
         usage();
         return 1;
+      default:
+        abort();
     }
+  }
 
-    data_filename = argv[0];
-    verity_filename = argv[1];
+  argc -= optind;
+  argv += optind;
 
-    int fd = open(data_filename, O_RDONLY);
-    if (fd < 0) {
-        FATAL("failed to open %s\n", data_filename);
-    }
+  HashTreeBuilder builder(kBlockSize);
 
-    struct sparse_file *file;
-    if (sparse) {
-        file = sparse_file_import(fd, false, false);
-    } else {
-        file = sparse_file_import_auto(fd, false, verbose);
-    }
+  if (salt.empty()) {
+    salt.resize(builder.hash_size());
 
-    if (!file) {
-        FATAL("failed to read file %s\n", data_filename);
+    int random_fd = open("/dev/urandom", O_RDONLY);
+    if (random_fd < 0) {
+      FATAL("failed to open /dev/urandom\n");
     }
 
-    int64_t len = sparse_file_len(file, false, false);
-    if (len % block_size != 0) {
-        FATAL("file size %" PRIu64 " is not a multiple of %zu bytes\n",
-                len, block_size);
+    ssize_t ret = read(random_fd, salt.data(), salt.size());
+    if (ret != static_cast<ssize_t>(salt.size())) {
+      FATAL("failed to read %zu bytes from /dev/urandom: %zd %d\n", salt.size(),
+            ret, errno);
     }
+    close(random_fd);
+  }
 
-    int levels = 0;
+  // TODO(xunchang) move the size calculation to HashTreeBuilder.
+  if (calculate_size) {
+    if (argc != 0) {
+      usage();
+      return 1;
+    }
     size_t verity_blocks = 0;
     size_t level_blocks;
-
+    size_t levels = 0;
     do {
-        level_blocks = verity_tree_blocks(len, block_size, hash_size, levels);
-        levels++;
-        verity_blocks += level_blocks;
+      level_blocks = verity_tree_blocks(calculate_size, kBlockSize,
+                                        builder.hash_size(), levels);
+      levels++;
+      verity_blocks += level_blocks;
     } while (level_blocks > 1);
 
-    unsigned char *verity_tree = new unsigned char[verity_blocks * block_size]();
-    unsigned char **verity_tree_levels = new unsigned char *[levels + 1]();
-    size_t *verity_tree_level_blocks = new size_t[levels]();
-    if (verity_tree == NULL || verity_tree_levels == NULL || verity_tree_level_blocks == NULL) {
-        FATAL("failed to allocate memory for verity tree\n");
-    }
-
-    unsigned char *ptr = verity_tree;
-    for (int i = levels - 1; i >= 0; i--) {
-        verity_tree_levels[i] = ptr;
-        verity_tree_level_blocks[i] = verity_tree_blocks(len, block_size, hash_size, i);
-        ptr += verity_tree_level_blocks[i] * block_size;
-    }
-    assert(ptr == verity_tree + verity_blocks * block_size);
-    assert(verity_tree_level_blocks[levels - 1] == 1);
-
-    unsigned char zero_block_hash[hash_size];
-    unsigned char zero_block[block_size];
-    memset(zero_block, 0, block_size);
-    hash_block(md, zero_block, block_size, salt.data(), salt.size(), zero_block_hash, NULL);
-
-    unsigned char root_hash[hash_size];
-    verity_tree_levels[levels] = root_hash;
-
-    struct sparse_hash_ctx ctx;
-    ctx.hashes = verity_tree_levels[0];
-    ctx.salt = salt.data();
-    ctx.salt_size = salt.size();
-    ctx.hash_size = hash_size;
-    ctx.block_size = block_size;
-    ctx.zero_block_hash = zero_block_hash;
-    ctx.md = md;
-
-    sparse_file_callback(file, false, false, hash_chunk, &ctx);
-
-    sparse_file_destroy(file);
-    close(fd);
-
-    for (int i = 0; i < levels; i++) {
-        size_t out_size;
-        hash_blocks(md,
-                verity_tree_levels[i], verity_tree_level_blocks[i] * block_size,
-                verity_tree_levels[i + 1], &out_size,
-                salt.data(), salt.size(), block_size);
-          if (i < levels - 1) {
-              assert(div_round_up(out_size, block_size) == verity_tree_level_blocks[i + 1]);
-          } else {
-              assert(out_size == hash_size);
-          }
-    }
-
-    for (size_t i = 0; i < hash_size; i++) {
-        printf("%02x", root_hash[i]);
-    }
-    printf(" ");
-    for (size_t i = 0; i < salt.size(); i++) {
-        printf("%02x", salt[i]);
-    }
-    printf("\n");
-
-    fd = open(verity_filename, O_WRONLY|O_CREAT, 0666);
-    if (fd < 0) {
-        FATAL("failed to open output file '%s'\n", verity_filename);
-    }
-    if (!android::base::WriteFully(fd, verity_tree, verity_blocks * block_size)) {
-        FATAL("failed to write '%s'\n", verity_filename);
-    }
-    close(fd);
-
-    delete[] verity_tree_levels;
-    delete[] verity_tree_level_blocks;
-    delete[] verity_tree;
+    printf("%" PRIu64 "\n", (uint64_t)verity_blocks * kBlockSize);
+    return 0;
+  }
+
+  if (argc != 2) {
+    usage();
+    return 1;
+  }
+
+  data_filename = argv[0];
+  verity_filename = argv[1];
+
+  android::base::unique_fd data_fd(open(data_filename, O_RDONLY));
+  if (data_fd == -1) {
+    FATAL("failed to open %s\n", data_filename);
+  }
+
+  struct sparse_file* file;
+  if (sparse) {
+    file = sparse_file_import(data_fd, false, false);
+  } else {
+    file = sparse_file_import_auto(data_fd, false, verbose);
+  }
+
+  if (!file) {
+    FATAL("failed to read file %s\n", data_filename);
+  }
+
+  int64_t len = sparse_file_len(file, false, false);
+  if (len % kBlockSize != 0) {
+    FATAL("file size %" PRIu64 " is not a multiple of %zu bytes\n", len,
+          kBlockSize);
+  }
+
+  // Initialize the builder to compute the hash tree.
+  if (!builder.Initialize(len, salt)) {
+    LOG(ERROR) << "Failed to initialize HashTreeBuilder";
+    return 1;
+  }
+
+  auto hash_callback = [](void* priv, const void* data, size_t len) {
+    auto sparse_hasher = static_cast<HashTreeBuilder*>(priv);
+    return sparse_hasher->Update(static_cast<const unsigned char*>(data), len)
+               ? 0
+               : 1;
+  };
+  sparse_file_callback(file, false, false, hash_callback, &builder);
+  sparse_file_destroy(file);
+
+  if (!builder.BuildHashTree()) {
+    return 1;
+  }
+
+  if (!builder.WriteHashTreeToFile(verity_filename)) {
+    return 1;
+  }
+
+  // Output the root hash and the salt.
+  for (const auto& c : builder.root_hash()) {
+    printf("%02x", c);
+  }
+  printf(" ");
+  for (const auto& c : salt) {
+    printf("%02x", c);
+  }
+  printf("\n");
+
+  return 0;
 }
diff --git a/verity/build_verity_tree_utils.h b/verity/build_verity_tree_utils.h
new file mode 100644
index 00000000..fbfce7ad
--- /dev/null
+++ b/verity/build_verity_tree_utils.h
@@ -0,0 +1,28 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __BUILD_VERITY_TREE_UTILS_H__
+#define __BUILD_VERITY_TREE_UTILS_H__
+
+#include <inttypes.h>
+#include <stddef.h>
+
+inline uint64_t div_round_up(uint64_t x, uint64_t y) { return (x + y - 1) / y; }
+
+size_t verity_tree_blocks(uint64_t data_size, size_t block_size,
+                          size_t hash_size, size_t level);
+
+#endif  // __BUILD_VERITY_TREE_UTILS_H__
diff --git a/verity/hash_tree_builder.cpp b/verity/hash_tree_builder.cpp
new file mode 100644
index 00000000..3c3692fb
--- /dev/null
+++ b/verity/hash_tree_builder.cpp
@@ -0,0 +1,176 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "hash_tree_builder.h"
+
+#include <android-base/file.h>
+#include <android-base/logging.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()) {
+  CHECK(md_ != nullptr) << "Failed to initialize md";
+
+  hash_size_ = EVP_MD_size(md_);
+  CHECK_LT(hash_size_ * 2, block_size_);
+}
+
+bool HashTreeBuilder::Initialize(int64_t expected_data_size,
+                                 const std::vector<unsigned char>& salt) {
+  data_size_ = expected_data_size;
+  salt_ = salt;
+
+  if (data_size_ % block_size_ != 0) {
+    LOG(ERROR) << "file size " << data_size_
+               << " is not a multiple of block size " << block_size_;
+    return false;
+  }
+
+  // Reserve enough space for the hash of the input data.
+  size_t base_level_blocks =
+      verity_tree_blocks(data_size_, block_size_, hash_size_, 0);
+  std::vector<unsigned char> base_level;
+  base_level.reserve(base_level_blocks * block_size_);
+  verity_tree_.emplace_back(std::move(base_level));
+
+  // Save the hash of the zero block to avoid future recalculation.
+  std::vector<unsigned char> zero_block(block_size_, 0);
+  zero_block_hash_.resize(hash_size_);
+  HashBlock(zero_block.data(), zero_block_hash_.data());
+
+  return true;
+}
+
+bool HashTreeBuilder::HashBlock(const unsigned char* block,
+                                unsigned char* out) {
+  unsigned int s;
+  int ret = 1;
+
+  EVP_MD_CTX* mdctx = EVP_MD_CTX_create();
+  CHECK(mdctx != nullptr);
+  ret &= EVP_DigestInit_ex(mdctx, md_, nullptr);
+  ret &= EVP_DigestUpdate(mdctx, salt_.data(), salt_.size());
+  ret &= EVP_DigestUpdate(mdctx, block, block_size_);
+  ret &= EVP_DigestFinal_ex(mdctx, out, &s);
+  EVP_MD_CTX_destroy(mdctx);
+
+  CHECK_EQ(1, ret);
+  CHECK_EQ(hash_size_, s);
+  return true;
+}
+
+bool HashTreeBuilder::HashBlocks(const unsigned char* data, size_t len,
+                                 std::vector<unsigned char>* output_vector) {
+  if (len == 0) {
+    return true;
+  }
+  CHECK_EQ(0, len % block_size_);
+
+  if (data == nullptr) {
+    for (size_t i = 0; i < len; i += block_size_) {
+      output_vector->insert(output_vector->end(), zero_block_hash_.begin(),
+                            zero_block_hash_.end());
+    }
+    return true;
+  }
+
+  for (size_t i = 0; i < len; i += block_size_) {
+    unsigned char hash_buffer[hash_size_];
+    if (!HashBlock(data + i, hash_buffer)) {
+      return false;
+    }
+    output_vector->insert(output_vector->end(), hash_buffer,
+                          hash_buffer + hash_size_);
+  }
+
+  return true;
+}
+
+bool HashTreeBuilder::Update(const unsigned char* data, size_t len) {
+  CHECK_GT(data_size_, 0);
+
+  return HashBlocks(data, len, &verity_tree_[0]);
+}
+
+bool HashTreeBuilder::BuildHashTree() {
+  // Expects only the base level in the verity_tree_.
+  CHECK_EQ(1, verity_tree_.size());
+
+  // Expects the base level to have the same size as the total hash size of
+  // input data.
+  AppendPaddings(&verity_tree_.back());
+  size_t base_level_blocks =
+      verity_tree_blocks(data_size_, block_size_, hash_size_, 0);
+  CHECK_EQ(base_level_blocks * block_size_, verity_tree_[0].size());
+
+  while (verity_tree_.back().size() > block_size_) {
+    const auto& current_level = verity_tree_.back();
+    // Computes the next level of the verity tree based on the hash of the
+    // current level.
+    size_t next_level_blocks =
+        verity_tree_blocks(current_level.size(), block_size_, hash_size_, 0);
+    std::vector<unsigned char> next_level;
+    next_level.reserve(next_level_blocks * block_size_);
+
+    HashBlocks(current_level.data(), current_level.size(), &next_level);
+    AppendPaddings(&next_level);
+
+    // Checks the size of the next level.
+    CHECK_EQ(next_level_blocks * block_size_, next_level.size());
+    verity_tree_.emplace_back(std::move(next_level));
+  }
+
+  CHECK_EQ(block_size_, verity_tree_.back().size());
+  HashBlocks(verity_tree_.back().data(), block_size_, &root_hash_);
+
+  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));
+  if (output_fd == -1) {
+    PLOG(ERROR) << "Failed to open output file " << output;
+    return false;
+  }
+
+  return WriteHashTreeToFd(output_fd);
+}
+
+bool HashTreeBuilder::WriteHashTreeToFd(int fd) const {
+  CHECK(!verity_tree_.empty());
+
+  // 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())) {
+      PLOG(ERROR) << "Failed to write the hash tree level " << i;
+      return false;
+    }
+  }
+
+  return true;
+}
+
+void HashTreeBuilder::AppendPaddings(std::vector<unsigned char>* data) {
+  size_t remainder = data->size() % block_size_;
+  if (remainder != 0) {
+    data->resize(data->size() + block_size_ - remainder, 0);
+  }
+}
diff --git a/verity/hash_tree_builder.h b/verity/hash_tree_builder.h
new file mode 100644
index 00000000..b83f8791
--- /dev/null
+++ b/verity/hash_tree_builder.h
@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __HASH_TREE_BUILDER_H__
+#define __HASH_TREE_BUILDER_H__
+
+#include <inttypes.h>
+#include <stddef.h>
+
+#include <string>
+#include <vector>
+
+#include <openssl/evp.h>
+
+// This class builds a verity hash tree based on the input data and a salt with
+// the length of hash size. It also supports the streaming of input data while
+// 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);
+  // Gets ready for the hash tree computation. We expect |expected_data_size|
+  // bytes source data.
+  bool Initialize(int64_t expected_data_size,
+                  const std::vector<unsigned char>& salt);
+  // Streams |len| bytes of source data to the hash tree builder, and the |len|
+  // is expected to be block aligned. This function can be called multiple until
+  // we processed all the source data. And the accumulated data_size is expected
+  // to be exactly the |data_size_| when we build the hash tree.
+  bool Update(const unsigned char* data, size_t len);
+  // Computes the upper levels of the hash tree based on the 0th level.
+  bool BuildHashTree();
+  // Writes the computed hash tree top-down to |output|.
+  bool WriteHashTreeToFile(const std::string& output) const;
+  bool WriteHashTreeToFd(int fd) const;
+
+  size_t hash_size() const { return hash_size_; }
+  const std::vector<unsigned char>& root_hash() const { return root_hash_; }
+
+ private:
+  // Calculates the hash of one single block. Write the result to |out|, a
+  // buffer allocated by the caller.
+  bool HashBlock(const unsigned char* block, unsigned char* out);
+  // Calculates the hash of |len| bytes of data starting from |data|. Append the
+  // result to |output_vector|.
+  bool HashBlocks(const unsigned char* data, size_t len,
+                  std::vector<unsigned char>* output_vector);
+  // Aligns |data| with block_size by padding 0s to the end.
+  void AppendPaddings(std::vector<unsigned char>* data);
+
+  size_t block_size_;
+  // Expected size of the source data, which is used to compute the hash for the
+  // base level.
+  uint64_t data_size_;
+  std::vector<unsigned char> salt_;
+  const EVP_MD* md_;
+  size_t hash_size_;
+
+  // Pre-calculated hash of a zero block.
+  std::vector<unsigned char> zero_block_hash_;
+  std::vector<unsigned char> root_hash_;
+  // Storage of the verity tree. The base level hash stores in verity_tree_[0]
+  // and the top level hash stores in verity_tree_.back().
+  std::vector<std::vector<unsigned char>> verity_tree_;
+};
+
+#endif  // __HASH_TREE_BUILDER_H__