diff --git a/fs_mgr/Android.bp b/fs_mgr/Android.bp
index dd61272..bbd068b 100644
--- a/fs_mgr/Android.bp
+++ b/fs_mgr/Android.bp
@@ -181,6 +181,10 @@
     ramdisk_available: true,
     vendor_ramdisk_available: true,
     recovery_available: true,
+    apex_available: [
+        "//apex_available:anyapex",
+        "//apex_available:platform",
+    ],
     host_supported: true,
     defaults: ["fs_mgr_defaults"],
     srcs: [
@@ -206,6 +210,7 @@
         "libbase_headers",
         "libgsi_headers",
     ],
+    min_sdk_version: "31",
 }
 
 cc_binary {
diff --git a/fs_mgr/fs_mgr.cpp b/fs_mgr/fs_mgr.cpp
index debcb3e..7f4959f 100644
--- a/fs_mgr/fs_mgr.cpp
+++ b/fs_mgr/fs_mgr.cpp
@@ -1849,18 +1849,14 @@
     return ret;
 }
 
-// If tmp_mount_point is non-null, mount the filesystem there.  This is for the
-// tmp mount we do to check the user password
 // If multiple fstab entries are to be mounted on "n_name", it will try to mount each one
 // in turn, and stop on 1st success, or no more match.
-static int fs_mgr_do_mount_helper(Fstab* fstab, const std::string& n_name,
-                                  const std::string& n_blk_device, const char* tmp_mount_point,
-                                  int needs_checkpoint, bool metadata_encrypted,
-                                  bool needs_encrypt) {
+int fs_mgr_do_mount(Fstab* fstab, const std::string& n_name, const std::string& n_blk_device,
+                    int needs_checkpoint, bool needs_encrypt) {
     int mount_errors = 0;
     int first_mount_errno = 0;
     std::string mount_point;
-    CheckpointManager checkpoint_manager(needs_checkpoint, metadata_encrypted, needs_encrypt);
+    CheckpointManager checkpoint_manager(needs_checkpoint, true, needs_encrypt);
     AvbUniquePtr avb_handle(nullptr);
 
     if (!fstab) {
@@ -1902,11 +1898,7 @@
         }
 
         // Now mount it where requested */
-        if (tmp_mount_point) {
-            mount_point = tmp_mount_point;
-        } else {
-            mount_point = fstab_entry.mount_point;
-        }
+        mount_point = fstab_entry.mount_point;
 
         int fs_stat = prepare_fs_for_mount(n_blk_device, fstab_entry, mount_point);
 
@@ -1963,35 +1955,6 @@
     return FS_MGR_DOMNT_FAILED;
 }
 
-int fs_mgr_do_mount(Fstab* fstab, const char* n_name, char* n_blk_device, char* tmp_mount_point) {
-    return fs_mgr_do_mount_helper(fstab, n_name, n_blk_device, tmp_mount_point, -1, false, false);
-}
-
-int fs_mgr_do_mount(Fstab* fstab, const char* n_name, char* n_blk_device, char* tmp_mount_point,
-                    bool needs_checkpoint, bool metadata_encrypted, bool needs_encrypt) {
-    return fs_mgr_do_mount_helper(fstab, n_name, n_blk_device, tmp_mount_point, needs_checkpoint,
-                                  metadata_encrypted, needs_encrypt);
-}
-
-/*
- * mount a tmpfs filesystem at the given point.
- * return 0 on success, non-zero on failure.
- */
-int fs_mgr_do_tmpfs_mount(const char *n_name)
-{
-    int ret;
-
-    ret = mount("tmpfs", n_name, "tmpfs", MS_NOATIME | MS_NOSUID | MS_NODEV | MS_NOEXEC,
-                CRYPTO_TMPFS_OPTIONS);
-    if (ret < 0) {
-        LERROR << "Cannot mount tmpfs filesystem at " << n_name;
-        return -1;
-    }
-
-    /* Success */
-    return 0;
-}
-
 static bool ConfigureIoScheduler(const std::string& device_path) {
     if (!StartsWith(device_path, "/dev/")) {
         LERROR << __func__ << ": invalid argument " << device_path;
diff --git a/fs_mgr/include/fs_mgr.h b/fs_mgr/include/fs_mgr.h
index 37515d4..bc4a7a6 100644
--- a/fs_mgr/include/fs_mgr.h
+++ b/fs_mgr/include/fs_mgr.h
@@ -85,14 +85,10 @@
 #define FS_MGR_DOMNT_FAILED (-1)
 #define FS_MGR_DOMNT_BUSY (-2)
 #define FS_MGR_DOMNT_SUCCESS 0
-int fs_mgr_do_mount(android::fs_mgr::Fstab* fstab, const char* n_name, char* n_blk_device,
-                    char* tmp_mount_point);
-int fs_mgr_do_mount(android::fs_mgr::Fstab* fstab, const char* n_name, char* n_blk_device,
-                    char* tmp_mount_point, bool need_cp, bool metadata_encrypted,
-                    bool need_encrypted);
+int fs_mgr_do_mount(android::fs_mgr::Fstab* fstab, const std::string& n_name,
+                    const std::string& n_blk_device, int needs_checkpoint, bool needs_encrypt);
 int fs_mgr_do_mount_one(const android::fs_mgr::FstabEntry& entry,
                         const std::string& mount_point = "");
-int fs_mgr_do_tmpfs_mount(const char *n_name);
 bool fs_mgr_load_verity_state(int* mode);
 // Returns true if verity is enabled on this particular FstabEntry.
 bool fs_mgr_is_verity_enabled(const android::fs_mgr::FstabEntry& entry);
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
index dd626bc..3a81f63 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
@@ -166,6 +166,13 @@
 static constexpr uint8_t kCowReadAheadInProgress = 1;
 static constexpr uint8_t kCowReadAheadDone = 2;
 
+static inline uint64_t GetCowOpSourceInfoData(const CowOperation* op) {
+    return op->source;
+}
+static inline bool GetCowOpSourceInfoCompression(const CowOperation* op) {
+    return op->compression != kCowCompressNone;
+}
+
 struct CowFooter {
     CowFooterOperation op;
     uint8_t unused[64];
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
index 3890b17..f4ce52f 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
@@ -133,6 +146,7 @@
 
     CowHeader& GetHeader() override { return header_; }
 
+    bool GetRawBytes(const CowOperation* op, void* buffer, size_t len, size_t* read);
     bool GetRawBytes(uint64_t offset, void* buffer, size_t len, size_t* read);
 
     // Returns the total number of data ops that should be merged. This is the
@@ -154,6 +168,7 @@
     bool ParseOps(std::optional<uint64_t> label);
     bool PrepMergeOps();
     uint64_t FindNumCopyops();
+    uint8_t GetCompressionType(const CowOperation* op);
 
     android::base::unique_fd owned_fd_;
     android::base::borrowed_fd fd_;
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
index c2a7fdb..489669a 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
@@ -230,7 +230,7 @@
             size_t seq_len = current_op.data_length / sizeof(uint32_t);
 
             merge_op_blocks->resize(merge_op_blocks->size() + seq_len);
-            if (!GetRawBytes(current_op.source, &merge_op_blocks->data()[num_seqs],
+            if (!GetRawBytes(&current_op, &merge_op_blocks->data()[num_seqs],
                              current_op.data_length, &read)) {
                 PLOG(ERROR) << "Failed to read sequence op!";
                 return false;
@@ -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"
@@ -516,6 +512,18 @@
                                             ignore_progress ? 0 : merge_op_start_);
 }
 
+bool CowReader::GetRawBytes(const CowOperation* op, void* buffer, size_t len, size_t* read) {
+    switch (op->type) {
+        case kCowSequenceOp:
+        case kCowReplaceOp:
+        case kCowXorOp:
+            return GetRawBytes(GetCowOpSourceInfoData(op), buffer, len, read);
+        default:
+            LOG(ERROR) << "Cannot get raw bytes of non-data op: " << *op;
+            return false;
+    }
+}
+
 bool CowReader::GetRawBytes(uint64_t offset, void* buffer, size_t len, size_t* read) {
     // Validate the offset, taking care to acknowledge possible overflow of offset+len.
     if (offset < header_.prefix.header_size || offset >= fd_size_ - sizeof(CowFooter) ||
@@ -566,10 +574,14 @@
     size_t remaining_;
 };
 
+uint8_t CowReader::GetCompressionType(const CowOperation* op) {
+    return op->compression;
+}
+
 ssize_t CowReader::ReadData(const CowOperation* op, void* buffer, size_t buffer_size,
                             size_t ignore_bytes) {
     std::unique_ptr<IDecompressor> decompressor;
-    switch (op->compression) {
+    switch (GetCompressionType(op)) {
         case kCowCompressNone:
             break;
         case kCowCompressGz:
@@ -589,7 +601,7 @@
             }
             break;
         default:
-            LOG(ERROR) << "Unknown compression type: " << op->compression;
+            LOG(ERROR) << "Unknown compression type: " << GetCompressionType(op);
             return -1;
     }
 
@@ -597,7 +609,7 @@
     if (op->type == kCowXorOp) {
         offset = data_loc_->at(op->new_block);
     } else {
-        offset = op->source;
+        offset = GetCowOpSourceInfoData(op);
     }
 
     if (!decompressor) {
@@ -610,5 +622,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 = GetCowOpSourceInfoData(op) * header_.block_size;
+            return true;
+        case kCowXorOp:
+            *source_offset = GetCowOpSourceInfoData(op);
+            return true;
+        default:
+            return false;
+    }
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/inspect_cow.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/inspect_cow.cpp
index 148ecb0..a6dee4f 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/inspect_cow.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/inspect_cow.cpp
@@ -57,7 +57,7 @@
     size_t count;
     auto buffer = std::make_unique<uint8_t[]>(op->data_length);
 
-    if (!reader.GetRawBytes(op->source, buffer.get(), op->data_length, &count)) {
+    if (!reader.GetRawBytes(op, buffer.get(), op->data_length, &count)) {
         std::cerr << "Failed to read at all!\n";
     } else {
         std::cout << "The Block data is:\n";
@@ -199,7 +199,7 @@
             std::vector<uint32_t> merge_op_blocks;
             size_t seq_len = op->data_length / sizeof(uint32_t);
             merge_op_blocks.resize(seq_len);
-            if (!reader.GetRawBytes(op->source, merge_op_blocks.data(), op->data_length, &read)) {
+            if (!reader.GetRawBytes(op, merge_op_blocks.data(), op->data_length, &read)) {
                 PLOG(ERROR) << "Failed to read sequence op!";
                 return false;
             }
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/libsnapshot_cow/test_v2.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/test_v2.cpp
index 120b2c0..31b9a58 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/test_v2.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/test_v2.cpp
@@ -145,7 +145,7 @@
     op = iter->Get();
 
     ASSERT_EQ(op->type, kCowReplaceOp);
-    ASSERT_EQ(op->compression, kCowCompressNone);
+    ASSERT_FALSE(GetCowOpSourceInfoCompression(op));
     ASSERT_EQ(op->data_length, 4096);
     ASSERT_EQ(op->new_block, 50);
     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
@@ -224,10 +224,10 @@
     op = iter->Get();
 
     ASSERT_EQ(op->type, kCowXorOp);
-    ASSERT_EQ(op->compression, kCowCompressNone);
+    ASSERT_FALSE(GetCowOpSourceInfoCompression(op));
     ASSERT_EQ(op->data_length, 4096);
     ASSERT_EQ(op->new_block, 50);
-    ASSERT_EQ(op->source, 98314);  // 4096 * 24 + 10
+    ASSERT_EQ(GetCowOpSourceInfoData(op), 98314);  // 4096 * 24 + 10
     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
     ASSERT_EQ(sink, data);
 
@@ -283,7 +283,7 @@
     std::string sink(data.size(), '\0');
 
     ASSERT_EQ(op->type, kCowReplaceOp);
-    ASSERT_EQ(op->compression, kCowCompressGz);
+    ASSERT_TRUE(GetCowOpSourceInfoCompression(op));
     ASSERT_EQ(op->data_length, 56);  // compressed!
     ASSERT_EQ(op->new_block, 50);
     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
@@ -339,7 +339,7 @@
             total_blocks += 1;
             std::string sink(xor_data.size(), '\0');
             ASSERT_EQ(op->new_block, 50);
-            ASSERT_EQ(op->source, 98314);  // 4096 * 24 + 10
+            ASSERT_EQ(GetCowOpSourceInfoData(op), 98314);  // 4096 * 24 + 10
             ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
             ASSERT_EQ(sink, xor_data);
         }
@@ -528,7 +528,7 @@
     std::string sink(data.size(), '\0');
 
     ASSERT_EQ(op->type, kCowReplaceOp);
-    ASSERT_EQ(op->compression, kCowCompressGz);
+    ASSERT_TRUE(GetCowOpSourceInfoCompression(op));
     ASSERT_EQ(op->data_length, 56);  // compressed!
     ASSERT_EQ(op->new_block, 50);
     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
@@ -546,7 +546,7 @@
 
     sink = {};
     sink.resize(data2.size(), '\0');
-    ASSERT_EQ(op->compression, kCowCompressGz);
+    ASSERT_TRUE(GetCowOpSourceInfoCompression(op));
     ASSERT_EQ(op->data_length, 41);  // compressed!
     ASSERT_EQ(op->new_block, 51);
     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
@@ -591,7 +591,7 @@
 
     auto op = iter->Get();
     ASSERT_EQ(op->type, kCowReplaceOp);
-    ASSERT_EQ(op->compression, kCowCompressGz);
+    ASSERT_TRUE(GetCowOpSourceInfoCompression(op));
     ASSERT_EQ(op->new_block, 51);
     ASSERT_TRUE(ReadData(reader, op, sink.data(), sink.size()));
 }
diff --git a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd.cpp b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd.cpp
index da9bd11..978a7f2 100644
--- a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd.cpp
@@ -508,7 +508,7 @@
             // the merge of operations are done based on the ops present
             // in the file.
             //===========================================================
-            uint64_t block_source = cow_op->source;
+            uint64_t block_source = GetCowOpSourceInfoData(cow_op);
             if (prev_id.has_value()) {
                 if (dest_blocks.count(cow_op->new_block) || source_blocks.count(block_source)) {
                     break;
diff --git a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_readahead.cpp b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_readahead.cpp
index a32c2bf..d5fbe91 100644
--- a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_readahead.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_readahead.cpp
@@ -172,7 +172,7 @@
 }
 
 void ReadAheadThread::CheckOverlap(const CowOperation* cow_op) {
-    uint64_t source_block = cow_op->source;
+    uint64_t source_block = GetCowOpSourceInfoData(cow_op);
     if (dest_blocks_.count(cow_op->new_block) || source_blocks_.count(source_block)) {
         overlap_ = true;
     }
@@ -191,7 +191,7 @@
         // Get the first block with offset
         const CowOperation* cow_op = GetRAOpIter();
         CHECK_NE(cow_op, nullptr);
-        *source_offset = cow_op->source;
+        *source_offset = GetCowOpSourceInfoData(cow_op);
         if (cow_op->type == kCowCopyOp) {
             *source_offset *= BLOCK_SZ;
         }
@@ -210,7 +210,7 @@
         while (!RAIterDone() && num_ops) {
             const CowOperation* op = GetRAOpIter();
             CHECK_NE(op, nullptr);
-            uint64_t next_offset = op->source;
+            uint64_t next_offset = GetCowOpSourceInfoData(op);
             if (op->type == kCowCopyOp) {
                 next_offset *= BLOCK_SZ;
             }
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_core.h b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
index c16ad24..cf38875 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
@@ -120,7 +120,7 @@
 
     // Functions interacting with dm-user
     bool ReadDmUserHeader();
-    bool WriteDmUserPayload(size_t size, bool header_response);
+    bool WriteDmUserPayload(size_t size);
     bool DmuserReadRequest();
 
     // IO Path
@@ -130,11 +130,11 @@
     bool ReadDataFromBaseDevice(sector_t sector, size_t read_size);
     bool ReadFromSourceDevice(const CowOperation* cow_op);
 
-    bool ReadAlignedSector(sector_t sector, size_t sz, bool header_response);
+    bool ReadAlignedSector(sector_t sector, size_t sz);
     bool ReadUnalignedSector(sector_t sector, size_t size);
     int ReadUnalignedSector(sector_t sector, size_t size,
                             std::vector<std::pair<sector_t, const CowOperation*>>::iterator& it);
-    bool RespondIOError(bool header_response);
+    bool RespondIOError();
 
     // Processing COW operations
     bool ProcessCowOp(const CowOperation* cow_op);
@@ -176,6 +176,7 @@
     unique_fd backing_store_fd_;
     unique_fd base_path_merge_fd_;
     unique_fd ctrl_fd_;
+    bool header_response_ = false;
 
     std::unique_ptr<ICowOpIter> cowop_iter_;
     size_t ra_block_index_ = 0;
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..c505c32 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)
@@ -305,10 +306,10 @@
 }
 
 // Send the payload/data back to dm-user misc device.
-bool Worker::WriteDmUserPayload(size_t size, bool header_response) {
+bool Worker::WriteDmUserPayload(size_t size) {
     size_t payload_size = size;
     void* buf = bufsink_.GetPayloadBufPtr();
-    if (header_response) {
+    if (header_response_) {
         payload_size += sizeof(struct dm_user_header);
         buf = bufsink_.GetBufPtr();
     }
@@ -318,6 +319,9 @@
         return false;
     }
 
+    // After the first header is sent in response to a request, we cannot
+    // send any additional headers.
+    header_response_ = false;
     return true;
 }
 
@@ -340,7 +344,7 @@
     return true;
 }
 
-bool Worker::ReadAlignedSector(sector_t sector, size_t sz, bool header_response) {
+bool Worker::ReadAlignedSector(sector_t sector, size_t sz) {
     struct dm_user_header* header = bufsink_.GetHeaderPtr();
     size_t remaining_size = sz;
     std::vector<std::pair<sector_t, const CowOperation*>>& chunk_vec = snapuserd_->GetChunkVec();
@@ -388,7 +392,7 @@
 
             // Just return the header if it is an error
             if (header->type == DM_USER_RESP_ERROR) {
-                if (!RespondIOError(header_response)) {
+                if (!RespondIOError()) {
                     return false;
                 }
 
@@ -403,14 +407,12 @@
         }
 
         if (!io_error) {
-            if (!WriteDmUserPayload(total_bytes_read, header_response)) {
+            if (!WriteDmUserPayload(total_bytes_read)) {
                 return false;
             }
 
             SNAP_LOG(DEBUG) << "WriteDmUserPayload success total_bytes_read: " << total_bytes_read
-                            << " header-response: " << header_response
                             << " remaining_size: " << remaining_size;
-            header_response = false;
             remaining_size -= total_bytes_read;
         }
     } while (remaining_size > 0 && !io_error);
@@ -483,7 +485,6 @@
     // to any COW ops. In that case, we just need to read from the base
     // device.
     bool merge_complete = false;
-    bool header_response = true;
     if (it == chunk_vec.end()) {
         if (chunk_vec.size() > 0) {
             // I/O request beyond the last mapped sector
@@ -502,7 +503,7 @@
             --it;
         }
     } else {
-        return ReadAlignedSector(sector, size, header_response);
+        return ReadAlignedSector(sector, size);
     }
 
     loff_t requested_offset = sector << SECTOR_SHIFT;
@@ -536,7 +537,7 @@
         if (ret < 0) {
             SNAP_LOG(ERROR) << "ReadUnalignedSector failed for sector: " << sector
                             << " size: " << size << " it->sector: " << it->first;
-            return RespondIOError(header_response);
+            return RespondIOError();
         }
 
         remaining_size -= ret;
@@ -544,14 +545,13 @@
         sector += (ret >> SECTOR_SHIFT);
 
         // Send the data back
-        if (!WriteDmUserPayload(total_bytes_read, header_response)) {
+        if (!WriteDmUserPayload(total_bytes_read)) {
             return false;
         }
 
-        header_response = false;
         // If we still have pending data to be processed, this will be aligned I/O
         if (remaining_size) {
-            return ReadAlignedSector(sector, remaining_size, header_response);
+            return ReadAlignedSector(sector, remaining_size);
         }
     } else {
         // This is all about handling I/O request to be routed to base device
@@ -565,21 +565,21 @@
         CHECK(diff_size <= BLOCK_SZ);
         if (remaining_size < diff_size) {
             if (!ReadDataFromBaseDevice(sector, remaining_size)) {
-                return RespondIOError(header_response);
+                return RespondIOError();
             }
             total_bytes_read += remaining_size;
 
-            if (!WriteDmUserPayload(total_bytes_read, header_response)) {
+            if (!WriteDmUserPayload(total_bytes_read)) {
                 return false;
             }
         } else {
             if (!ReadDataFromBaseDevice(sector, diff_size)) {
-                return RespondIOError(header_response);
+                return RespondIOError();
             }
 
             total_bytes_read += diff_size;
 
-            if (!WriteDmUserPayload(total_bytes_read, header_response)) {
+            if (!WriteDmUserPayload(total_bytes_read)) {
                 return false;
             }
 
@@ -587,17 +587,16 @@
             size_t num_sectors_read = (diff_size >> SECTOR_SHIFT);
             sector += num_sectors_read;
             CHECK(IsBlockAligned(sector << SECTOR_SHIFT));
-            header_response = false;
 
             // If we still have pending data to be processed, this will be aligned I/O
-            return ReadAlignedSector(sector, remaining_size, header_response);
+            return ReadAlignedSector(sector, remaining_size);
         }
     }
 
     return true;
 }
 
-bool Worker::RespondIOError(bool header_response) {
+bool Worker::RespondIOError() {
     struct dm_user_header* header = bufsink_.GetHeaderPtr();
     header->type = DM_USER_RESP_ERROR;
     // This is an issue with the dm-user interface. There
@@ -609,9 +608,9 @@
     // this back to dm-user.
     //
     // TODO: Fix the interface
-    CHECK(header_response);
+    CHECK(header_response_);
 
-    if (!WriteDmUserPayload(0, header_response)) {
+    if (!WriteDmUserPayload(0)) {
         return false;
     }
 
@@ -628,7 +627,7 @@
         return ReadUnalignedSector(header->sector, header->len);
     }
 
-    return ReadAlignedSector(header->sector, header->len, true);
+    return ReadAlignedSector(header->sector, header->len);
 }
 
 bool Worker::ProcessIORequest() {
@@ -644,6 +643,8 @@
     SNAP_LOG(DEBUG) << "Daemon: msg->type: " << std::dec << header->type;
     SNAP_LOG(DEBUG) << "Daemon: msg->flags: " << std::dec << header->flags;
 
+    header_response_ = true;
+
     switch (header->type) {
         case DM_USER_REQ_MAP_READ: {
             if (!DmuserReadRequest()) {
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_;
diff --git a/init/Android.bp b/init/Android.bp
index b4184d0..1af398a 100644
--- a/init/Android.bp
+++ b/init/Android.bp
@@ -214,8 +214,8 @@
     visibility: [":__subpackages__"],
 }
 
-cc_library_static {
-    name: "libinit",
+cc_defaults {
+    name: "libinit_defaults",
     recovery_available: true,
     defaults: [
         "init_defaults",
@@ -251,10 +251,17 @@
             ],
         },
     },
-    visibility: [
-        "//system/apex/apexd",
-        "//frameworks/native/cmds/installd",
-    ],
+}
+
+cc_library_static {
+    name: "libinit",
+    defaults: ["libinit_defaults"],
+}
+
+cc_library_static {
+    name: "libinit.microdroid",
+    defaults: ["libinit_defaults"],
+    cflags: ["-DMICRODROID=1"],
 }
 
 phony {
@@ -269,7 +276,6 @@
     recovery_available: true,
     stem: "init",
     defaults: ["init_defaults"],
-    static_libs: ["libinit"],
     srcs: ["main.cpp"],
     symlinks: ["ueventd"],
     target: {
@@ -308,12 +314,14 @@
 cc_binary {
     name: "init_second_stage",
     defaults: ["init_second_stage_defaults"],
+    static_libs: ["libinit"],
 }
 
 cc_binary {
     name: "init_second_stage.microdroid",
     defaults: ["init_second_stage_defaults"],
-    cflags: ["-DMICRODROID"],
+    static_libs: ["libinit.microdroid"],
+    cflags: ["-DMICRODROID=1"],
     installable: false,
     visibility: ["//packages/modules/Virtualization/microdroid"],
 }
@@ -459,7 +467,7 @@
 cc_binary {
     name: "init_first_stage.microdroid",
     defaults: ["init_first_stage_defaults"],
-    cflags: ["-DMICRODROID"],
+    cflags: ["-DMICRODROID=1"],
     installable: false,
 }
 
diff --git a/init/fuzzer/Android.bp b/init/fuzzer/Android.bp
index c21a196..856ca8c 100644
--- a/init/fuzzer/Android.bp
+++ b/init/fuzzer/Android.bp
@@ -18,7 +18,7 @@
 }
 
 cc_defaults {
-    name: "libinit_defaults",
+    name: "libinit_fuzzer_defaults",
     static_libs: [
         "libc++fs",
         "liblmkd_utils",
@@ -53,7 +53,7 @@
     ],
     shared_libs: ["libhidlmetadata",],
     defaults: [
-        "libinit_defaults",
+        "libinit_fuzzer_defaults",
     ],
 }
 
@@ -62,7 +62,7 @@
     srcs: [
         "init_property_fuzzer.cpp",
     ],
-    defaults: ["libinit_defaults"],
+    defaults: ["libinit_fuzzer_defaults"],
 }
 
 cc_fuzz {
@@ -71,6 +71,6 @@
         "init_ueventHandler_fuzzer.cpp",
     ],
     defaults: [
-        "libinit_defaults",
+        "libinit_fuzzer_defaults",
     ],
 }
diff --git a/init/selinux.cpp b/init/selinux.cpp
index a936532..e0ef491 100644
--- a/init/selinux.cpp
+++ b/init/selinux.cpp
@@ -300,8 +300,6 @@
 }
 
 constexpr const char plat_policy_cil_file[] = "/system/etc/selinux/plat_sepolicy.cil";
-constexpr const char kMicrodroidPrecompiledSepolicy[] =
-        "/system/etc/selinux/microdroid_precompiled_sepolicy";
 
 bool IsSplitPolicyDevice() {
     return access(plat_policy_cil_file, R_OK) != -1;
@@ -499,19 +497,14 @@
 
 bool OpenMonolithicPolicy(PolicyFile* policy_file) {
     static constexpr char kSepolicyFile[] = "/sepolicy";
-    // In Microdroid the precompiled sepolicy is located on /system, since there is no vendor code.
-    // TODO(b/287206497): refactor once we start conditionally compiling init for Microdroid.
-    std::string monolithic_policy_file = access(kMicrodroidPrecompiledSepolicy, R_OK) == 0
-                                                 ? kMicrodroidPrecompiledSepolicy
-                                                 : kSepolicyFile;
 
-    LOG(INFO) << "Opening SELinux policy from monolithic file " << monolithic_policy_file;
-    policy_file->fd.reset(open(monolithic_policy_file.c_str(), O_RDONLY | O_CLOEXEC | O_NOFOLLOW));
+    LOG(INFO) << "Opening SELinux policy from monolithic file " << kSepolicyFile;
+    policy_file->fd.reset(open(kSepolicyFile, O_RDONLY | O_CLOEXEC | O_NOFOLLOW));
     if (policy_file->fd < 0) {
         PLOG(ERROR) << "Failed to open monolithic SELinux policy";
         return false;
     }
-    policy_file->path = monolithic_policy_file;
+    policy_file->path = kSepolicyFile;
     return true;
 }
 
@@ -858,6 +851,10 @@
 }
 
 int SelinuxGetVendorAndroidVersion() {
+    if (IsMicrodroid()) {
+        // As of now Microdroid doesn't have any vendor code.
+        return __ANDROID_API_FUTURE__;
+    }
     static int vendor_android_version = [] {
         if (!IsSplitPolicyDevice()) {
             // If this device does not split sepolicy files, it's not a Treble device and therefore,
@@ -961,6 +958,26 @@
     }
 }
 
+// Encapsulates steps to load SELinux policy in Microdroid.
+// So far the process is very straightforward - just load the precompiled policy from /system.
+void LoadSelinuxPolicyMicrodroid() {
+    constexpr const char kMicrodroidPrecompiledSepolicy[] =
+            "/system/etc/selinux/microdroid_precompiled_sepolicy";
+
+    LOG(INFO) << "Opening SELinux policy from " << kMicrodroidPrecompiledSepolicy;
+    unique_fd policy_fd(open(kMicrodroidPrecompiledSepolicy, O_RDONLY | O_CLOEXEC | O_NOFOLLOW));
+    if (policy_fd < 0) {
+        PLOG(FATAL) << "Failed to open " << kMicrodroidPrecompiledSepolicy;
+    }
+
+    std::string policy;
+    if (!android::base::ReadFdToString(policy_fd, &policy)) {
+        PLOG(FATAL) << "Failed to read policy file: " << kMicrodroidPrecompiledSepolicy;
+    }
+
+    LoadSelinuxPolicy(policy);
+}
+
 // The SELinux setup process is carefully orchestrated around snapuserd. Policy
 // must be loaded off dynamic partitions, and during an OTA, those partitions
 // cannot be read without snapuserd. But, with kernel-privileged snapuserd
@@ -976,20 +993,9 @@
 //  (5) Re-launch snapuserd and attach it to the dm-user devices from step (2).
 //
 // After this sequence, it is safe to enable enforcing mode and continue booting.
-int SetupSelinux(char** argv) {
-    SetStdioToDevNull(argv);
-    InitKernelLogging(argv);
-
-    if (REBOOT_BOOTLOADER_ON_PANIC) {
-        InstallRebootSignalHandlers();
-    }
-
-    boot_clock::time_point start_time = boot_clock::now();
-
+void LoadSelinuxPolicyAndroid() {
     MountMissingSystemPartitions();
 
-    SelinuxSetupKernelLogging();
-
     LOG(INFO) << "Opening SELinux policy";
 
     PrepareApexSepolicy();
@@ -1001,9 +1007,8 @@
 
     auto snapuserd_helper = SnapuserdSelinuxHelper::CreateIfNeeded();
     if (snapuserd_helper) {
-        // Kill the old snapused to avoid audit messages. After this we cannot
-        // read from /system (or other dynamic partitions) until we call
-        // FinishTransition().
+        // Kill the old snapused to avoid audit messages. After this we cannot read from /system
+        // (or other dynamic partitions) until we call FinishTransition().
         snapuserd_helper->StartTransition();
     }
 
@@ -1021,6 +1026,26 @@
     if (selinux_android_restorecon("/dev/selinux/", SELINUX_ANDROID_RESTORECON_RECURSE) == -1) {
         PLOG(FATAL) << "restorecon failed of /dev/selinux failed";
     }
+}
+
+int SetupSelinux(char** argv) {
+    SetStdioToDevNull(argv);
+    InitKernelLogging(argv);
+
+    if (REBOOT_BOOTLOADER_ON_PANIC) {
+        InstallRebootSignalHandlers();
+    }
+
+    boot_clock::time_point start_time = boot_clock::now();
+
+    SelinuxSetupKernelLogging();
+
+    // TODO(b/287206497): refactor into different headers to only include what we need.
+    if (IsMicrodroid()) {
+        LoadSelinuxPolicyMicrodroid();
+    } else {
+        LoadSelinuxPolicyAndroid();
+    }
 
     SelinuxSetEnforcement();
 
diff --git a/init/util.cpp b/init/util.cpp
index bc8ea6e..d0478e8 100644
--- a/init/util.cpp
+++ b/init/util.cpp
@@ -732,11 +732,6 @@
     is_default_mount_namespace_ready = true;
 }
 
-bool IsMicrodroid() {
-    static bool is_microdroid = android::base::GetProperty("ro.hardware", "") == "microdroid";
-    return is_microdroid;
-}
-
 bool Has32BitAbi() {
     static bool has = !android::base::GetProperty("ro.product.cpu.abilist32", "").empty();
     return has;
diff --git a/init/util.h b/init/util.h
index e58e70e..3f0a4e0 100644
--- a/init/util.h
+++ b/init/util.h
@@ -105,7 +105,14 @@
 bool IsDefaultMountNamespaceReady();
 void SetDefaultMountNamespaceReady();
 
-bool IsMicrodroid();
+inline constexpr bool IsMicrodroid() {
+#ifdef MICRODROID
+    return MICRODROID;
+#else
+    return false;
+#endif
+}
+
 bool Has32BitAbi();
 
 std::string GetApexNameFromFileName(const std::string& path);