libsnapshot: Add Initialize and InitializeAppend methods to ISnapshotWriter.

This is so update engine can resume from the correct label.

Bug: 168554689
Test: vts_libsnapshot_test
Change-Id: Ib04e80e8219f954f105d5a85f86efa7bb9097579
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 4bbdca3..5b50ca9 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -185,6 +185,9 @@
     // must be suffixed. If a source partition exists, it must be specified as well. The source
     // partition will only be used if raw bytes are needed. The source partition should be an
     // absolute path to the device, not a partition name.
+    //
+    // After calling OpenSnapshotWriter, the caller must invoke Initialize or InitializeForAppend
+    // before invoking write operations.
     virtual std::unique_ptr<ISnapshotWriter> OpenSnapshotWriter(
             const android::fs_mgr::CreateLogicalPartitionParams& params,
             const std::optional<std::string>& source_device) = 0;
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
index 45c7ec4..4732c2d 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
@@ -14,6 +14,8 @@
 
 #pragma once
 
+#include <optional>
+
 #include <android-base/unique_fd.h>
 
 #include <libsnapshot/cow_writer.h>
@@ -37,6 +39,13 @@
     // device is only opened on the first operation that requires it.
     void SetSourceDevice(const std::string& source_device);
 
+    // Open the writer in write mode (no append).
+    virtual bool Initialize() = 0;
+
+    // Open the writer in append mode, optionally with the last label to resume
+    // from. See CowWriter::InitializeAppend.
+    virtual bool InitializeAppend(std::optional<uint64_t> label = {}) = 0;
+
     virtual std::unique_ptr<FileDescriptor> OpenReader() = 0;
 
   protected:
@@ -56,6 +65,8 @@
     // Sets the COW device; this is required.
     bool SetCowDevice(android::base::unique_fd&& cow_device);
 
+    bool Initialize() override;
+    bool InitializeAppend(std::optional<uint64_t> label = {}) override;
     bool Finalize() override;
     uint64_t GetCowSize() override;
     std::unique_ptr<FileDescriptor> OpenReader() override;
@@ -80,6 +91,9 @@
     // Set the device used for all writes.
     void SetSnapshotDevice(android::base::unique_fd&& snapshot_fd, uint64_t cow_size);
 
+    bool Initialize() override { return true; }
+    bool InitializeAppend(std::optional<uint64_t>) override { return true; }
+
     bool Finalize() override;
     uint64_t GetCowSize() override { return cow_size_; }
     std::unique_ptr<FileDescriptor> OpenReader() override;
diff --git a/fs_mgr/libsnapshot/snapshot_reader_test.cpp b/fs_mgr/libsnapshot/snapshot_reader_test.cpp
index f54eefa..4202d22 100644
--- a/fs_mgr/libsnapshot/snapshot_reader_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_reader_test.cpp
@@ -157,6 +157,7 @@
     auto writer = std::make_unique<CompressedSnapshotWriter>(options);
     writer->SetSourceDevice(base_->path);
     ASSERT_TRUE(writer->SetCowDevice(std::move(cow_fd)));
+    ASSERT_TRUE(writer->Initialize());
     ASSERT_NO_FATAL_FAILURE(WriteCow(writer.get()));
     ASSERT_NO_FATAL_FAILURE(TestReads(writer.get()));
 }
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 7327629..ec92dcd 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -907,6 +907,9 @@
         if (!result) {
             return AssertionFailure() << "Cannot open snapshot for writing: " << name;
         }
+        if (!result->Initialize()) {
+            return AssertionFailure() << "Cannot initialize snapshot for writing: " << name;
+        }
 
         if (writer) {
             *writer = std::move(result);
diff --git a/fs_mgr/libsnapshot/snapshot_writer.cpp b/fs_mgr/libsnapshot/snapshot_writer.cpp
index d6ecbfd..85ed156 100644
--- a/fs_mgr/libsnapshot/snapshot_writer.cpp
+++ b/fs_mgr/libsnapshot/snapshot_writer.cpp
@@ -56,9 +56,9 @@
 bool CompressedSnapshotWriter::SetCowDevice(android::base::unique_fd&& cow_device) {
     cow_device_ = std::move(cow_device);
     cow_ = std::make_unique<CowWriter>(options_);
-
-    return cow_->Initialize(cow_device_);
+    return true;
 }
+
 bool CompressedSnapshotWriter::Finalize() {
     return cow_->Finalize();
 }
@@ -112,6 +112,17 @@
     return cow_->AddLabel(label);
 }
 
+bool CompressedSnapshotWriter::Initialize() {
+    return cow_->Initialize(cow_device_, CowWriter::OpenMode::WRITE);
+}
+
+bool CompressedSnapshotWriter::InitializeAppend(std::optional<uint64_t> label) {
+    if (label) {
+        return cow_->InitializeAppend(cow_device_, *label);
+    }
+    return cow_->Initialize(cow_device_, CowWriter::OpenMode::APPEND);
+}
+
 OnlineKernelSnapshotWriter::OnlineKernelSnapshotWriter(const CowOptions& options)
     : ISnapshotWriter(options) {}