Add op count check before attempting to write operations

Test: th
Bug: 313962438
Change-Id: I0e288a42984d737d327236693a6b69c03a7ecc6e
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
index de097f5..a95b442 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
@@ -216,7 +216,23 @@
     return true;
 }
 
+bool CowWriterV3::CheckOpCount(size_t op_count) {
+    if (IsEstimating()) {
+        return true;
+    }
+    if (header_.op_count + op_count > header_.op_count_max) {
+        LOG(ERROR) << "Current number of ops on disk: " << header_.op_count
+                   << ", number of ops attempting to write: " << op_count
+                   << ", this will exceed max op count " << header_.op_count_max;
+        return false;
+    }
+    return true;
+}
+
 bool CowWriterV3::EmitCopy(uint64_t new_block, uint64_t old_block, uint64_t num_blocks) {
+    if (!CheckOpCount(num_blocks)) {
+        return false;
+    }
     std::vector<CowOperationV3> ops(num_blocks);
     for (size_t i = 0; i < num_blocks; i++) {
         CowOperationV3& op = ops[i];
@@ -232,11 +248,17 @@
 }
 
 bool CowWriterV3::EmitRawBlocks(uint64_t new_block_start, const void* data, size_t size) {
+    if (!CheckOpCount(size / header_.block_size)) {
+        return false;
+    }
     return EmitBlocks(new_block_start, data, size, 0, 0, kCowReplaceOp);
 }
 
 bool CowWriterV3::EmitXorBlocks(uint32_t new_block_start, const void* data, size_t size,
                                 uint32_t old_block, uint16_t offset) {
+    if (!CheckOpCount(size / header_.block_size)) {
+        return false;
+    }
     return EmitBlocks(new_block_start, data, size, old_block, offset, kCowXorOp);
 }
 
@@ -313,8 +335,11 @@
 }
 
 bool CowWriterV3::EmitZeroBlocks(uint64_t new_block_start, const uint64_t num_blocks) {
+    if (!CheckOpCount(num_blocks)) {
+        return false;
+    }
     std::vector<CowOperationV3> ops(num_blocks);
-    for (uint64_t i = 0; i < ops.size(); i++) {
+    for (uint64_t i = 0; i < num_blocks; i++) {
         auto& op = ops[i];
         op.set_type(kCowZeroOp);
         op.new_block = new_block_start + i;
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.h b/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.h
index 93f1d24..613bbe5 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.h
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.h
@@ -53,6 +53,7 @@
     bool EmitBlocks(uint64_t new_block_start, const void* data, size_t size, uint64_t old_block,
                     uint16_t offset, CowOperationType type);
     bool CompressBlocks(size_t num_blocks, const void* data);
+    bool CheckOpCount(size_t op_count);
 
   private:
     CowHeaderV3 header_{};