From 35ff136a256e23dc1362e8836e32617fb6fd5096 Mon Sep 17 00:00:00 2001
From: Daniel Rosenberg <drosen@google.com>
Date: Mon, 12 Apr 2021 21:27:44 -0700
Subject: [PATCH] libsnapshot: Check cluster size on appends

When appending, if the cluster should end after the given label, ensure
that it does.

Bug: 183985866
Test: cow_api_test#ResumeEndCluster
Change-Id: Ie93d09b3431755d0b9b92761619d55df7f9f6151
---
 fs_mgr/libsnapshot/cow_api_test.cpp           | 61 +++++++++++++++++++
 fs_mgr/libsnapshot/cow_writer.cpp             | 17 +++---
 .../include/libsnapshot/cow_writer.h          |  1 +
 3 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/fs_mgr/libsnapshot/cow_api_test.cpp b/fs_mgr/libsnapshot/cow_api_test.cpp
index d06427e4b..b75b154bf 100644
--- a/fs_mgr/libsnapshot/cow_api_test.cpp
+++ b/fs_mgr/libsnapshot/cow_api_test.cpp
@@ -869,6 +869,67 @@ TEST_F(CowTest, ResumeMidCluster) {
     ASSERT_EQ(num_clusters, 2);
 }
 
+TEST_F(CowTest, ResumeEndCluster) {
+    CowOptions options;
+    int cluster_ops = 5;
+    options.cluster_ops = cluster_ops;
+    auto writer = std::make_unique<CowWriter>(options);
+    ASSERT_TRUE(writer->Initialize(cow_->fd));
+
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 1, "Block 1"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 2, "Block 2"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 3, "Block 3"));
+    ASSERT_TRUE(writer->AddLabel(1));
+    ASSERT_TRUE(writer->Finalize());
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 4, "Block 4"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 5, "Block 5"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 6, "Block 6"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 7, "Block 7"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 8, "Block 8"));
+    ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
+
+    writer = std::make_unique<CowWriter>(options);
+    ASSERT_TRUE(writer->InitializeAppend(cow_->fd, 1));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 4, "Block 4"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 5, "Block 5"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 6, "Block 6"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 7, "Block 7"));
+    ASSERT_TRUE(WriteDataBlock(writer.get(), 8, "Block 8"));
+    ASSERT_TRUE(writer->AddLabel(2));
+    ASSERT_TRUE(writer->Finalize());
+    ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
+
+    CowReader reader;
+    ASSERT_TRUE(reader.Parse(cow_->fd));
+
+    auto iter = reader.GetOpIter();
+    size_t num_replace = 0;
+    size_t max_in_cluster = 0;
+    size_t num_in_cluster = 0;
+    size_t num_clusters = 0;
+    while (!iter->Done()) {
+        const auto& op = iter->Get();
+
+        num_in_cluster++;
+        max_in_cluster = std::max(max_in_cluster, num_in_cluster);
+
+        if (op.type == kCowReplaceOp) {
+            num_replace++;
+
+            ASSERT_EQ(op.new_block, num_replace);
+            ASSERT_TRUE(CompareDataBlock(&reader, op, "Block " + std::to_string(num_replace)));
+        } else if (op.type == kCowClusterOp) {
+            num_in_cluster = 0;
+            num_clusters++;
+        }
+
+        iter->Next();
+    }
+    ASSERT_EQ(num_replace, 8);
+    ASSERT_EQ(max_in_cluster, cluster_ops);
+    ASSERT_EQ(num_clusters, 3);
+}
+
 TEST_F(CowTest, DeleteMidCluster) {
     CowOptions options;
     options.cluster_ops = 7;
diff --git a/fs_mgr/libsnapshot/cow_writer.cpp b/fs_mgr/libsnapshot/cow_writer.cpp
index 8fd1cc4a3..645ae9d0d 100644
--- a/fs_mgr/libsnapshot/cow_writer.cpp
+++ b/fs_mgr/libsnapshot/cow_writer.cpp
@@ -236,7 +236,7 @@ bool CowWriter::OpenForAppend(uint64_t label) {
         PLOG(ERROR) << "lseek failed";
         return false;
     }
-    return true;
+    return EmitClusterIfNeeded();
 }
 
 bool CowWriter::EmitCopy(uint64_t new_block, uint64_t old_block) {
@@ -315,6 +315,14 @@ bool CowWriter::EmitCluster() {
     return WriteOperation(op);
 }
 
+bool CowWriter::EmitClusterIfNeeded() {
+    // If there isn't room for another op and the cluster end op, end the current cluster
+    if (cluster_size_ && cluster_size_ < current_cluster_size_ + 2 * sizeof(CowOperation)) {
+        if (!EmitCluster()) return false;
+    }
+    return true;
+}
+
 std::basic_string<uint8_t> CowWriter::Compress(const void* data, size_t length) {
     switch (compression_) {
         case kCowCompressGz: {
@@ -467,12 +475,7 @@ bool CowWriter::WriteOperation(const CowOperation& op, const void* data, size_t
         if (!WriteRawData(data, size)) return false;
     }
     AddOperation(op);
-    // If there isn't room for another op and the cluster end op, end the current cluster
-    if (cluster_size_ && op.type != kCowClusterOp &&
-        cluster_size_ < current_cluster_size_ + 2 * sizeof(op)) {
-        if (!EmitCluster()) return false;
-    }
-    return true;
+    return EmitClusterIfNeeded();
 }
 
 void CowWriter::AddOperation(const CowOperation& op) {
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
index 6ffd5d8f3..a9efad8a8 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
@@ -115,6 +115,7 @@ class CowWriter : public ICowWriter {
 
   private:
     bool EmitCluster();
+    bool EmitClusterIfNeeded();
     void SetupHeaders();
     bool ParseOptions();
     bool OpenForWrite();