Move functions in build_verity_tree to a static library

Move the functions to calculate the hash tree size & generate the tree
to a seperate static library. Also switch to libbase logging.

Bug: 25170618
Test: run build_verity_tree to calculate size & generate a tree
Change-Id: I88a2ce077292e35fe069d00000a9e59dbeb95317

7 files changed, 318 insertions, 200 deletions

diff --git a/verity/Android.bp b/verity/Android.bp
index 3bc6e754..e2594d7f 100644
--- a/verity/Android.bp
+++ b/verity/Android.bp
@@ -66,14 +66,39 @@ cc_binary_host {
     ],
 }
 
+cc_library_static {
+    name: "libverity_tree",
+    srcs: [
+        "build_verity_tree.cpp",
+        "build_verity_tree_utils.cpp",
+        "hash_tree_builder.cpp",
+    ],
+
+    host_supported: true,
+
+    static_libs: [
+        "libsparse",
+        "libz",
+    ],
+    shared_libs: [
+        "libcrypto",
+        "libbase",
+    ],
+    cflags: [
+        "-D_FILE_OFFSET_BITS=64",
+        "-Wall",
+        "-Werror",
+    ],
+}
+
 cc_binary_host {
     name: "build_verity_tree",
     srcs: [
-        "hash_tree_builder.cpp",
-        "build_verity_tree.cpp"
+        "build_verity_tree_main.cpp"
     ],
 
     static_libs: [
+        "libverity_tree",
         "libsparse",
         "libz",
     ],
@@ -82,6 +107,7 @@ cc_binary_host {
         "libbase",
     ],
     cflags: [
+        "-D_FILE_OFFSET_BITS=64",
         "-Wall",
         "-Werror",
     ],
diff --git a/verity/build_verity_tree.cpp b/verity/build_verity_tree.cpp
index e9978cd8..27f9255f 100644
--- a/verity/build_verity_tree.cpp
+++ b/verity/build_verity_tree.cpp
@@ -14,183 +14,23 @@
  * limitations under the License.
  */
 
-#include <openssl/bn.h>
-#include <sparse/sparse.h>
-
-#undef NDEBUG
-
-#include <errno.h>
-#include <getopt.h>
-#include <fcntl.h>
-#include <inttypes.h>
-#include <stdio.h>
-#include <string.h>
-#include <unistd.h>
-
-#include <limits>
-#include <vector>
+#include "build_verity_tree.h"
 
-#include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/unique_fd.h>
+#include <sparse/sparse.h>
 
-#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, 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);
-
-  do {
-    level_blocks = div_round_up(level_blocks, hashes_per_block);
-  } while (level--);
-
-  CHECK_LE(level_blocks, std::numeric_limits<size_t>::max());
-  return level_blocks;
-}
-
-void usage(void)
-{
-    printf("usage: build_verity_tree [ <options> ] -s <size> | <data> <verity>\n"
-           "options:\n"
-           "  -a,--salt-str=<string>       set salt to <string>\n"
-           "  -A,--salt-hex=<hex digits>   set salt to <hex digits>\n"
-           "  -h                           show this help\n"
-           "  -s,--verity-size=<data size> print the size of the verity tree\n"
-           "  -v,                          enable verbose logging\n"
-           "  -S                           treat <data image> as a sparse file\n"
-        );
-}
-
-int main(int argc, char **argv)
-{
-  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");
-        }
-        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();
-    }
-  }
-
-  argc -= optind;
-  argv += optind;
-
-  HashTreeBuilder builder(kBlockSize);
-
-  if (salt.empty()) {
-    salt.resize(builder.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);
-    }
-    close(random_fd);
-  }
-
-  // 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(calculate_size, kBlockSize,
-                                        builder.hash_size(), levels);
-      levels++;
-      verity_blocks += level_blocks;
-    } while (level_blocks > 1);
-
-    printf("%" PRIu64 "\n", (uint64_t)verity_blocks * kBlockSize);
-    return 0;
-  }
-
-  if (argc != 2) {
-    usage();
-    return 1;
-  }
-
-  data_filename = argv[0];
-  verity_filename = argv[1];
+#undef NDEBUG
 
-  android::base::unique_fd data_fd(open(data_filename, O_RDONLY));
+bool generate_verity_tree(const std::string& data_filename,
+                          const std::string& verity_filename,
+                          HashTreeBuilder* builder,
+                          const std::vector<unsigned char>& salt_content,
+                          size_t block_size, bool sparse, bool verbose) {
+  android::base::unique_fd data_fd(open(data_filename.c_str(), O_RDONLY));
   if (data_fd == -1) {
-    FATAL("failed to open %s\n", data_filename);
+    PLOG(ERROR) << "failed to open " << data_filename;
+    return false;
   }
 
   struct sparse_file* file;
@@ -201,19 +41,21 @@ int main(int argc, char **argv)
   }
 
   if (!file) {
-    FATAL("failed to read file %s\n", data_filename);
+    LOG(ERROR) << "failed to read file " << data_filename;
+    return false;
   }
 
   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);
+  if (len % block_size != 0) {
+    LOG(ERROR) << "file size " << len << " is not a multiple of " << block_size
+               << " byte";
+    return false;
   }
 
   // Initialize the builder to compute the hash tree.
-  if (!builder.Initialize(len, salt)) {
+  if (!builder->Initialize(len, salt_content)) {
     LOG(ERROR) << "Failed to initialize HashTreeBuilder";
-    return 1;
+    return false;
   }
 
   auto hash_callback = [](void* priv, const void* data, size_t len) {
@@ -222,26 +64,12 @@ int main(int argc, char **argv)
                ? 0
                : 1;
   };
-  sparse_file_callback(file, false, false, hash_callback, &builder);
+  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);
+  if (!builder->BuildHashTree()) {
+    return false;
   }
-  printf("\n");
 
-  return 0;
+  return builder->WriteHashTreeToFile(verity_filename);
 }
diff --git a/verity/build_verity_tree.h b/verity/build_verity_tree.h
new file mode 100644
index 00000000..9ca40005
--- /dev/null
+++ b/verity/build_verity_tree.h
@@ -0,0 +1,33 @@
+/*
+ * 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_H__
+#define __BUILD_VERITY_TREE_H__
+
+#include <inttypes.h>
+
+#include <string>
+#include <vector>
+
+#include "hash_tree_builder.h"
+
+bool generate_verity_tree(const std::string& data_filename,
+                          const std::string& verity_filename,
+                          HashTreeBuilder* hasher,
+                          const std::vector<unsigned char>& salt_content,
+                          size_t block_size, bool sparse, bool verbose);
+
+#endif  // __BUILD_VERITY_TREE_H__
diff --git a/verity/build_verity_tree_main.cpp b/verity/build_verity_tree_main.cpp
new file mode 100644
index 00000000..3635a975
--- /dev/null
+++ b/verity/build_verity_tree_main.cpp
@@ -0,0 +1,163 @@
+/*
+ * 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 <fcntl.h>
+#include <getopt.h>
+#include <inttypes.h>
+#include <stdio.h>
+#include <string.h>
+#include <unistd.h>
+
+#include <limits>
+#include <string>
+#include <vector>
+
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/parseint.h>
+#include <android-base/unique_fd.h>
+#include <openssl/bn.h>
+
+#include "build_verity_tree.h"
+
+static void usage(void) {
+  printf(
+      "usage: build_verity_tree [ <options> ] -s <size> | <data> <verity>\n"
+      "options:\n"
+      "  -a,--salt-str=<string>       set salt to <string>\n"
+      "  -A,--salt-hex=<hex digits>   set salt to <hex digits>\n"
+      "  -h                           show this help\n"
+      "  -s,--verity-size=<data size> print the size of the verity tree\n"
+      "  -v,                          enable verbose logging\n"
+      "  -S                           treat <data image> as a sparse file\n");
+}
+
+int main(int argc, char** argv) {
+  constexpr size_t kBlockSize = 4096;
+
+  std::vector<unsigned char> salt;
+  bool sparse = false;
+  uint64_t calculate_size = 0;
+  bool verbose = false;
+
+  while (1) {
+    constexpr 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'},
+        {nullptr, 0, 0, 0}};
+    int c = getopt_long(argc, argv, "a:A:hSs:v", long_options, nullptr);
+    if (c < 0) {
+      break;
+    }
+
+    switch (c) {
+      case 'a':
+        salt.clear();
+        salt.insert(salt.end(), optarg, &optarg[strlen(optarg)]);
+        break;
+      case 'A': {
+        BIGNUM* bn = nullptr;
+        if (!BN_hex2bn(&bn, optarg)) {
+          LOG(ERROR) << "Failed to convert salt from hex";
+          return 1;
+        }
+        size_t salt_size = BN_num_bytes(bn);
+        salt.resize(salt_size);
+        if (BN_bn2bin(bn, salt.data()) != salt_size) {
+          LOG(ERROR) << "Failed to convert salt to bytes";
+          return 1;
+        }
+      } break;
+      case 'h':
+        usage();
+        return 1;
+      case 'S':
+        sparse = true;
+        break;
+      case 's': {
+        if (!android::base::ParseUint(optarg, &calculate_size,
+                                      std::numeric_limits<uint64_t>::max())) {
+          LOG(ERROR) << "Invalid input size: " << optarg;
+          return 1;
+        }
+
+      } break;
+      case 'v':
+        verbose = true;
+        break;
+      case '?':
+        usage();
+        return 1;
+      default:
+        abort();
+    }
+  }
+
+  argc -= optind;
+  argv += optind;
+
+  HashTreeBuilder builder(kBlockSize);
+
+  if (calculate_size) {
+    if (argc != 0) {
+      usage();
+      return 1;
+    }
+
+    uint64_t tree_size = builder.CalculateSize(calculate_size);
+    printf("%" PRIu64 "\n", tree_size);
+    return 0;
+  }
+
+  if (argc != 2) {
+    usage();
+    return 1;
+  }
+
+  if (salt.empty()) {
+    salt.resize(builder.hash_size());
+
+    android::base::unique_fd random_fd(open("/dev/urandom", O_RDONLY));
+    if (random_fd < 0) {
+      PLOG(ERROR) << "failed to open /dev/urandom";
+      return 1;
+    }
+
+    ssize_t ret = read(random_fd, salt.data(), salt.size());
+    if (ret != static_cast<ssize_t>(salt.size())) {
+      PLOG(ERROR) << "failed to read " << salt.size()
+                  << " bytes from /dev/urandom: " << ret;
+      return 1;
+    }
+  }
+
+  if (!generate_verity_tree(argv[0], argv[1], &builder, salt, kBlockSize,
+                            sparse, verbose)) {
+    return 1;
+  }
+
+  // Output the root hash and the salt.
+  std::string root_hash_string =
+      HashTreeBuilder::BytesArrayToString(builder.root_hash());
+  std::string salt_string = HashTreeBuilder::BytesArrayToString(salt);
+  printf("%s %s\n", root_hash_string.c_str(), salt_string.c_str());
+
+  return 0;
+}
diff --git a/verity/build_verity_tree_utils.cpp b/verity/build_verity_tree_utils.cpp
new file mode 100644
index 00000000..2b3005f9
--- /dev/null
+++ b/verity/build_verity_tree_utils.cpp
@@ -0,0 +1,34 @@
+/*
+ * 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 "build_verity_tree_utils.h"
+
+#include <limits>
+
+#include <android-base/logging.h>
+
+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);
+
+  do {
+    level_blocks = div_round_up(level_blocks, hashes_per_block);
+  } while (level--);
+
+  CHECK_LE(level_blocks, std::numeric_limits<size_t>::max());
+  return level_blocks;
+}
\ No newline at end of file
diff --git a/verity/hash_tree_builder.cpp b/verity/hash_tree_builder.cpp
index 3c3692fb..4b06a5f2 100644
--- a/verity/hash_tree_builder.cpp
+++ b/verity/hash_tree_builder.cpp
@@ -18,6 +18,7 @@
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/stringprintf.h>
 #include <android-base/unique_fd.h>
 
 #include "build_verity_tree_utils.h"
@@ -30,6 +31,29 @@ HashTreeBuilder::HashTreeBuilder(size_t block_size)
   CHECK_LT(hash_size_ * 2, block_size_);
 }
 
+std::string HashTreeBuilder::BytesArrayToString(
+    const std::vector<unsigned char>& bytes) {
+  std::string result;
+  for (const auto& c : bytes) {
+    result += android::base::StringPrintf("%02x", c);
+  }
+  return result;
+}
+
+uint64_t HashTreeBuilder::CalculateSize(uint64_t input_size) const {
+  size_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);
+    levels++;
+    verity_blocks += level_blocks;
+  } while (level_blocks > 1);
+
+  return verity_blocks * block_size_;
+}
+
 bool HashTreeBuilder::Initialize(int64_t expected_data_size,
                                  const std::vector<unsigned char>& salt) {
   data_size_ = expected_data_size;
@@ -149,12 +173,17 @@ bool HashTreeBuilder::WriteHashTreeToFile(const std::string& output) const {
     return false;
   }
 
-  return WriteHashTreeToFd(output_fd);
+  return WriteHashTreeToFd(output_fd, 0);
 }
 
-bool HashTreeBuilder::WriteHashTreeToFd(int fd) const {
+bool HashTreeBuilder::WriteHashTreeToFd(int fd, off_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;
+  }
+
   // 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];
diff --git a/verity/hash_tree_builder.h b/verity/hash_tree_builder.h
index b83f8791..bb217a1d 100644
--- a/verity/hash_tree_builder.h
+++ b/verity/hash_tree_builder.h
@@ -34,6 +34,8 @@
 class HashTreeBuilder {
  public:
   explicit HashTreeBuilder(size_t block_size);
+  // 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|
   // bytes source data.
   bool Initialize(int64_t expected_data_size,
@@ -47,10 +49,13 @@ class HashTreeBuilder {
   bool BuildHashTree();
   // Writes the computed hash tree top-down to |output|.
   bool WriteHashTreeToFile(const std::string& output) const;
-  bool WriteHashTreeToFd(int fd) const;
+  bool WriteHashTreeToFd(int fd, off_t offset) const;
 
   size_t hash_size() const { return hash_size_; }
   const std::vector<unsigned char>& root_hash() const { return root_hash_; }
+  // Converts |bytes| to string for hexdump.
+  static std::string BytesArrayToString(
+      const std::vector<unsigned char>& bytes);
 
  private:
   // Calculates the hash of one single block. Write the result to |out|, a