Add checks before writing hashtree/verity

Verity that we read everything in hashtree_data_extent before writing
hash tree.

Bug: 173432386
Test: th
Change-Id: I00ab8053de71b13991adaa243b6cb6c7efd6e60f
diff --git a/payload_consumer/filesystem_verifier_action_unittest.cc b/payload_consumer/filesystem_verifier_action_unittest.cc
index f618626..7d3ee9e 100644
--- a/payload_consumer/filesystem_verifier_action_unittest.cc
+++ b/payload_consumer/filesystem_verifier_action_unittest.cc
@@ -33,6 +33,7 @@
 #include "update_engine/common/mock_dynamic_partition_control.h"
 #include "update_engine/common/test_utils.h"
 #include "update_engine/common/utils.h"
+#include "update_engine/payload_consumer/fake_file_descriptor.h"
 #include "update_engine/payload_consumer/install_plan.h"
 
 using brillo::MessageLoop;
@@ -251,8 +252,8 @@
   processor_.set_delegate(&delegate);
 
   processor_.StartProcessing();
-  EXPECT_FALSE(processor_.IsRunning());
-  EXPECT_TRUE(delegate.ran_);
+  ASSERT_FALSE(processor_.IsRunning());
+  ASSERT_TRUE(delegate.ran_);
   EXPECT_EQ(ErrorCode::kError, delegate.code_);
 }
 
@@ -397,12 +398,12 @@
           base::Unretained(&processor_)));
   loop_.Run();
 
-  EXPECT_FALSE(processor_.IsRunning());
-  EXPECT_TRUE(delegate.ran());
-  EXPECT_EQ(ErrorCode::kSuccess, delegate.code());
+  ASSERT_FALSE(processor_.IsRunning());
+  ASSERT_TRUE(delegate.ran());
+  ASSERT_EQ(ErrorCode::kSuccess, delegate.code());
 }
 
-TEST_F(FilesystemVerifierActionTest, RunWithVABC) {
+TEST_F(FilesystemVerifierActionTest, RunWithVABCNoVerity) {
   InstallPlan install_plan;
   InstallPlan::Partition& part = install_plan.partitions.emplace_back();
   part.name = "fake_part";
@@ -410,15 +411,20 @@
   part.target_size = 4096 * 4096;
   part.block_size = 4096;
   part.source_path = "/dev/fake_source_path";
+  part.fec_size = 0;
+  part.hash_tree_size = 0;
+  part.target_hash.clear();
+  part.source_hash.clear();
 
   NiceMock<MockDynamicPartitionControl> dynamic_control;
+  auto fake_fd = std::make_shared<FakeFileDescriptor>();
 
   ON_CALL(dynamic_control, GetDynamicPartitionsFeatureFlag())
       .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
   ON_CALL(dynamic_control, UpdateUsesSnapshotCompression())
       .WillByDefault(Return(true));
   ON_CALL(dynamic_control, OpenCowReader(_, _, _))
-      .WillByDefault(Return(nullptr));
+      .WillByDefault(Return(fake_fd));
   ON_CALL(dynamic_control, IsDynamicPartition(part.name, _))
       .WillByDefault(Return(true));
 
@@ -443,10 +449,18 @@
           base::Unretained(&processor_)));
   loop_.Run();
 
-  EXPECT_FALSE(processor_.IsRunning());
-  EXPECT_TRUE(delegate.ran());
-  // Filesystem verifier will fail, because we returned nullptr as CowReader
-  EXPECT_EQ(ErrorCode::kFilesystemVerifierError, delegate.code());
+  ASSERT_FALSE(processor_.IsRunning());
+  ASSERT_TRUE(delegate.ran());
+  // Filesystem verifier will fail, because we set an empty hash
+  ASSERT_EQ(ErrorCode::kNewRootfsVerificationError, delegate.code());
+  const auto& read_pos = fake_fd->GetReadOps();
+  size_t expected_offset = 0;
+  for (const auto& [off, size] : read_pos) {
+    ASSERT_EQ(off, expected_offset);
+    expected_offset += size;
+  }
+  const auto actual_read_size = expected_offset;
+  ASSERT_EQ(actual_read_size, part.target_size);
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/verity_writer_android.cc b/payload_consumer/verity_writer_android.cc
index 01d8977..4b23b83 100644
--- a/payload_consumer/verity_writer_android.cc
+++ b/payload_consumer/verity_writer_android.cc
@@ -67,12 +67,18 @@
       return false;
     }
   }
+  total_offset_ = 0;
   return true;
 }
 
-bool VerityWriterAndroid::Update(uint64_t offset,
+bool VerityWriterAndroid::Update(const uint64_t offset,
                                  const uint8_t* buffer,
                                  size_t size) {
+  if (offset != total_offset_) {
+    LOG(ERROR) << "Sequential read expected, expected to read at: "
+               << total_offset_ << " actual read occurs at: " << offset;
+    return false;
+  }
   if (partition_->hash_tree_size != 0) {
     const uint64_t hash_tree_data_end =
         partition_->hash_tree_data_offset + partition_->hash_tree_data_size;
@@ -96,12 +102,22 @@
       }
     }
   }
+  total_offset_ += size;
 
   return true;
 }
 
 bool VerityWriterAndroid::Finalize(FileDescriptorPtr read_fd,
                                    FileDescriptorPtr write_fd) {
+  const auto hash_tree_data_end =
+      partition_->hash_tree_data_offset + partition_->hash_tree_data_size;
+  if (total_offset_ < hash_tree_data_end) {
+    LOG(ERROR) << "Read up to " << total_offset_
+               << " when we are expecting to read everything "
+                  "before "
+               << hash_tree_data_end;
+    return false;
+  }
   // All hash tree data blocks has been hashed, write hash tree to disk.
   LOG(INFO) << "Writing verity hash tree to " << partition_->target_path;
   TEST_AND_RETURN_FALSE(hash_tree_builder_->BuildHashTree());
diff --git a/payload_consumer/verity_writer_android.h b/payload_consumer/verity_writer_android.h
index e95f188..8339528 100644
--- a/payload_consumer/verity_writer_android.h
+++ b/payload_consumer/verity_writer_android.h
@@ -64,6 +64,7 @@
   const InstallPlan::Partition* partition_ = nullptr;
 
   std::unique_ptr<HashTreeBuilder> hash_tree_builder_;
+  uint64_t total_offset_ = 0;
   DISALLOW_COPY_AND_ASSIGN(VerityWriterAndroid);
 };
 
diff --git a/payload_consumer/verity_writer_android_unittest.cc b/payload_consumer/verity_writer_android_unittest.cc
index 81d8ca2..75da0ae 100644
--- a/payload_consumer/verity_writer_android_unittest.cc
+++ b/payload_consumer/verity_writer_android_unittest.cc
@@ -76,6 +76,15 @@
   ASSERT_TRUE(verity_writer_.Update(8192, part_data.data(), part_data.size()));
 }
 
+TEST_F(VerityWriterAndroidTest, DiscontinuedRead) {
+  partition_.hash_tree_data_size = 8192;
+  partition_.hash_tree_size = 4096;
+  brillo::Blob part_data(4096);
+  ASSERT_TRUE(verity_writer_.Init(partition_));
+  ASSERT_TRUE(verity_writer_.Update(0, part_data.data(), part_data.size()));
+  ASSERT_FALSE(verity_writer_.Update(8192, part_data.data(), part_data.size()));
+}
+
 TEST_F(VerityWriterAndroidTest, InvalidHashAlgorithmTest) {
   partition_.hash_tree_algorithm = "sha123";
   ASSERT_FALSE(verity_writer_.Init(partition_));
@@ -106,6 +115,32 @@
   ASSERT_EQ(part_data, actual_part);
 }
 
+TEST_F(VerityWriterAndroidTest, NonZeroOffsetSHA256Test) {
+  partition_.hash_tree_algorithm = "sha256";
+  partition_.hash_tree_data_offset = 100;
+  partition_.hash_tree_offset =
+      partition_.hash_tree_data_offset + partition_.hash_tree_data_size;
+  brillo::Blob part_data(8192 + partition_.hash_tree_data_offset);
+  test_utils::WriteFileVector(partition_.target_path, part_data);
+  ASSERT_TRUE(verity_writer_.Init(partition_));
+  ASSERT_TRUE(verity_writer_.Update(0, part_data.data(), 4096));
+  ASSERT_TRUE(verity_writer_.Update(4096, part_data.data() + 4096, 4096));
+  ASSERT_TRUE(verity_writer_.Update(
+      8192, part_data.data() + 8192, partition_.hash_tree_data_offset));
+  ASSERT_TRUE(verity_writer_.Finalize(partition_fd_, partition_fd_));
+  brillo::Blob actual_part;
+  utils::ReadFile(partition_.target_path, &actual_part);
+  // dd if=/dev/zero bs=4096 count=1 2>/dev/null | sha256sum | xxd -r -p |
+  //     hexdump -v -e '/1 "0x%02x, "'
+  brillo::Blob hash = {0xad, 0x7f, 0xac, 0xb2, 0x58, 0x6f, 0xc6, 0xe9,
+                       0x66, 0xc0, 0x04, 0xd7, 0xd1, 0xd1, 0x6b, 0x02,
+                       0x4f, 0x58, 0x05, 0xff, 0x7c, 0xb4, 0x7c, 0x7a,
+                       0x85, 0xda, 0xbd, 0x8b, 0x48, 0x89, 0x2c, 0xa7};
+  memcpy(
+      part_data.data() + partition_.hash_tree_offset, hash.data(), hash.size());
+  ASSERT_EQ(part_data, actual_part);
+}
+
 TEST_F(VerityWriterAndroidTest, FECTest) {
   partition_.fec_data_offset = 0;
   partition_.fec_data_size = 4096;