Merge "libsnapshot: Add a maximum block count to CowWriter." am: 2be4a85a2e

Original change: https://android-review.googlesource.com/c/platform/system/core/+/1432900

Change-Id: Ib9195dad176e4e1d9768c62df486fe6b97bed964

fs_mgr/libsnapshot/cow_writer.cpp

@@ -32,6 +32,45 @@ namespace snapshot {
 
 static_assert(sizeof(off_t) == sizeof(uint64_t));
 
+bool ICowWriter::AddCopy(uint64_t new_block, uint64_t old_block) {
+    if (!ValidateNewBlock(new_block)) {
+        return false;
+    }
+    return EmitCopy(new_block, old_block);
+}
+
+bool ICowWriter::AddRawBlocks(uint64_t new_block_start, const void* data, size_t size) {
+    if (size % options_.block_size != 0) {
+        LOG(ERROR) << "AddRawBlocks: size " << size << " is not a multiple of "
+                   << options_.block_size;
+        return false;
+    }
+
+    uint64_t num_blocks = size / options_.block_size;
+    uint64_t last_block = new_block_start + num_blocks - 1;
+    if (!ValidateNewBlock(last_block)) {
+        return false;
+    }
+    return EmitRawBlocks(new_block_start, data, size);
+}
+
+bool ICowWriter::AddZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) {
+    uint64_t last_block = new_block_start + num_blocks - 1;
+    if (!ValidateNewBlock(last_block)) {
+        return false;
+    }
+    return EmitZeroBlocks(new_block_start, num_blocks);
+}
+
+bool ICowWriter::ValidateNewBlock(uint64_t new_block) {
+    if (options_.max_blocks && new_block >= options_.max_blocks.value()) {
+        LOG(ERROR) << "New block " << new_block << " exceeds maximum block count "
+                   << options_.max_blocks.value();
+        return false;
+    }
+    return true;
+}
+
 CowWriter::CowWriter(const CowOptions& options) : ICowWriter(options), fd_(-1) {
     SetupHeaders();
 }
@@ -134,7 +173,7 @@ bool CowWriter::OpenForAppend() {
     return true;
 }
 
-bool CowWriter::AddCopy(uint64_t new_block, uint64_t old_block) {
+bool CowWriter::EmitCopy(uint64_t new_block, uint64_t old_block) {
     CowOperation op = {};
     op.type = kCowCopyOp;
     op.new_block = new_block;
@@ -143,13 +182,7 @@ bool CowWriter::AddCopy(uint64_t new_block, uint64_t old_block) {
     return true;
 }
 
-bool CowWriter::AddRawBlocks(uint64_t new_block_start, const void* data, size_t size) {
-    if (size % header_.block_size != 0) {
-        LOG(ERROR) << "AddRawBlocks: size " << size << " is not a multiple of "
-                   << header_.block_size;
-        return false;
-    }
-
+bool CowWriter::EmitRawBlocks(uint64_t new_block_start, const void* data, size_t size) {
     uint64_t pos;
     if (!GetDataPos(&pos)) {
         return false;
@@ -195,7 +228,7 @@ bool CowWriter::AddRawBlocks(uint64_t new_block_start, const void* data, size_t
     return true;
 }
 
-bool CowWriter::AddZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) {
+bool CowWriter::EmitZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) {
     for (uint64_t i = 0; i < num_blocks; i++) {
         CowOperation op = {};
         op.type = kCowZeroOp;
@@ -291,7 +324,7 @@ bool CowWriter::Flush() {
     return true;
 }
 
-size_t CowWriter::GetCowSize() {
+uint64_t CowWriter::GetCowSize() {
     return header_.ops_offset + header_.num_ops * sizeof(CowOperation);
 }
 

fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h

@@ -16,6 +16,7 @@
 
 #include <stdint.h>
 
+#include <optional>
 #include <string>
 
 #include <android-base/unique_fd.h>
@@ -27,6 +28,9 @@ namespace snapshot {
 struct CowOptions {
     uint32_t block_size = 4096;
     std::string compression;
+
+    // Maximum number of blocks that can be written.
+    std::optional<uint64_t> max_blocks;
 };
 
 // Interface for writing to a snapuserd COW. All operations are ordered; merges
@@ -39,20 +43,29 @@ class ICowWriter {
 
     // Encode an operation that copies the contents of |old_block| to the
     // location of |new_block|.
-    virtual bool AddCopy(uint64_t new_block, uint64_t old_block) = 0;
+    bool AddCopy(uint64_t new_block, uint64_t old_block);
 
     // Encode a sequence of raw blocks. |size| must be a multiple of the block size.
-    virtual bool AddRawBlocks(uint64_t new_block_start, const void* data, size_t size) = 0;
+    bool AddRawBlocks(uint64_t new_block_start, const void* data, size_t size);
 
     // Encode a sequence of zeroed blocks. |size| must be a multiple of the block size.
-    virtual bool AddZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) = 0;
+    bool AddZeroBlocks(uint64_t new_block_start, uint64_t num_blocks);
 
     // Flush all pending writes. This must be called before closing the writer
     // to ensure that the correct headers and footers are written.
     virtual bool Flush() = 0;
 
     // Return number of bytes the cow image occupies on disk.
-    virtual size_t GetCowSize() = 0;
+    virtual uint64_t GetCowSize() = 0;
+
+    const CowOptions& options() { return options_; }
+
+  protected:
+    virtual bool EmitCopy(uint64_t new_block, uint64_t old_block) = 0;
+    virtual bool EmitRawBlocks(uint64_t new_block_start, const void* data, size_t size) = 0;
+    virtual bool EmitZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) = 0;
+
+    bool ValidateNewBlock(uint64_t new_block);
 
   protected:
     CowOptions options_;
@@ -68,13 +81,14 @@ class CowWriter : public ICowWriter {
     bool Initialize(android::base::unique_fd&& fd, OpenMode mode = OpenMode::WRITE);
     bool Initialize(android::base::borrowed_fd fd, OpenMode mode = OpenMode::WRITE);
 
-    bool AddCopy(uint64_t new_block, uint64_t old_block) override;
-    bool AddRawBlocks(uint64_t new_block_start, const void* data, size_t size) override;
-    bool AddZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) override;
-
     bool Flush() override;
 
-    size_t GetCowSize() override;
+    uint64_t GetCowSize() override;
+
+  protected:
+    virtual bool EmitCopy(uint64_t new_block, uint64_t old_block) override;
+    virtual bool EmitRawBlocks(uint64_t new_block_start, const void* data, size_t size) override;
+    virtual bool EmitZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) override;
 
   private:
     void SetupHeaders();