libsnapshot: Implement OpenSnapshotWriter for compressed snapshots.

This does not implement OpenReader() yet, however, it implements enough
for some vts_libsnapshot_tests to start passing with Virtual A/B
Compression enabled.

Note that OpenSnapshotWriter() has been disabled when linking with init
or recovery. This is to avoid pulling in all the compression libraries
in those places. OpenSnapshotWriter is only designed to be called by
update_engine and will not work outside of normal operation.

Bug: 168554689
Test: vts_libsnapshot_test
Test: full OTA with update_device.py
Test: incremental OTA with update_device.py
Change-Id: I9737d28bdd5c5f4914bc30a2bb72f357d5f44d2b
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 7985fcd..49d4563 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -28,6 +28,7 @@
         "liblog",
     ],
     static_libs: [
+        "libbrotli",
         "libdm",
         "libfstab",
         "libsnapshot_cow",
@@ -109,6 +110,9 @@
     defaults: ["libsnapshot_defaults"],
     srcs: [":libsnapshot_sources"],
     recovery_available: true,
+    cflags: [
+        "-DLIBSNAPSHOT_NO_COW_WRITE",
+    ],
     static_libs: [
         "libfs_mgr",
     ],
@@ -122,6 +126,9 @@
     ],
     srcs: [":libsnapshot_sources"],
     recovery_available: true,
+    cflags: [
+        "-DLIBSNAPSHOT_NO_COW_WRITE",
+    ],
     static_libs: [
         "libfs_mgr",
     ],
@@ -244,6 +251,7 @@
     static_libs: [
         "android.hardware.boot@1.0",
         "android.hardware.boot@1.1",
+        "libbrotli",
         "libfs_mgr",
         "libgsi",
         "libgmock",
@@ -276,9 +284,11 @@
         "snapshotctl.cpp",
     ],
     static_libs: [
+        "libbrotli",
         "libfstab",
         "libsnapshot",
         "libsnapshot_cow",
+        "libz",
         "update_metadata-protos",
     ],
     shared_libs: [
@@ -332,6 +342,7 @@
     ],
     static_libs: [
         "libbase",
+        "libbrotli",
         "libcrypto_static",
         "libcutils",
         "libext2_uuid",
@@ -345,6 +356,7 @@
         "libsnapshot_cow",
         "libsnapshot_test_helpers",
         "libprotobuf-mutator",
+        "libz",
     ],
     header_libs: [
         "libchrome",
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 1bc972e..403e350 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -532,8 +532,8 @@
         // Target/base device (eg system_b), always present.
         std::string target_device;
 
-        // COW path (eg system_cow). Not present if no COW is needed.
-        std::string cow_device;
+        // COW name (eg system_cow). Not present if no COW is needed.
+        std::string cow_device_name;
 
         // dm-snapshot instance. Not present in Update mode for VABC.
         std::string snapshot_device;
@@ -626,6 +626,9 @@
     bool GetMappedImageDeviceStringOrPath(const std::string& device_name,
                                           std::string* device_string_or_mapped_path);
 
+    // Same as above, but for paths only (no major:minor device strings).
+    bool GetMappedImageDevicePath(const std::string& device_name, std::string* device_path);
+
     std::string gsid_dir_;
     std::string metadata_dir_;
     std::unique_ptr<IDeviceInfo> device_;
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
index bf57a00..da058f9 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
@@ -42,6 +42,29 @@
     android::base::unique_fd source_fd_;
 };
 
+// Send writes to a COW or a raw device directly, based on a threshold.
+class CompressedSnapshotWriter : public ISnapshotWriter {
+  public:
+    CompressedSnapshotWriter(const CowOptions& options);
+
+    // Sets the COW device, if needed.
+    bool SetCowDevice(android::base::unique_fd&& cow_device);
+
+    bool Flush() override;
+    uint64_t GetCowSize() override;
+    std::unique_ptr<FileDescriptor> OpenReader() override;
+
+  protected:
+    bool EmitCopy(uint64_t new_block, uint64_t old_block) override;
+    bool EmitRawBlocks(uint64_t new_block_start, const void* data, size_t size) override;
+    bool EmitZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) override;
+
+  private:
+    android::base::unique_fd cow_device_;
+
+    std::unique_ptr<CowWriter> cow_;
+};
+
 // Write directly to a dm-snapshot device.
 class OnlineKernelSnapshotWriter : public ISnapshotWriter {
   public:
@@ -52,7 +75,7 @@
 
     bool Flush() override;
     uint64_t GetCowSize() override { return cow_size_; }
-    virtual std::unique_ptr<FileDescriptor> OpenReader() override;
+    std::unique_ptr<FileDescriptor> OpenReader() override;
 
   protected:
     bool EmitRawBlocks(uint64_t new_block_start, const void* data, size_t size) override;
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 8112b5f..3541ef6 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -271,7 +271,7 @@
         return false;
     }
 
-    if (!EnsureNoOverflowSnapshot(lock.get())) {
+    if (!IsCompressionEnabled() && !EnsureNoOverflowSnapshot(lock.get())) {
         LOG(ERROR) << "Cannot ensure there are no overflow snapshots.";
         return false;
     }
@@ -1716,7 +1716,7 @@
         return false;
     }
     if (paths) {
-        paths->cow_device = cow_device;
+        paths->cow_device_name = cow_name;
     }
 
     remaining_time = GetRemainingTime(params.timeout_ms, begin);
@@ -1724,6 +1724,7 @@
 
     if (context == SnapshotContext::Update && IsCompressionEnabled()) {
         // Stop here, we can't run dm-user yet, the COW isn't built.
+        created_devices.Release();
         return true;
     }
 
@@ -2471,6 +2472,12 @@
 
 std::unique_ptr<ISnapshotWriter> SnapshotManager::OpenSnapshotWriter(
         const android::fs_mgr::CreateLogicalPartitionParams& params) {
+#if defined(LIBSNAPSHOT_NO_COW_WRITE)
+    (void)params;
+
+    LOG(ERROR) << "Snapshots cannot be written in first-stage init or recovery";
+    return nullptr;
+#else
     // First unmap any existing mapping.
     auto lock = LockShared();
     if (!lock) return nullptr;
@@ -2486,7 +2493,7 @@
     }
 
     SnapshotStatus status;
-    if (!paths.cow_device.empty()) {
+    if (!paths.cow_device_name.empty()) {
         if (!ReadSnapshotStatus(lock.get(), params.GetPartitionName(), &status)) {
             return nullptr;
         }
@@ -2504,12 +2511,49 @@
         return OpenCompressedSnapshotWriter(lock.get(), params.GetPartitionName(), status, paths);
     }
     return OpenKernelSnapshotWriter(lock.get(), params.GetPartitionName(), status, paths);
+#endif
 }
 
+#if !defined(LIBSNAPSHOT_NO_COW_WRITE)
 std::unique_ptr<ISnapshotWriter> SnapshotManager::OpenCompressedSnapshotWriter(
-        LockedFile*, const std::string&, const SnapshotStatus&, const SnapshotPaths&) {
-    LOG(ERROR) << "OpenSnapshotWriter not yet implemented for compression";
-    return nullptr;
+        LockedFile* lock, [[maybe_unused]] const std::string& partition_name,
+        const SnapshotStatus& status, const SnapshotPaths& paths) {
+    CHECK(lock);
+
+    CowOptions cow_options;
+    cow_options.compression = "gz";
+    cow_options.max_blocks = {status.device_size() / cow_options.block_size};
+
+    // Currently we don't support partial snapshots, since partition_cow_creator
+    // never creates this scenario.
+    CHECK(status.snapshot_size() == status.device_size());
+
+    auto writer = std::make_unique<CompressedSnapshotWriter>(cow_options);
+
+    unique_fd base_fd(open(paths.target_device.c_str(), O_RDWR | O_CLOEXEC));
+    if (base_fd < 0) {
+        PLOG(ERROR) << "OpenCompressedSnapshotWriter: open " << paths.target_device;
+        return nullptr;
+    }
+    writer->SetSourceDevice(std::move(base_fd));
+
+    std::string cow_path;
+    if (!GetMappedImageDevicePath(paths.cow_device_name, &cow_path)) {
+        LOG(ERROR) << "Could not determine path for " << paths.cow_device_name;
+        return nullptr;
+    }
+
+    unique_fd cow_fd(open(cow_path.c_str(), O_RDWR | O_CLOEXEC));
+    if (cow_fd < 0) {
+        PLOG(ERROR) << "OpenCompressedSnapshotWriter: open " << cow_path;
+        return nullptr;
+    }
+    if (!writer->SetCowDevice(std::move(cow_fd))) {
+        LOG(ERROR) << "Could not create COW writer from " << cow_path;
+        return nullptr;
+    }
+
+    return writer;
 }
 
 std::unique_ptr<ISnapshotWriter> SnapshotManager::OpenKernelSnapshotWriter(
@@ -2534,6 +2578,7 @@
 
     return writer;
 }
+#endif  // !defined(LIBSNAPSHOT_NO_COW_WRITE)
 
 bool SnapshotManager::UnmapUpdateSnapshot(const std::string& target_partition_name) {
     auto lock = LockShared();
@@ -2872,6 +2917,22 @@
     return SnapshotMergeStats::GetInstance(*this);
 }
 
+// This is only to be used in recovery or normal Android (not first-stage init).
+// We don't guarantee dm paths are available in first-stage init, because ueventd
+// isn't running yet.
+bool SnapshotManager::GetMappedImageDevicePath(const std::string& device_name,
+                                               std::string* device_path) {
+    auto& dm = DeviceMapper::Instance();
+
+    // Try getting the device string if it is a device mapper device.
+    if (dm.GetState(device_name) != DmDeviceState::INVALID) {
+        return dm.GetDmDevicePathByName(device_name, device_path);
+    }
+
+    // Otherwise, get path from IImageManager.
+    return images_->GetMappedImageDevice(device_name, device_path);
+}
+
 bool SnapshotManager::GetMappedImageDeviceStringOrPath(const std::string& device_name,
                                                        std::string* device_string_or_mapped_path) {
     auto& dm = DeviceMapper::Instance();
diff --git a/fs_mgr/libsnapshot/snapshot_writer.cpp b/fs_mgr/libsnapshot/snapshot_writer.cpp
index 584f15e..aa49ab1 100644
--- a/fs_mgr/libsnapshot/snapshot_writer.cpp
+++ b/fs_mgr/libsnapshot/snapshot_writer.cpp
@@ -33,6 +33,40 @@
     source_fd_ = std::move(source_fd);
 }
 
+CompressedSnapshotWriter::CompressedSnapshotWriter(const CowOptions& options)
+    : ISnapshotWriter(options) {}
+
+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_);
+}
+bool CompressedSnapshotWriter::Flush() {
+    return cow_->Flush();
+}
+
+uint64_t CompressedSnapshotWriter::GetCowSize() {
+    return cow_->GetCowSize();
+}
+
+std::unique_ptr<FileDescriptor> CompressedSnapshotWriter::OpenReader() {
+    return nullptr;
+}
+
+bool CompressedSnapshotWriter::EmitCopy(uint64_t new_block, uint64_t old_block) {
+    return cow_->AddCopy(new_block, old_block);
+}
+
+bool CompressedSnapshotWriter::EmitRawBlocks(uint64_t new_block_start, const void* data,
+                                             size_t size) {
+    return cow_->AddRawBlocks(new_block_start, data, size);
+}
+
+bool CompressedSnapshotWriter::EmitZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) {
+    return cow_->AddZeroBlocks(new_block_start, num_blocks);
+}
+
 OnlineKernelSnapshotWriter::OnlineKernelSnapshotWriter(const CowOptions& options)
     : ISnapshotWriter(options) {}