Add unit tests to build_verity_tree

Also factor out the cc_defaults in the bp file.

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

4 files changed, 175 insertions, 23 deletions

diff --git a/verity/Android.bp b/verity/Android.bp
index e2594d7f..d1d81fac 100644
--- a/verity/Android.bp
+++ b/verity/Android.bp
@@ -66,16 +66,14 @@ cc_binary_host {
     ],
 }
 
-cc_library_static {
-    name: "libverity_tree",
-    srcs: [
-        "build_verity_tree.cpp",
-        "build_verity_tree_utils.cpp",
-        "hash_tree_builder.cpp",
+cc_defaults {
+    name: "verity_tree_defaults",
+    cflags: [
+        "-D_FILE_OFFSET_BITS=64",
+        "-Wall",
+        "-Werror",
     ],
 
-    host_supported: true,
-
     static_libs: [
         "libsparse",
         "libz",
@@ -84,31 +82,49 @@ cc_library_static {
         "libcrypto",
         "libbase",
     ],
-    cflags: [
-        "-D_FILE_OFFSET_BITS=64",
-        "-Wall",
-        "-Werror",
+
+    host_supported: true,
+}
+
+cc_library_static {
+    name: "libverity_tree",
+    defaults: [
+        "verity_tree_defaults",
+    ],
+
+    srcs: [
+        "build_verity_tree.cpp",
+        "build_verity_tree_utils.cpp",
+        "hash_tree_builder.cpp",
     ],
 }
 
 cc_binary_host {
     name: "build_verity_tree",
+    defaults: [
+        "verity_tree_defaults",
+    ],
+
     srcs: [
-        "build_verity_tree_main.cpp"
+        "build_verity_tree_main.cpp",
     ],
 
     static_libs: [
         "libverity_tree",
-        "libsparse",
-        "libz",
     ],
-    shared_libs: [
-        "libcrypto",
-        "libbase",
+}
+
+cc_test {
+    name: "build_verity_tree_test",
+    defaults: [
+        "verity_tree_defaults",
     ],
-    cflags: [
-        "-D_FILE_OFFSET_BITS=64",
-        "-Wall",
-        "-Werror",
+
+    srcs: [
+        "build_verity_tree_test.cpp",
+    ],
+
+    static_libs: [
+        "libverity_tree",
     ],
 }
diff --git a/verity/build_verity_tree_test.cpp b/verity/build_verity_tree_test.cpp
new file mode 100644
index 00000000..36aa34cb
--- /dev/null
+++ b/verity/build_verity_tree_test.cpp
@@ -0,0 +1,135 @@
+/*
+ * 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 <stdlib.h>
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include "build_verity_tree_utils.h"
+#include "hash_tree_builder.h"
+
+// The hex string we are using in build_image.py
+// aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7
+constexpr unsigned char kSaltHex[] = {
+    0xae, 0xe0, 0x87, 0xa5, 0xbe, 0x3b, 0x98, 0x29, 0x78, 0xc9, 0x23,
+    0xf5, 0x66, 0xa9, 0x46, 0x13, 0x49, 0x6b, 0x41, 0x7f, 0x2a, 0xf5,
+    0x92, 0x63, 0x9b, 0xc8, 0x0d, 0x14, 0x1e, 0x34, 0xdf, 0xe7};
+
+class BuildVerityTreeTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    salt_hex =
+        std::vector<unsigned char>(kSaltHex, kSaltHex + sizeof(kSaltHex));
+    builder.reset(new HashTreeBuilder(4096));
+  }
+
+  const std::vector<unsigned char>& zero_block_hash() const {
+    return builder->zero_block_hash_;
+  }
+  const std::vector<std::vector<unsigned char>>& verity_tree() const {
+    return builder->verity_tree_;
+  }
+
+  void GenerateHashTree(const std::vector<unsigned char>& data,
+                        const std::vector<unsigned char>& salt) {
+    ASSERT_TRUE(builder->Initialize(data.size(), salt));
+    ASSERT_TRUE(builder->Update(data.data(), data.size()));
+    ASSERT_TRUE(builder->BuildHashTree());
+  }
+
+  std::vector<unsigned char> salt_hex;
+  std::unique_ptr<HashTreeBuilder> builder;
+};
+
+TEST_F(BuildVerityTreeTest, CalculateSize) {
+  // The hash of source data will occupy at least one block
+  ASSERT_EQ(4096u, builder->CalculateSize(1));
+  // Sha256 hash of 128 blocks fits into one block.
+  ASSERT_EQ(4096u, builder->CalculateSize(128 * 4096));
+  // Requires 3 blocks to hold more data.
+  ASSERT_EQ(12288u, builder->CalculateSize(128 * 4096 + 1));
+
+  ASSERT_EQ(20811776u, builder->CalculateSize(2641915904));
+}
+
+TEST_F(BuildVerityTreeTest, InitializeBuilder) {
+  // data_size should be divisible by 4096
+  ASSERT_FALSE(builder->Initialize(4095, salt_hex));
+
+  ASSERT_TRUE(builder->Initialize(4096, salt_hex));
+  ASSERT_EQ(1u, verity_tree().size());
+  ASSERT_EQ("6eb8c4e1bce842d137f18b27beb857d3b43899d178090537ad7a0fbe3bf4126a",
+            HashTreeBuilder::BytesArrayToString(zero_block_hash()));
+}
+
+TEST_F(BuildVerityTreeTest, HashSingleBlock) {
+  std::vector<unsigned char> data(4096, 1);
+
+  GenerateHashTree(data, salt_hex);
+
+  ASSERT_EQ(1u, verity_tree().size());
+  ASSERT_EQ("e69eb527b16f933483768e92de9bca45f6cc09208525d408436bb362eb865d32",
+            HashTreeBuilder::BytesArrayToString(builder->root_hash()));
+}
+
+TEST_F(BuildVerityTreeTest, HashSingleLevel) {
+  std::vector<unsigned char> data(128 * 4096, 0);
+
+  GenerateHashTree(data, salt_hex);
+
+  ASSERT_EQ(1u, verity_tree().size());
+  ASSERT_EQ("62a4fbe8c9036168ba77fe3e3fd78dd6ed963aeb8aaaa36e84f5c7f9107c6b78",
+            HashTreeBuilder::BytesArrayToString(builder->root_hash()));
+}
+
+TEST_F(BuildVerityTreeTest, HashMultipleLevels) {
+  std::vector<unsigned char> data(129 * 4096, 0xff);
+
+  GenerateHashTree(data, salt_hex);
+
+  ASSERT_EQ(2u, verity_tree().size());
+  ASSERT_EQ(2 * 4096u, verity_tree()[0].size());
+  ASSERT_EQ("9e74f2d47a990c276093760f01de5e9039883e808286ee9492c9cafe9e4ff825",
+            HashTreeBuilder::BytesArrayToString(builder->root_hash()));
+}
+
+TEST_F(BuildVerityTreeTest, StreamingDataMultipleBlocks) {
+  std::vector<unsigned char> data(256 * 4096);
+  for (size_t i = 0; i < 256; i++) {
+    std::fill_n(data.begin() + i * 4096, 4096, i);
+  }
+
+  ASSERT_TRUE(builder->Initialize(data.size(), salt_hex));
+
+  size_t offset = 0;
+  while (offset < data.size()) {
+    size_t data_length =
+        std::min<size_t>(rand() % 10 * 4096, data.size() - offset);
+    ASSERT_TRUE(builder->Update(data.data() + offset, data_length));
+    offset += data_length;
+  }
+
+  ASSERT_TRUE(builder->BuildHashTree());
+  ASSERT_EQ(2u, verity_tree().size());
+  ASSERT_EQ(2 * 4096u, verity_tree()[0].size());
+  ASSERT_EQ("6e73d59b0b6baf026e921814979a7db02244c95a46b869a17aa1310dad066deb",
+            HashTreeBuilder::BytesArrayToString(builder->root_hash()));
+}
diff --git a/verity/hash_tree_builder.cpp b/verity/hash_tree_builder.cpp
index 4b06a5f2..05c06ada 100644
--- a/verity/hash_tree_builder.cpp
+++ b/verity/hash_tree_builder.cpp
@@ -41,7 +41,7 @@ std::string HashTreeBuilder::BytesArrayToString(
 }
 
 uint64_t HashTreeBuilder::CalculateSize(uint64_t input_size) const {
-  size_t verity_blocks = 0;
+  uint64_t verity_blocks = 0;
   size_t level_blocks;
   size_t levels = 0;
   do {
diff --git a/verity/hash_tree_builder.h b/verity/hash_tree_builder.h
index bb217a1d..8f180f4b 100644
--- a/verity/hash_tree_builder.h
+++ b/verity/hash_tree_builder.h
@@ -58,6 +58,7 @@ class HashTreeBuilder {
       const std::vector<unsigned char>& bytes);
 
  private:
+  friend class BuildVerityTreeTest;
   // 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);