libsnapshot: Add helpers for accessing CowOperation offsets.

These helpers avoid manual inspection of CowOperation::source, and
fiddling with block sizes.

Bug: 280529365
Test: cow_api_test
      vts_libsnapshot_test
      apply OTA
Change-Id: I1ba276e155eb8232a3115066caaa11030c171e11
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
index f599501..d0369e5 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
@@ -73,8 +73,20 @@
     // The operation pointer must derive from ICowOpIter::Get().
     virtual ssize_t ReadData(const CowOperation* op, void* buffer, size_t buffer_size,
                              size_t ignore_bytes = 0) = 0;
+
+    // Get the absolute source offset, in bytes, of a CowOperation. Returns
+    // false if the operation does not read from source partitions.
+    virtual bool GetSourceOffset(const CowOperation* op, uint64_t* source_offset) = 0;
 };
 
+static constexpr uint64_t GetBlockFromOffset(const CowHeader& header, uint64_t offset) {
+    return offset / header.block_size;
+}
+
+static constexpr uint64_t GetBlockRelativeOffset(const CowHeader& header, uint64_t offset) {
+    return offset % header.block_size;
+}
+
 // Iterate over a sequence of COW operations. The iterator is bidirectional.
 class ICowOpIter {
   public:
@@ -119,6 +131,7 @@
     bool VerifyMergeOps() override;
     bool GetFooter(CowFooter* footer) override;
     bool GetLastLabel(uint64_t* label) override;
+    bool GetSourceOffset(const CowOperation* op, uint64_t* source_offset) override;
 
     // Create a CowOpIter object which contains footer_.num_ops
     // CowOperation objects. Get() returns a unique CowOperation object
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
index 4ac89ae..cffc8de 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
@@ -312,19 +312,15 @@
     std::unordered_map<uint64_t, const CowOperation*> overwritten_blocks;
     while (!itr->AtEnd()) {
         const auto& op = itr->Get();
-        uint64_t block;
-        bool offset;
-        if (op->type == kCowCopyOp) {
-            block = op->source;
-            offset = false;
-        } else if (op->type == kCowXorOp) {
-            block = op->source / header_.block_size;
-            offset = (op->source % header_.block_size) != 0;
-        } else {
+        uint64_t offset;
+        if (!GetSourceOffset(op, &offset)) {
             itr->Next();
             continue;
         }
 
+        uint64_t block = GetBlockFromOffset(header_, offset);
+        bool misaligned = (GetBlockRelativeOffset(header_, offset) != 0);
+
         const CowOperation* overwrite = nullptr;
         if (overwritten_blocks.count(block)) {
             overwrite = overwritten_blocks[block];
@@ -332,7 +328,7 @@
                        << op << "\noverwritten by previously merged op:\n"
                        << *overwrite;
         }
-        if (offset && overwritten_blocks.count(block + 1)) {
+        if (misaligned && overwritten_blocks.count(block + 1)) {
             overwrite = overwritten_blocks[block + 1];
             LOG(ERROR) << "Invalid Sequence! Block needed for op:\n"
                        << op << "\noverwritten by previously merged op:\n"
@@ -622,5 +618,18 @@
     return decompressor->Decompress(buffer, buffer_size, header_.block_size, ignore_bytes);
 }
 
+bool CowReader::GetSourceOffset(const CowOperation* op, uint64_t* source_offset) {
+    switch (op->type) {
+        case kCowCopyOp:
+            *source_offset = op->source * header_.block_size;
+            return true;
+        case kCowXorOp:
+            *source_offset = op->source;
+            return true;
+        default:
+            return false;
+    }
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/snapshot_reader.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/snapshot_reader.cpp
index f9cdbc0..a3e40d9 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/snapshot_reader.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/snapshot_reader.cpp
@@ -155,7 +155,12 @@
         }
 
         if (op) {
-            chunk = op->source;
+            uint64_t source_offset;
+            if (!cow_->GetSourceOffset(op, &source_offset)) {
+                LOG(ERROR) << "GetSourceOffset failed in CompressedSnapshotReader for op: " << *op;
+                return false;
+            }
+            chunk = GetBlockFromOffset(cow_->GetHeader(), source_offset);
         }
 
         off64_t offset = (chunk * block_size_) + start_offset;
@@ -179,7 +184,12 @@
             return -1;
         }
 
-        off64_t offset = op->source + start_offset;
+        uint64_t source_offset;
+        if (!cow_->GetSourceOffset(op, &source_offset)) {
+            LOG(ERROR) << "GetSourceOffset failed in CompressedSnapshotReader for op: " << *op;
+            return false;
+        }
+        off64_t offset = source_offset + start_offset;
 
         std::string data(bytes_to_read, '\0');
         if (!android::base::ReadFullyAtOffset(fd, data.data(), data.size(), offset)) {
diff --git a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_worker.cpp b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_worker.cpp
index 922df34..559dcc7 100644
--- a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_worker.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_worker.cpp
@@ -115,12 +115,13 @@
         SNAP_LOG(ERROR) << "ReadFromBaseDevice: Failed to get payload buffer";
         return false;
     }
-    SNAP_LOG(DEBUG) << " ReadFromBaseDevice...: new-block: " << cow_op->new_block
-                    << " Source: " << cow_op->source;
-    uint64_t offset = cow_op->source;
-    if (cow_op->type == kCowCopyOp) {
-        offset *= BLOCK_SZ;
+    uint64_t offset;
+    if (!reader_->GetSourceOffset(cow_op, &offset)) {
+        SNAP_LOG(ERROR) << "ReadFromBaseDevice: Failed to get source offset";
+        return false;
     }
+    SNAP_LOG(DEBUG) << " ReadFromBaseDevice...: new-block: " << cow_op->new_block
+                    << " Source: " << offset;
     if (!android::base::ReadFullyAtOffset(backing_store_fd_, buffer, BLOCK_SZ, offset)) {
         SNAP_PLOG(ERROR) << "Copy op failed. Read from backing store: " << backing_store_device_
                          << "at block :" << offset / BLOCK_SZ << " offset:" << offset % BLOCK_SZ;
@@ -508,7 +509,7 @@
                 if (read_ahead_buffer_map.find(cow_op->new_block) == read_ahead_buffer_map.end()) {
                     SNAP_LOG(ERROR)
                             << " Block: " << cow_op->new_block << " not found in read-ahead cache"
-                            << " Source: " << cow_op->source;
+                            << " Op: " << *cow_op;
                     return -1;
                 }
                 // If this is a final block merged in the read-ahead buffer
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
index 7858216..44b7319 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
@@ -94,12 +94,13 @@
         SNAP_LOG(ERROR) << "ReadFromBaseDevice: Failed to get payload buffer";
         return false;
     }
-    SNAP_LOG(DEBUG) << " ReadFromBaseDevice...: new-block: " << cow_op->new_block
-                    << " Source: " << cow_op->source;
-    uint64_t offset = cow_op->source;
-    if (cow_op->type == kCowCopyOp) {
-        offset *= BLOCK_SZ;
+    uint64_t offset;
+    if (!reader_->GetSourceOffset(cow_op, &offset)) {
+        SNAP_LOG(ERROR) << "ReadFromSourceDevice: Failed to get source offset";
+        return false;
     }
+    SNAP_LOG(DEBUG) << " ReadFromBaseDevice...: new-block: " << cow_op->new_block
+                    << " Op: " << *cow_op;
     if (!android::base::ReadFullyAtOffset(backing_store_fd_, buffer, BLOCK_SZ, offset)) {
         std::string op;
         if (cow_op->type == kCowCopyOp)
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp
index af24286..399f7b8 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp
@@ -34,14 +34,17 @@
 }
 
 void ReadAhead::CheckOverlap(const CowOperation* cow_op) {
-    uint64_t source_block = cow_op->source;
-    uint64_t source_offset = 0;
-    if (cow_op->type == kCowXorOp) {
-        source_block /= BLOCK_SZ;
-        source_offset = cow_op->source % BLOCK_SZ;
+    uint64_t source_offset;
+    if (!reader_->GetSourceOffset(cow_op, &source_offset)) {
+        SNAP_LOG(ERROR) << "ReadAhead operation has no source offset: " << *cow_op;
+        return;
     }
+
+    uint64_t source_block = GetBlockFromOffset(header_, source_offset);
+    bool misaligned = (GetBlockRelativeOffset(header_, source_offset) != 0);
+
     if (dest_blocks_.count(cow_op->new_block) || source_blocks_.count(source_block) ||
-        (source_offset > 0 && source_blocks_.count(source_block + 1))) {
+        (misaligned && source_blocks_.count(source_block + 1))) {
         overlap_ = true;
     }
 
@@ -66,11 +69,12 @@
 
     // Get the first block with offset
     const CowOperation* cow_op = GetRAOpIter();
-    *source_offset = cow_op->source;
 
-    if (cow_op->type == kCowCopyOp) {
-        *source_offset *= BLOCK_SZ;
-    } else if (cow_op->type == kCowXorOp) {
+    if (!reader_->GetSourceOffset(cow_op, source_offset)) {
+        SNAP_LOG(ERROR) << "PrepareNextReadAhead operation has no source offset: " << *cow_op;
+        return nr_consecutive;
+    }
+    if (cow_op->type == kCowXorOp) {
         xor_op_vec.push_back(cow_op);
     }
 
@@ -88,10 +92,10 @@
      */
     while (!RAIterDone() && num_ops) {
         const CowOperation* op = GetRAOpIter();
-        uint64_t next_offset = op->source;
-
-        if (cow_op->type == kCowCopyOp) {
-            next_offset *= BLOCK_SZ;
+        uint64_t next_offset;
+        if (!reader_->GetSourceOffset(op, &next_offset)) {
+            SNAP_LOG(ERROR) << "PrepareNextReadAhead operation has no source offset: " << *cow_op;
+            break;
         }
 
         // Check for consecutive blocks
@@ -803,6 +807,7 @@
     if (!reader_->InitForMerge(std::move(cow_fd_))) {
         return false;
     }
+    header_ = reader_->GetHeader();
     return true;
 }
 
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.h b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.h
index 5e94de0..d3ba126 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.h
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.h
@@ -85,6 +85,7 @@
 
     std::shared_ptr<SnapshotHandler> snapuserd_;
     std::unique_ptr<CowReader> reader_;
+    CowHeader header_;
 
     std::unordered_set<uint64_t> dest_blocks_;
     std::unordered_set<uint64_t> source_blocks_;