Implement VABCPartitionWriter

Test: treehugger

Change-Id: I7ec47adba0bc095879989b43ec642ea6a7852ab6
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index ab4897f..1578f29 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -17,33 +17,74 @@
 #include "update_engine/payload_consumer/vabc_partition_writer.h"
 
 #include <memory>
+#include <vector>
 
 #include <libsnapshot/cow_writer.h>
 
+#include "update_engine/common/cow_operation_convert.h"
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/extent_writer.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/partition_writer.h"
+#include "update_engine/payload_consumer/snapshot_extent_writer.h"
 
 namespace chromeos_update_engine {
 bool VABCPartitionWriter::Init(const InstallPlan* install_plan,
                                bool source_may_exist) {
+  TEST_AND_RETURN_FALSE(install_plan != nullptr);
   TEST_AND_RETURN_FALSE(PartitionWriter::Init(install_plan, source_may_exist));
+  cow_writer_ = dynamic_control_->OpenCowWriter(
+      install_part_.name, install_part_.source_path, install_plan->is_resume);
+  TEST_AND_RETURN_FALSE(cow_writer_ != nullptr);
 
-  // TODO(zhangkelvin) Add code specific to VABC. E.x. Convert InstallOps to
-  // CowOps, perform all SOURCE_COPY upfront according to merge sequence.
+  // TODO(zhangkelvin) Emit a label before writing SOURCE_COPY. When resuming,
+  // use pref or CowWriter::GetLastLabel to determine if the SOURCE_COPY ops are
+  // written. No need to handle SOURCE_COPY operations when resuming.
+
+  // ===== Resume case handling code goes here ====
+
+  // ==============================================
+
+  // TODO(zhangkelvin) Rewrite this in C++20 coroutine once that's available.
+  auto converted = ConvertToCowOperations(partition_update_.operations(),
+                                          partition_update_.merge_operations());
+  std::vector<uint8_t> buffer(block_size_);
+  for (const auto& cow_op : converted) {
+    switch (cow_op.op) {
+      case CowOperation::CowCopy:
+        TEST_AND_RETURN_FALSE(
+            cow_writer_->AddCopy(cow_op.dst_block, cow_op.src_block));
+        break;
+      case CowOperation::CowReplace:
+        ssize_t bytes_read = 0;
+        TEST_AND_RETURN_FALSE(utils::PReadAll(source_fd_,
+                                              buffer.data(),
+                                              block_size_,
+                                              cow_op.src_block * block_size_,
+                                              &bytes_read));
+        if (bytes_read <= 0 || static_cast<size_t>(bytes_read) != block_size_) {
+          LOG(ERROR) << "source_fd->Read failed: " << bytes_read;
+          return false;
+        }
+        TEST_AND_RETURN_FALSE(cow_writer_->AddRawBlocks(
+            cow_op.dst_block, buffer.data(), block_size_));
+        break;
+    }
+  }
   return true;
 }
 
 std::unique_ptr<ExtentWriter> VABCPartitionWriter::CreateBaseExtentWriter() {
-  // TODO(zhangkelvin) Return a SnapshotExtentWriter
-  return std::make_unique<DirectExtentWriter>();
+  return std::make_unique<SnapshotExtentWriter>(cow_writer_.get());
 }
 
 [[nodiscard]] bool VABCPartitionWriter::PerformZeroOrDiscardOperation(
     const InstallOperation& operation) {
-  // TODO(zhangkelvin) Create a COW_ZERO operation and send it to CowWriter
-  return PartitionWriter::PerformZeroOrDiscardOperation(operation);
+  for (const auto& extent : operation.dst_extents()) {
+    TEST_AND_RETURN_FALSE(
+        cow_writer_->AddZeroBlocks(extent.start_block(), extent.num_blocks()));
+  }
+  return true;
 }
 
 [[nodiscard]] bool VABCPartitionWriter::PerformSourceCopyOperation(
@@ -53,6 +94,13 @@
   return true;
 }
 
-VABCPartitionWriter::~VABCPartitionWriter() = default;
+bool VABCPartitionWriter::Flush() {
+  // No need to do anything, as CowWriter automatically flushes every OP added.
+  return true;
+}
+
+VABCPartitionWriter::~VABCPartitionWriter() {
+  cow_writer_->Finalize();
+}
 
 }  // namespace chromeos_update_engine