Merge "libsnapshot: Implement OpenSnapshotWriter for compressed snapshots."
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index b373d99..b239f31 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) {}