Add GetLastLabel and InitializeAppend

GetLastLabel returns the last Label that a reader is confident about.
InitializeAppend starts a writer up to append data after the last given
label, assuming all later labels are not relevant data.

Change-Id: I3339d5527bae833d9293cbbc63126136b94bd976
Bug: 168829493
Test: cow_api_test
diff --git a/fs_mgr/libsnapshot/cow_api_test.cpp b/fs_mgr/libsnapshot/cow_api_test.cpp
index ffee52f..3d0321f 100644
--- a/fs_mgr/libsnapshot/cow_api_test.cpp
+++ b/fs_mgr/libsnapshot/cow_api_test.cpp
@@ -300,7 +300,9 @@
 
     // Read back both operations.
     CowReader reader;
+    uint64_t label;
     ASSERT_TRUE(reader.Parse(cow_->fd));
+    ASSERT_FALSE(reader.GetLastLabel(&label));
 
     StringSink sink;
 
@@ -432,6 +434,15 @@
 
     ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
 
+    // Get the last known good label
+    CowReader label_reader;
+    uint64_t label;
+    ASSERT_TRUE(label_reader.Parse(cow_->fd));
+    ASSERT_TRUE(label_reader.GetLastLabel(&label));
+    ASSERT_EQ(label, 5);
+
+    ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
+
     writer = std::make_unique<CowWriter>(options);
     ASSERT_TRUE(writer->Initialize(cow_->fd, CowWriter::OpenMode::APPEND));
 
@@ -459,6 +470,78 @@
     ASSERT_TRUE(iter->Done());
 }
 
+TEST_F(CowTest, AppendbyLabel) {
+    CowOptions options;
+    auto writer = std::make_unique<CowWriter>(options);
+    ASSERT_TRUE(writer->Initialize(cow_->fd));
+
+    ASSERT_TRUE(writer->AddLabel(4));
+
+    ASSERT_TRUE(writer->AddLabel(5));
+    std::string data = "This is some data, believe it";
+    data.resize(options.block_size * 2, '\0');
+    ASSERT_TRUE(writer->AddRawBlocks(50, data.data(), data.size()));
+
+    ASSERT_TRUE(writer->AddLabel(6));
+    ASSERT_TRUE(writer->AddZeroBlocks(50, 2));
+
+    ASSERT_EQ(lseek(cow_->fd, 0, SEEK_SET), 0);
+
+    writer = std::make_unique<CowWriter>(options);
+    ASSERT_FALSE(writer->InitializeAppend(cow_->fd, 12));
+    ASSERT_TRUE(writer->InitializeAppend(cow_->fd, 5));
+
+    // This should drop label 6
+    ASSERT_TRUE(writer->Finalize());
+
+    struct stat buf;
+    ASSERT_EQ(fstat(cow_->fd, &buf), 0);
+    ASSERT_EQ(buf.st_size, writer->GetCowSize());
+
+    // Read back all ops
+    CowReader reader;
+    ASSERT_TRUE(reader.Parse(cow_->fd));
+
+    StringSink sink;
+
+    auto iter = reader.GetOpIter();
+    ASSERT_NE(iter, nullptr);
+
+    ASSERT_FALSE(iter->Done());
+    auto op = &iter->Get();
+    ASSERT_EQ(op->type, kCowLabelOp);
+    ASSERT_EQ(op->source, 4);
+
+    iter->Next();
+
+    ASSERT_FALSE(iter->Done());
+    op = &iter->Get();
+    ASSERT_EQ(op->type, kCowLabelOp);
+    ASSERT_EQ(op->source, 5);
+
+    iter->Next();
+
+    ASSERT_FALSE(iter->Done());
+    op = &iter->Get();
+    ASSERT_EQ(op->type, kCowReplaceOp);
+    ASSERT_TRUE(reader.ReadData(*op, &sink));
+    ASSERT_EQ(sink.stream(), data.substr(0, options.block_size));
+
+    iter->Next();
+    sink.Reset();
+
+    ASSERT_FALSE(iter->Done());
+    op = &iter->Get();
+    ASSERT_EQ(op->type, kCowReplaceOp);
+    ASSERT_TRUE(reader.ReadData(*op, &sink));
+    ASSERT_EQ(sink.stream(), data.substr(options.block_size, 2 * options.block_size));
+
+    iter->Next();
+    sink.Reset();
+
+    ASSERT_TRUE(iter->Done());
+}
+
 }  // namespace snapshot
 }  // namespace android
 
diff --git a/fs_mgr/libsnapshot/cow_reader.cpp b/fs_mgr/libsnapshot/cow_reader.cpp
index 6d6a22d..b1667e3 100644
--- a/fs_mgr/libsnapshot/cow_reader.cpp
+++ b/fs_mgr/libsnapshot/cow_reader.cpp
@@ -30,7 +30,14 @@
 namespace android {
 namespace snapshot {
 
-CowReader::CowReader() : fd_(-1), header_(), footer_(), fd_size_(0), has_footer_(false) {}
+CowReader::CowReader()
+    : fd_(-1),
+      header_(),
+      footer_(),
+      fd_size_(0),
+      has_footer_(false),
+      last_label_(0),
+      has_last_label_(false) {}
 
 static void SHA256(const void*, size_t, uint8_t[]) {
 #if 0
@@ -101,6 +108,65 @@
         return false;
     }
     has_footer_ = (footer_.op.type == kCowFooterOp);
+    return ParseOps();
+}
+
+bool CowReader::ParseOps() {
+    uint64_t pos = lseek(fd_.get(), sizeof(header_), SEEK_SET);
+    if (pos != sizeof(header_)) {
+        PLOG(ERROR) << "lseek ops failed";
+        return false;
+    }
+    uint64_t next_last_label = 0;
+    bool has_next = false;
+    auto ops_buffer = std::make_shared<std::vector<CowOperation>>();
+    if (has_footer_) ops_buffer->reserve(footer_.op.num_ops);
+    uint64_t current_op_num = 0;
+    // Look until we reach the last possible non-footer position.
+    uint64_t last_pos = fd_size_ - (has_footer_ ? sizeof(footer_) : sizeof(CowOperation));
+
+    // Alternating op and data
+    while (pos < last_pos) {
+        ops_buffer->resize(current_op_num + 1);
+        if (!android::base::ReadFully(fd_, ops_buffer->data() + current_op_num,
+                                      sizeof(CowOperation))) {
+            PLOG(ERROR) << "read op failed";
+            return false;
+        }
+        auto& current_op = ops_buffer->data()[current_op_num];
+        pos = lseek(fd_.get(), GetNextOpOffset(current_op), SEEK_CUR);
+        if (pos < 0) {
+            PLOG(ERROR) << "lseek next op failed";
+            return false;
+        }
+        current_op_num++;
+        if (current_op.type == kCowLabelOp) {
+            // If we don't have a footer, the last label may be incomplete
+            if (has_footer_) {
+                has_last_label_ = true;
+                last_label_ = current_op.source;
+            } else {
+                last_label_ = next_last_label;
+                if (has_next) has_last_label_ = true;
+                next_last_label = current_op.source;
+                has_next = true;
+            }
+        }
+    }
+
+    uint8_t csum[32];
+    memset(csum, 0, sizeof(uint8_t) * 32);
+
+    if (has_footer_) {
+        SHA256(ops_buffer.get()->data(), footer_.op.ops_size, csum);
+        if (memcmp(csum, footer_.data.ops_checksum, sizeof(csum)) != 0) {
+            LOG(ERROR) << "ops checksum does not match";
+            return false;
+        }
+    } else {
+        LOG(INFO) << "No Footer, recovered data";
+    }
+    ops_ = ops_buffer;
     return true;
 }
 
@@ -115,21 +181,27 @@
     return true;
 }
 
+bool CowReader::GetLastLabel(uint64_t* label) {
+    if (!has_last_label_) return false;
+    *label = last_label_;
+    return true;
+}
+
 class CowOpIter final : public ICowOpIter {
   public:
-    CowOpIter(std::unique_ptr<std::vector<CowOperation>>&& ops);
+    CowOpIter(std::shared_ptr<std::vector<CowOperation>>& ops);
 
     bool Done() override;
     const CowOperation& Get() override;
     void Next() override;
 
   private:
-    std::unique_ptr<std::vector<CowOperation>> ops_;
+    std::shared_ptr<std::vector<CowOperation>> ops_;
     std::vector<CowOperation>::iterator op_iter_;
 };
 
-CowOpIter::CowOpIter(std::unique_ptr<std::vector<CowOperation>>&& ops) {
-    ops_ = std::move(ops);
+CowOpIter::CowOpIter(std::shared_ptr<std::vector<CowOperation>>& ops) {
+    ops_ = ops;
     op_iter_ = ops_.get()->begin();
 }
 
@@ -148,43 +220,7 @@
 }
 
 std::unique_ptr<ICowOpIter> CowReader::GetOpIter() {
-    uint64_t pos = lseek(fd_.get(), sizeof(header_), SEEK_SET);
-    if (pos != sizeof(header_)) {
-        PLOG(ERROR) << "lseek ops failed";
-        return nullptr;
-    }
-    auto ops_buffer = std::make_unique<std::vector<CowOperation>>();
-    if (has_footer_) ops_buffer->reserve(footer_.op.num_ops);
-    uint64_t current_op_num = 0;
-    uint64_t last_pos = fd_size_ - (has_footer_ ? sizeof(footer_) : sizeof(CowOperation));
-
-    // Alternating op and data
-    while (pos < last_pos) {
-        ops_buffer->resize(current_op_num + 1);
-        if (!android::base::ReadFully(fd_, ops_buffer->data() + current_op_num,
-                                      sizeof(CowOperation))) {
-            PLOG(ERROR) << "read op failed";
-            return nullptr;
-        }
-        auto current_op = ops_buffer->data()[current_op_num];
-        pos = lseek(fd_.get(), GetNextOpOffset(current_op), SEEK_CUR);
-        current_op_num++;
-    }
-
-    uint8_t csum[32];
-    memset(csum, 0, sizeof(uint8_t) * 32);
-
-    if (has_footer_) {
-        SHA256(ops_buffer.get()->data(), footer_.op.ops_size, csum);
-        if (memcmp(csum, footer_.data.ops_checksum, sizeof(csum)) != 0) {
-            LOG(ERROR) << "ops checksum does not match";
-            return nullptr;
-        }
-    } else {
-        LOG(INFO) << "No Footer, recovered data";
-    }
-
-    return std::make_unique<CowOpIter>(std::move(ops_buffer));
+    return std::make_unique<CowOpIter>(ops_);
 }
 
 bool CowReader::GetRawBytes(uint64_t offset, void* buffer, size_t len, size_t* read) {
diff --git a/fs_mgr/libsnapshot/cow_writer.cpp b/fs_mgr/libsnapshot/cow_writer.cpp
index 679b55e..ec2dc96 100644
--- a/fs_mgr/libsnapshot/cow_writer.cpp
+++ b/fs_mgr/libsnapshot/cow_writer.cpp
@@ -133,6 +133,21 @@
     }
 }
 
+bool CowWriter::InitializeAppend(android::base::unique_fd&& fd, uint64_t label) {
+    owned_fd_ = std::move(fd);
+    return InitializeAppend(android::base::borrowed_fd{owned_fd_}, label);
+}
+
+bool CowWriter::InitializeAppend(android::base::borrowed_fd fd, uint64_t label) {
+    fd_ = fd;
+
+    if (!ParseOptions()) {
+        return false;
+    }
+
+    return OpenForAppend(label);
+}
+
 bool CowWriter::OpenForWrite() {
     // This limitation is tied to the data field size in CowOperation.
     if (header_.block_size > std::numeric_limits<uint16_t>::max()) {
@@ -205,6 +220,52 @@
     return true;
 }
 
+bool CowWriter::OpenForAppend(uint64_t label) {
+    auto reader = std::make_unique<CowReader>();
+    std::queue<CowOperation> toAdd;
+    if (!reader->Parse(fd_) || !reader->GetHeader(&header_)) {
+        return false;
+    }
+
+    options_.block_size = header_.block_size;
+    bool found_label = false;
+
+    // Reset this, since we're going to reimport all operations.
+    footer_.op.num_ops = 0;
+    next_op_pos_ = sizeof(header_);
+
+    auto iter = reader->GetOpIter();
+    while (!iter->Done()) {
+        CowOperation op = iter->Get();
+        if (op.type == kCowFooterOp) break;
+        if (op.type == kCowLabelOp) {
+            if (found_label) break;
+            if (op.source == label) found_label = true;
+        }
+        AddOperation(op);
+        iter->Next();
+    }
+
+    if (!found_label) {
+        PLOG(ERROR) << "Failed to find last label";
+        return false;
+    }
+
+    // Free reader so we own the descriptor position again.
+    reader = nullptr;
+
+    // Position for new writing
+    if (ftruncate(fd_.get(), next_op_pos_) != 0) {
+        PLOG(ERROR) << "Failed to trim file";
+        return false;
+    }
+    if (lseek(fd_.get(), 0, SEEK_END) < 0) {
+        PLOG(ERROR) << "lseek failed";
+        return false;
+    }
+    return true;
+}
+
 bool CowWriter::EmitCopy(uint64_t new_block, uint64_t old_block) {
     CowOperation op = {};
     op.type = kCowCopyOp;
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
index a4360aa..814e104 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
@@ -61,6 +61,9 @@
     // Return the file footer.
     virtual bool GetFooter(CowFooter* footer) = 0;
 
+    // Return the last valid label
+    virtual bool GetLastLabel(uint64_t* label) = 0;
+
     // Return an iterator for retrieving CowOperation entries.
     virtual std::unique_ptr<ICowOpIter> GetOpIter() = 0;
 
@@ -94,21 +97,28 @@
     bool GetHeader(CowHeader* header) override;
     bool GetFooter(CowFooter* footer) override;
 
-    // Create a CowOpIter object which contains header_.num_ops
+    bool GetLastLabel(uint64_t* label) override;
+
+    // Create a CowOpIter object which contains footer_.num_ops
     // CowOperation objects. Get() returns a unique CowOperation object
-    // whose lifeteime depends on the CowOpIter object
+    // whose lifetime depends on the CowOpIter object
     std::unique_ptr<ICowOpIter> GetOpIter() override;
     bool ReadData(const CowOperation& op, IByteSink* sink) override;
 
     bool GetRawBytes(uint64_t offset, void* buffer, size_t len, size_t* read);
 
   private:
+    bool ParseOps();
+
     android::base::unique_fd owned_fd_;
     android::base::borrowed_fd fd_;
     CowHeader header_;
     CowFooter footer_;
     uint64_t fd_size_;
     bool has_footer_;
+    uint64_t last_label_;
+    bool has_last_label_;
+    std::shared_ptr<std::vector<CowOperation>> ops_;
 };
 
 }  // namespace snapshot
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
index 4cfbaaa..c031d63 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
@@ -16,11 +16,13 @@
 
 #include <stdint.h>
 
+#include <memory>
 #include <optional>
 #include <string>
 
 #include <android-base/unique_fd.h>
 #include <libsnapshot/cow_format.h>
+#include <libsnapshot/cow_reader.h>
 
 namespace android {
 namespace snapshot {
@@ -85,8 +87,16 @@
     explicit CowWriter(const CowOptions& options);
 
     // Set up the writer.
+    // If opening for write, the file starts from the beginning.
+    // If opening for append, if the file has a footer, we start appending to the last op.
+    // If the footer isn't found, the last label is considered corrupt, and dropped.
     bool Initialize(android::base::unique_fd&& fd, OpenMode mode = OpenMode::WRITE);
     bool Initialize(android::base::borrowed_fd fd, OpenMode mode = OpenMode::WRITE);
+    // Set up a writer, assuming that the given label is the last valid label.
+    // This will result in dropping any labels that occur after the given on, and will fail
+    // if the given label does not appear.
+    bool InitializeAppend(android::base::unique_fd&&, uint64_t label);
+    bool InitializeAppend(android::base::borrowed_fd fd, uint64_t label);
 
     bool Finalize() override;
 
@@ -103,6 +113,8 @@
     bool ParseOptions();
     bool OpenForWrite();
     bool OpenForAppend();
+    bool OpenForAppend(uint64_t label);
+    bool ImportOps(std::unique_ptr<ICowOpIter> iter);
     bool GetDataPos(uint64_t* pos);
     bool WriteRawData(const void* data, size_t size);
     bool WriteOperation(const CowOperation& op, const void* data = nullptr, size_t size = 0);