Implement VABCPartitionWriter

Test: treehugger

Change-Id: I7ec47adba0bc095879989b43ec642ea6a7852ab6
diff --git a/Android.bp b/Android.bp
index 52019b6..a8fdf51 100644
--- a/Android.bp
+++ b/Android.bp
@@ -82,6 +82,28 @@
     },
 }
 
+// libcow_operation_convert (type: library)
+// ========================================================
+cc_library {
+    name: "libcow_operation_convert",
+    host_supported: true,
+    recovery_available: true,
+    defaults: [
+        "ue_defaults",
+        "update_metadata-protos_exports",
+    ],
+    srcs: [
+        "common/cow_operation_convert.cc",
+    ],
+    static_libs: [
+        "libsnapshot_cow",
+        "update_metadata-protos",
+        "libpayload_extent_ranges",
+        "libbrotli",
+        "libz",
+    ],
+}
+
 // update_metadata-protos (type: static_library)
 // ========================================================
 // Protobufs.
@@ -125,6 +147,10 @@
         "libpuffpatch",
         "libverity_tree",
         "libsnapshot_cow",
+        "libbrotli",
+        "libz",
+        "libpayload_extent_ranges",
+        "libcow_operation_convert",
     ],
     shared_libs: [
         "libbase",
@@ -182,6 +208,7 @@
         "payload_consumer/partition_writer.cc",
         "payload_consumer/partition_writer_factory_android.cc",
         "payload_consumer/vabc_partition_writer.cc",
+        "payload_consumer/snapshot_extent_writer.cc",
         "payload_consumer/postinstall_runner_action.cc",
         "payload_consumer/verity_writer_android.cc",
         "payload_consumer/xz_extent_writer.cc",
@@ -482,6 +509,7 @@
         "ue_defaults",
     ],
     host_supported: true,
+    recovery_available: true,
     srcs: [
         "payload_generator/extent_ranges.cc",
     ],
diff --git a/common/cow_operation_convert.cc b/common/cow_operation_convert.cc
new file mode 100644
index 0000000..a4eaba3
--- /dev/null
+++ b/common/cow_operation_convert.cc
@@ -0,0 +1,30 @@
+//
+// Copyright (C) 2020 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 "update_engine/common/cow_operation_convert.h"
+
+#include "update_engine/payload_generator/extent_ranges.h"
+
+namespace chromeos_update_engine {
+std::vector<CowOperation> ConvertToCowOperations(
+    const ::google::protobuf::RepeatedPtrField<
+        ::chromeos_update_engine::InstallOperation>& operations,
+    const ::google::protobuf::RepeatedPtrField<CowMergeOperation>&
+        merge_operations) {
+  // TODO(zhangkelvin) Implement this.
+  return {};
+}
+}  // namespace chromeos_update_engine
diff --git a/common/cow_operation_convert.h b/common/cow_operation_convert.h
new file mode 100644
index 0000000..bca10ac
--- /dev/null
+++ b/common/cow_operation_convert.h
@@ -0,0 +1,55 @@
+//
+// Copyright (C) 2020 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 __COW_OPERATION_CONVERT_H
+#define __COW_OPERATION_CONVERT_H
+
+#include <vector>
+
+#include <libsnapshot/cow_format.h>
+
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+struct CowOperation {
+  enum Type {
+    CowCopy = android::snapshot::kCowCopyOp,
+    CowReplace = android::snapshot::kCowReplaceOp,
+  };
+  Type op;
+  uint64_t src_block;
+  uint64_t dst_block;
+};
+
+// Convert SOURCE_COPY operations in `operations` list to a list of
+// CowOperations according to the merge sequence. This function only converts
+// SOURCE_COPY, other operations are ignored. If there's a merge conflict in
+// SOURCE_COPY operations, some blocks may be converted to COW_REPLACE instead
+// of COW_COPY.
+
+// The list returned does not necessarily preserve the order of
+// SOURCE_COPY in `operations`. The only guarantee about ordering in the
+// returned list is that if operations are applied in such order, there would be
+// no merge conflicts.
+
+// This funnction is intended to be used by delta_performer to perform
+// SOURCE_COPY operations on Virtual AB Compression devices.
+std::vector<CowOperation> ConvertToCowOperations(
+    const ::google::protobuf::RepeatedPtrField<
+        ::chromeos_update_engine::InstallOperation>& operations,
+    const ::google::protobuf::RepeatedPtrField<CowMergeOperation>&
+        merge_operations);
+}  // namespace chromeos_update_engine
+#endif
diff --git a/payload_consumer/snapshot_extent_writer.cc b/payload_consumer/snapshot_extent_writer.cc
new file mode 100644
index 0000000..882d1f7
--- /dev/null
+++ b/payload_consumer/snapshot_extent_writer.cc
@@ -0,0 +1,54 @@
+//
+// Copyright (C) 2020 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 "update_engine/payload_consumer/snapshot_extent_writer.h"
+
+#include <algorithm>
+#include <cstdint>
+
+#include <libsnapshot/cow_writer.h>
+
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+SnapshotExtentWriter::SnapshotExtentWriter(
+    android::snapshot::ICowWriter* cow_writer)
+    : cow_writer_(cow_writer) {
+  CHECK_NE(cow_writer, nullptr);
+}
+
+SnapshotExtentWriter::~SnapshotExtentWriter() {
+  CHECK(buffer_.empty());
+}
+
+bool SnapshotExtentWriter::Init(
+    FileDescriptorPtr /*fd*/,
+    const google::protobuf::RepeatedPtrField<Extent>& extents,
+    uint32_t /*block_size*/) {
+  // TODO(zhangkelvin) Implement this
+  return true;
+}
+
+// Returns true on success.
+// This will construct a COW_REPLACE operation and forward it to CowWriter. It
+// is important that caller does not perform SOURCE_COPY operation on this
+// class, otherwise raw data will be stored. Caller should find ways to use
+// COW_COPY whenever possible.
+bool SnapshotExtentWriter::Write(const void* bytes, size_t count) {
+  // TODO(zhangkelvin) Implement this
+  return true;
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_consumer/snapshot_extent_writer.h b/payload_consumer/snapshot_extent_writer.h
new file mode 100644
index 0000000..43a8317
--- /dev/null
+++ b/payload_consumer/snapshot_extent_writer.h
@@ -0,0 +1,47 @@
+//
+// Copyright (C) 2020 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 <cstdint>
+#include <vector>
+
+#include <libsnapshot/cow_writer.h>
+
+#include "update_engine/payload_consumer/extent_writer.h"
+
+namespace chromeos_update_engine {
+class SnapshotExtentWriter : public chromeos_update_engine::ExtentWriter {
+ public:
+  explicit SnapshotExtentWriter(android::snapshot::ICowWriter* cow_writer);
+  ~SnapshotExtentWriter();
+  // Returns true on success.
+  bool Init(FileDescriptorPtr fd,
+            const google::protobuf::RepeatedPtrField<Extent>& extents,
+            uint32_t block_size) override;
+  // Returns true on success.
+  // This will construct a COW_REPLACE operation and forward it to CowWriter. It
+  // is important that caller does not perform SOURCE_COPY operation on this
+  // class, otherwise raw data will be stored. Caller should find ways to use
+  // COW_COPY whenever possible.
+  bool Write(const void* bytes, size_t count) override;
+
+ private:
+  // It's a non-owning pointer, because PartitionWriter owns the CowWruter. This
+  // allows us to use a single instance of CowWriter for all operations applied
+  // to the same partition.
+  [[maybe_unused]] android::snapshot::ICowWriter* cow_writer_;
+  [[maybe_unused]] google::protobuf::RepeatedPtrField<Extent> extents_;
+  [[maybe_unused]] std::vector<uint8_t> buffer_;
+};
+}  // namespace chromeos_update_engine
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
diff --git a/payload_consumer/vabc_partition_writer.h b/payload_consumer/vabc_partition_writer.h
index 034fb57..d65ac4a 100644
--- a/payload_consumer/vabc_partition_writer.h
+++ b/payload_consumer/vabc_partition_writer.h
@@ -19,7 +19,7 @@
 
 #include <memory>
 
-#include <libsnapshot/cow_writer.h>
+#include <libsnapshot/snapshot_writer.h>
 
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/partition_writer.h"
@@ -41,9 +41,10 @@
       const InstallOperation& operation) override;
   [[nodiscard]] bool PerformSourceCopyOperation(
       const InstallOperation& operation, ErrorCode* error) override;
+  [[nodiscard]] bool Flush() override;
 
  private:
-  std::unique_ptr<android::snapshot::ICowWriter> cow_writer_;
+  std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer_;
 };
 
 }  // namespace chromeos_update_engine