verity: Support partial blocks in Update().

In update_engine we get data from a callback, usually this will be a
multiple of the block size, but that's not guaranteed.

Bug: 28171891
Test: build_verity_tree_test
Change-Id: I2fdf40e6e3750a86d409885917d3b5aeb77643d1

3 files changed, 60 insertions, 4 deletions

diff --git a/verity/build_verity_tree_test.cpp b/verity/build_verity_tree_test.cpp
index c1cb6a39..864ebdc9 100644
--- a/verity/build_verity_tree_test.cpp
+++ b/verity/build_verity_tree_test.cpp
@@ -135,6 +135,28 @@ TEST_F(BuildVerityTreeTest, StreamingDataMultipleBlocks) {
             HashTreeBuilder::BytesArrayToString(builder->root_hash()));
 }
 
+TEST_F(BuildVerityTreeTest, StreamingDataPartialBlocks) {
+  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() % 40960, 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()));
+}
+
 TEST_F(BuildVerityTreeTest, SHA1MultipleBlocks) {
   std::vector<unsigned char> data(128 * 4096, 0xff);
 
diff --git a/verity/hash_tree_builder.cpp b/verity/hash_tree_builder.cpp
index a776572f..197c38e5 100644
--- a/verity/hash_tree_builder.cpp
+++ b/verity/hash_tree_builder.cpp
@@ -182,6 +182,32 @@ bool HashTreeBuilder::HashBlocks(const unsigned char* data, size_t len,
 bool HashTreeBuilder::Update(const unsigned char* data, size_t len) {
   CHECK_GT(data_size_, 0);
 
+  if (!leftover_.empty()) {
+    CHECK_LT(leftover_.size(), block_size_);
+    size_t append_len = std::min(len, block_size_ - leftover_.size());
+    if (data == nullptr) {
+      leftover_.insert(leftover_.end(), append_len, 0);
+    } else {
+      leftover_.insert(leftover_.end(), data, data + append_len);
+    }
+    if (leftover_.size() < block_size_) {
+      return true;
+    }
+    if (!HashBlocks(leftover_.data(), leftover_.size(), &verity_tree_[0])) {
+      return false;
+    }
+    leftover_.clear();
+    data += append_len;
+    len -= append_len;
+  }
+  if (len % block_size_ != 0) {
+    if (data == nullptr) {
+      leftover_.assign(len % block_size_, 0);
+    } else {
+      leftover_.assign(data + len - len % block_size_, data + len);
+    }
+    len -= len % block_size_;
+  }
   return HashBlocks(data, len, &verity_tree_[0]);
 }
 
@@ -189,6 +215,11 @@ bool HashTreeBuilder::BuildHashTree() {
   // Expects only the base level in the verity_tree_.
   CHECK_EQ(1, verity_tree_.size());
 
+  if (!leftover_.empty()) {
+    LOG(ERROR) << leftover_.size() << " bytes data left from last Update().";
+    return false;
+  }
+
   // Expects the base level to have the same size as the total hash size of
   // input data.
   AppendPaddings(&verity_tree_.back());
diff --git a/verity/include/verity/hash_tree_builder.h b/verity/include/verity/hash_tree_builder.h
index c6f00c0e..da47d0fe 100644
--- a/verity/include/verity/hash_tree_builder.h
+++ b/verity/include/verity/hash_tree_builder.h
@@ -39,10 +39,10 @@ class HashTreeBuilder {
   // 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.
+  // Streams |len| bytes of source data to the hash tree builder. 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();
@@ -94,6 +94,9 @@ class HashTreeBuilder {
   // 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_;
+  // The remaining data passed to the last call to Update() that's less than a
+  // block.
+  std::vector<unsigned char> leftover_;
 };
 
 #endif  // __HASH_TREE_BUILDER_H__