//
// Copyright (C) 2020 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include <sys/types.h>
#include <unistd.h>

#include <algorithm>
#include <optional>
#include <utility>
#include <vector>

#include <android-base/file.h>
#include <android-base/logging.h>
#include <libsnapshot/cow_format.h>
#include <libsnapshot/cow_reader.h>
#include <storage_literals/storage_literals.h>
#include <zlib.h>

#include "cow_decompress.h"
#include "parser_v2.h"
#include "parser_v3.h"

namespace android {
namespace snapshot {

using namespace android::storage_literals;

bool ReadCowHeader(android::base::borrowed_fd fd, CowHeaderV3* header) {
    if (lseek(fd.get(), 0, SEEK_SET) < 0) {
        PLOG(ERROR) << "lseek header failed";
        return false;
    }

    memset(header, 0, sizeof(*header));

    if (!android::base::ReadFully(fd, &header->prefix, sizeof(header->prefix))) {
        return false;
    }
    if (header->prefix.magic != kCowMagicNumber) {
        LOG(ERROR) << "Header Magic corrupted. Magic: " << header->prefix.magic
                   << "Expected: " << kCowMagicNumber;
        return false;
    }
    if (header->prefix.header_size > sizeof(CowHeaderV3)) {
        LOG(ERROR) << "Unknown CowHeader size (got " << header->prefix.header_size
                   << " bytes, expected at most " << sizeof(CowHeaderV3) << " bytes)";
        return false;
    }

    if (lseek(fd.get(), 0, SEEK_SET) < 0) {
        PLOG(ERROR) << "lseek header failed";
        return false;
    }
    return android::base::ReadFully(fd, header, header->prefix.header_size);
}

CowReader::CowReader(ReaderFlags reader_flag, bool is_merge)
    : fd_(-1), header_(), fd_size_(0), reader_flag_(reader_flag), is_merge_(is_merge) {}

std::unique_ptr<CowReader> CowReader::CloneCowReader() {
    auto cow = std::make_unique<CowReader>();
    cow->owned_fd_.reset();
    cow->header_ = header_;
    cow->footer_ = footer_;
    cow->fd_size_ = fd_size_;
    cow->last_label_ = last_label_;
    cow->ops_ = ops_;
    cow->merge_op_start_ = merge_op_start_;
    cow->num_total_data_ops_ = num_total_data_ops_;
    cow->num_ordered_ops_to_merge_ = num_ordered_ops_to_merge_;
    cow->xor_data_loc_ = xor_data_loc_;
    cow->is_merge_ = is_merge_;
    return cow;
}

bool CowReader::InitForMerge(android::base::unique_fd&& fd) {
    owned_fd_ = std::move(fd);
    fd_ = owned_fd_.get();

    auto pos = lseek(fd_.get(), 0, SEEK_END);
    if (pos < 0) {
        PLOG(ERROR) << "lseek end failed";
        return false;
    }
    fd_size_ = pos;

    if (lseek(fd_.get(), 0, SEEK_SET) < 0) {
        PLOG(ERROR) << "lseek header failed";
        return false;
    }

    CHECK_GE(header_.prefix.header_size, sizeof(CowHeader));
    CHECK_LE(header_.prefix.header_size, sizeof(header_));

    if (!android::base::ReadFully(fd_, &header_, header_.prefix.header_size)) {
        PLOG(ERROR) << "read header failed";
        return false;
    }
    return true;
}

bool CowReader::Parse(android::base::unique_fd&& fd, std::optional<uint64_t> label, bool is_merge) {
    owned_fd_ = std::move(fd);
    return Parse(android::base::borrowed_fd{owned_fd_}, label, is_merge);
}

bool CowReader::Parse(android::base::borrowed_fd fd, std::optional<uint64_t> label, bool is_merge) {
    fd_ = fd;

    if (!ReadCowHeader(fd, &header_)) {
        return false;
    }

    std::unique_ptr<CowParserBase> parser;
    switch (header_.prefix.major_version) {
        case 1:
        case 2:
            parser = std::make_unique<CowParserV2>();
            break;
        case 3:
            parser = std::make_unique<CowParserV3>();
            break;
        default:
            LOG(ERROR) << "Unknown version: " << header_.prefix.major_version;
            return false;
    }
    if (!parser->Parse(fd, header_, label)) {
        return false;
    }

    TranslatedCowOps ops_info;
    if (!parser->Translate(&ops_info)) {
        return false;
    }

    header_ = ops_info.header;
    ops_ = std::move(ops_info.ops);
    footer_ = parser->footer();
    fd_size_ = parser->fd_size();
    last_label_ = parser->last_label();
    xor_data_loc_ = parser->xor_data_loc();

    // If we're resuming a write, we're not ready to merge
    if (label.has_value() || !is_merge) return true;
    return PrepMergeOps();
}

uint32_t CowReader::GetMaxCompressionSize() {
    switch (header_.prefix.major_version) {
        case 1:
        case 2:
            // Old versions supports only 4KB compression.
            return header_.block_size;
            ;
        case 3:
            return header_.max_compression_size;
        default:
            LOG(ERROR) << "Unknown version: " << header_.prefix.major_version;
            return 0;
    }
}

bool CowReader::VerifyMergeSequence(const std::vector<uint32_t>& vec) const {
    std::vector<bool> is_in_ops;
    for (const auto& op : *ops_) {
        if (IsMetadataOp(op)) {
            continue;
        }
        if (op.new_block >= is_in_ops.size()) {
            is_in_ops.resize(op.new_block + 1);
        }
        is_in_ops[op.new_block] = true;
    }
    for (const auto& block : vec) {
        if (block >= is_in_ops.size() || !is_in_ops[block]) {
            LOG(INFO) << "Block " << block << " is in merge sequence, but not part of any COW op.";
            return false;
        }
    }
    return true;
}

bool ByNewBlock(const CowOperation& a, const CowOperation& b) {
    const bool is_metadata_a = IsMetadataOp(a);
    const bool is_metadata_b = IsMetadataOp(b);
    if (is_metadata_a != is_metadata_b) {
        return is_metadata_a < is_metadata_b;
    }
    return a.new_block < b.new_block;
}

bool ByRevNewBlock(const CowOperation& a, const CowOperation& b) {
    return ByNewBlock(b, a);
}

bool CowReader::GetMergeOrder(std::vector<uint32_t>* merge_op_blocks) {
    switch (header_.prefix.major_version) {
        case 1:
        case 2:
            if (!GetSequenceDataV2(merge_op_blocks)) {
                LOG(ERROR) << "Failed to get V2 merge sequence data";
                return false;
            }
            break;
        case 3:
            if (!GetSequenceData(merge_op_blocks)) {
                LOG(ERROR) << "Failed to get V3 merge sequence data";
                return false;
            }
            break;
        default:
            break;
    }
    return true;
}

// Partition the ops array such that all ops which |p| returns true
// is before all ops which |p| returns false.
// std::partition does not guarantee relative ordering of elements
// But this implementation preserves relative ordering for elements
// which |p| returns true

template <class ForwardIt, class UnaryPred>
ForwardIt Partition(ForwardIt first, ForwardIt last, UnaryPred p) {
    first = std::find_if_not(first, last, p);
    if (first == last) return first;

    for (auto i = std::next(first); i != last; ++i)
        if (p(*i)) {
            std::iter_swap(i, first);
            ++first;
        }

    return first;
}
template <typename Pred>
size_t PartitionOps(std::vector<CowOperation>* ops, const Pred& p) {
    auto it = Partition(ops->begin(), ops->end(), p);
    size_t idx = std::distance(ops->begin(), it);
    return idx;
}

//
// This sets up the data needed for MergeOpIter. MergeOpIter presents
// data in the order we intend to merge in.
//
// We merge all order sensitive ops up front, and sort the rest to allow for
// batch merging. Order sensitive ops can either be presented in their proper
// order in the cow, or be ordered by sequence ops (kCowSequenceOp), in which
// case we want to merge those ops first, followed by any ops not specified by
// new_block value by the sequence op, in sorted order.
// We will re-arrange the vector in such a way that
// kernel can batch merge. Ex:
//
// Existing COW format; All the copy operations
// are at the beginning.
// =======================================
// Copy-op-1    - cow_op->new_block = 1
// Copy-op-2    - cow_op->new_block = 2
// Copy-op-3    - cow_op->new_block = 3
// Replace-op-4 - cow_op->new_block = 6
// Replace-op-5 - cow_op->new_block = 4
// Replace-op-6 - cow_op->new_block = 8
// Replace-op-7 - cow_op->new_block = 9
// Zero-op-8    - cow_op->new_block = 7
// Zero-op-9    - cow_op->new_block = 5
// =======================================
//
// First find the operation which isn't a copy-op
// and then sort all the operations in descending order
// with the key being cow_op->new_block (source block)
//
// The data-structure will look like:
//
// =======================================
// Copy-op-1    - cow_op->new_block = 1
// Copy-op-2    - cow_op->new_block = 2
// Copy-op-3    - cow_op->new_block = 3
// Replace-op-7 - cow_op->new_block = 9
// Replace-op-6 - cow_op->new_block = 8
// Zero-op-8    - cow_op->new_block = 7
// Replace-op-4 - cow_op->new_block = 6
// Zero-op-9    - cow_op->new_block = 5
// Replace-op-5 - cow_op->new_block = 4
// =======================================
//
// Daemon will read the above data-structure in reverse-order
// when reading metadata. Thus, kernel will get the metadata
// in the following order:
//
// ========================================
// Replace-op-5 - cow_op->new_block = 4
// Zero-op-9    - cow_op->new_block = 5
// Replace-op-4 - cow_op->new_block = 6
// Zero-op-8    - cow_op->new_block = 7
// Replace-op-6 - cow_op->new_block = 8
// Replace-op-7 - cow_op->new_block = 9
// Copy-op-3    - cow_op->new_block = 3
// Copy-op-2    - cow_op->new_block = 2
// Copy-op-1    - cow_op->new_block = 1
// ===========================================
//
// When merging begins, kernel will start from the last
// metadata which was read: In the above format, Copy-op-1
// will be the first merge operation.
//
// Now, batching of the merge operations happens only when
// 1: origin block numbers in the base device are contiguous
// (cow_op->new_block) and,
// 2: cow block numbers which are assigned by daemon in ReadMetadata()
// are contiguous. These are monotonically increasing numbers.
//
// When both (1) and (2) are true, kernel will batch merge the operations.
// In the above case, we have to ensure that the copy operations
// are merged first before replace operations are done. Hence,
// we will not change the order of copy operations. Since,
// cow_op->new_block numbers are contiguous, we will ensure that the
// cow block numbers assigned in ReadMetadata() for these respective copy
// operations are not contiguous forcing kernel to issue merge for each
// copy operations without batch merging.
//
// For all the other operations viz. Replace and Zero op, the cow block
// numbers assigned by daemon will be contiguous allowing kernel to batch
// merge.
//
// The final format after assiging COW block numbers by the daemon will
// look something like:
//
// =========================================================
// Replace-op-5 - cow_op->new_block = 4  cow-block-num = 2
// Zero-op-9    - cow_op->new_block = 5  cow-block-num = 3
// Replace-op-4 - cow_op->new_block = 6  cow-block-num = 4
// Zero-op-8    - cow_op->new_block = 7  cow-block-num = 5
// Replace-op-6 - cow_op->new_block = 8  cow-block-num = 6
// Replace-op-7 - cow_op->new_block = 9  cow-block-num = 7
// Copy-op-3    - cow_op->new_block = 3  cow-block-num = 9
// Copy-op-2    - cow_op->new_block = 2  cow-block-num = 11
// Copy-op-1    - cow_op->new_block = 1  cow-block-num = 13
// ==========================================================
//
// Merge sequence will look like:
//
// Merge-1 - Batch-merge { Copy-op-1, Copy-op-2, Copy-op-3 }
// Merge-2 - Batch-merge {Replace-op-7, Replace-op-6, Zero-op-8,
//                        Replace-op-4, Zero-op-9, Replace-op-5 }
//==============================================================

bool CowReader::PrepMergeOps() {
    std::vector<uint32_t> sequence_data;
    if (!GetMergeOrder(&sequence_data)) {
        return false;
    }
    // Metadata ops do not need to be merged, so strip
    ops_->erase(std::remove_if(ops_->begin(), ops_->end(), IsMetadataOp), ops_->end());
    size_t ordered_ops_count = 0;
    uint32_t max_sequence_block = 0;
    if (sequence_data.empty()) {
        ordered_ops_count = PartitionOps(ops_.get(), IsOrderedOp);
    } else {
        std::vector<bool> is_in_sequence;
        max_sequence_block = *std::max_element(sequence_data.begin(), sequence_data.end());
        is_in_sequence.resize(max_sequence_block + 1);
        for (const auto& block : sequence_data) {
            is_in_sequence[block] = true;
        }
        ordered_ops_count =
                PartitionOps(ops_.get(), [&is_in_sequence](const CowOperation& op) -> bool {
                    if (op.new_block >= is_in_sequence.size()) {
                        return false;
                    }
                    return is_in_sequence[op.new_block];
                });
        // This means there's some ops which are in merge sequence but not
        // found in COW image, run |VerifyMergeSequence| to print helpful debug
        // messages
        if (ordered_ops_count != sequence_data.size()) {
            LOG(ERROR) << "Invalid merge sequence";
            VerifyMergeSequence(sequence_data);
            return false;
        }
    }

    if (ordered_ops_count > header_.num_merge_ops) {
        num_ordered_ops_to_merge_ = ordered_ops_count - header_.num_merge_ops;
    } else {
        num_ordered_ops_to_merge_ = 0;
    }

    // Sort the vector in increasing order if merging in user-space as
    // we can batch merge them when iterating from forward.
    //
    // dm-snapshot-merge requires decreasing order as we iterate the blocks
    // in reverse order.
    if (reader_flag_ == ReaderFlags::USERSPACE_MERGE) {
        std::sort(ops_->begin() + ordered_ops_count, ops_->end(), ByNewBlock);
    } else {
        std::sort(ops_->begin() + ordered_ops_count, ops_->end(), ByRevNewBlock);
    }
    LOG(INFO) << "Merge sequence ordered ops " << ordered_ops_count << " other ops "
              << ops_->size() - ordered_ops_count;
    if (header_.num_merge_ops > 0) {
        merge_op_start_ = header_.num_merge_ops;
    }
    if (!sequence_data.empty()) {
        // For block X, what is its position in the merge sequence?
        std::vector<uint32_t> block_order;
        block_order.resize(max_sequence_block + 1, sequence_data.size());
        for (size_t i = 0; i < sequence_data.size(); i++) {
            block_order[sequence_data[i]] = i;
        }
        // Re-order ops array according to block numbers in merge sequence.
        auto& ops = *ops_;
        for (size_t i = 0; i < sequence_data.size(); i++) {
            CHECK_LT(ops[i].new_block, block_order.size()) << " " << ops[i];
            if (block_order[ops[i].new_block] >= sequence_data.size()) {
                LOG(ERROR) << "Op " << ops[i]
                           << " is not an ordered op, but get partition as an ordered op, this is "
                              "likely a bug in PartitionOps";
                return false;
            }
            while (i != block_order[ops[i].new_block]) {
                std::swap(ops[i], ops[block_order[ops[i].new_block]]);
            }
        }
    }
    num_total_data_ops_ = ops_->size();

    return true;
}

bool CowReader::GetSequenceDataV2(std::vector<uint32_t>* sequence_data) {
    size_t num_seqs = 0;
    size_t read;
    for (size_t i = 0; i < ops_->size(); i++) {
        auto& current_op = ops_->data()[i];

        if (current_op.type() == kCowSequenceOp) {
            size_t seq_len = current_op.data_length / sizeof(uint32_t);

            sequence_data->resize(sequence_data->size() + seq_len);
            if (!GetRawBytes(&current_op, &sequence_data->data()[num_seqs], current_op.data_length,
                             &read)) {
                PLOG(ERROR) << "Failed to read sequence op!";
                return false;
            }
            num_seqs += seq_len;
        }
    }
    return true;
}

bool CowReader::GetSequenceData(std::vector<uint32_t>* sequence_data) {
    // read sequence ops data
    sequence_data->resize(header_.sequence_data_count);
    if (!android::base::ReadFullyAtOffset(
                fd_, sequence_data->data(),
                header_.sequence_data_count * sizeof(sequence_data->at(0)),
                GetSequenceOffset(header_))) {
        PLOG(ERROR) << "failed to read sequence buffer. seq_data_count: "
                    << header_.sequence_data_count << " at offset: " << GetSequenceOffset(header_);
        return false;
    }
    return true;
}

const CowOperation* FindOpForBlock(uint64_t block, ICowOpIter* iter) {
    while (!iter->AtBegin()) {
        if (iter->Get()->new_block == block) {
            return iter->Get();
        }
        iter->Prev();
    }
    if (iter->Get()->new_block == block) {
        return iter->Get();
    }
    return nullptr;
}

bool CowReader::VerifyMergeOps() {
    auto itr = GetMergeOpIter(true);
    std::vector<bool> dirty_blocks;
    bool non_ordered_op_found = false;

    while (!itr->AtEnd()) {
        const auto& op = itr->Get();
        uint64_t offset;

        // Op should not be a metadata
        if (IsMetadataOp(*op)) {
            LOG(ERROR) << "Metadata op: " << *op << " found during merge sequence";
            return false;
        }

        // Sequence ops should contain all the ordered ops followed
        // by Replace and Zero ops. If we find the first op which
        // is not ordered, that means all ordered ops processing
        // has been completed.
        if (!IsOrderedOp(*op)) {
            non_ordered_op_found = true;
        }

        // Since, all ordered ops processing has been completed,
        // check that the subsequent ops are not ordered.
        if (non_ordered_op_found && IsOrderedOp(*op)) {
            LOG(ERROR) << "Invalid sequence - non-ordered and ordered ops"
                       << " cannot be mixed during sequence generation";
            return false;
        }

        if (!GetSourceOffset(op, &offset)) {
            itr->Next();
            continue;
        }

        uint64_t block = GetBlockFromOffset(header_, offset);
        bool misaligned = (GetBlockRelativeOffset(header_, offset) != 0);

        const CowOperation* overwrite = nullptr;
        if (block < dirty_blocks.size() && dirty_blocks[block]) {
            overwrite = FindOpForBlock(block, itr.get());
            CHECK(overwrite != nullptr);
            LOG(ERROR) << "Invalid Sequence! Block needed for op:\n"
                       << *op << "\noverwritten by previously merged op:\n"
                       << *overwrite;
        }
        if (misaligned && block + 1 < dirty_blocks.size() && dirty_blocks[block + 1]) {
            overwrite = FindOpForBlock(block + 1, itr.get());
            CHECK(overwrite != nullptr);
            LOG(ERROR) << "Invalid Sequence! Block needed for op:\n"
                       << op << "\noverwritten by previously merged op:\n"
                       << *overwrite;
        }
        if (overwrite != nullptr) return false;
        if (op->new_block >= dirty_blocks.size()) {
            dirty_blocks.resize(op->new_block + 1);
        }
        dirty_blocks[op->new_block] = true;
        itr->Next();
    }
    return true;
}

bool CowReader::GetFooter(CowFooter* footer) {
    if (!footer_) return false;
    *footer = footer_.value();
    return true;
}

bool CowReader::GetLastLabel(uint64_t* label) {
    if (!last_label_) return false;
    *label = last_label_.value();
    return true;
}

class CowOpIter final : public ICowOpIter {
  public:
    CowOpIter(std::shared_ptr<std::vector<CowOperation>>& ops, uint64_t start);

    bool AtEnd() override;
    const CowOperation* Get() override;
    void Next() override;

    void Prev() override;
    bool AtBegin() override;

  private:
    std::shared_ptr<std::vector<CowOperation>> ops_;
    std::vector<CowOperation>::iterator op_iter_;
};

CowOpIter::CowOpIter(std::shared_ptr<std::vector<CowOperation>>& ops, uint64_t start) {
    ops_ = ops;
    op_iter_ = ops_->begin() + start;
}

bool CowOpIter::AtBegin() {
    return op_iter_ == ops_->begin();
}

void CowOpIter::Prev() {
    CHECK(!AtBegin());
    op_iter_--;
}

bool CowOpIter::AtEnd() {
    return op_iter_ == ops_->end();
}

void CowOpIter::Next() {
    CHECK(!AtEnd());
    op_iter_++;
}

class CowRevOpIter final : public ICowOpIter {
  public:
    CowRevOpIter(std::shared_ptr<std::vector<CowOperation>>& ops, uint64_t start)
        : ops_(ops), start_(start) {
        op_iter_ = ops_->rbegin();
    }

    bool AtEnd() override { return op_iter_ == ops_->rend() - start_; }
    const CowOperation* Get() override { return &(*op_iter_); }
    void Next() override { ++op_iter_; }

    void Prev() override { --op_iter_; }
    bool AtBegin() override { return op_iter_ == ops_->rbegin(); }

  private:
    std::shared_ptr<std::vector<CowOperation>> ops_;
    std::vector<CowOperation>::reverse_iterator op_iter_;
    uint64_t start_;
};

const CowOperation* CowOpIter::Get() {
    CHECK(!AtEnd());
    return &(*op_iter_);
}

class CowRevMergeOpIter final : public ICowOpIter {
  public:
    explicit CowRevMergeOpIter(std::shared_ptr<std::vector<CowOperation>> ops,
                               std::shared_ptr<std::vector<int>> block_pos_index, uint64_t start);

    bool AtEnd() override;
    const CowOperation* Get() override;
    void Next() override;

    void Prev() override;
    bool AtBegin() override;

  private:
    std::shared_ptr<std::vector<CowOperation>> ops_;
    std::vector<int>::reverse_iterator block_riter_;
    std::shared_ptr<std::vector<int>> cow_op_index_vec_;
    uint64_t start_;
};

class CowMergeOpIter final : public ICowOpIter {
  public:
    explicit CowMergeOpIter(std::shared_ptr<std::vector<CowOperation>> ops,
                            std::shared_ptr<std::vector<int>> block_pos_index, uint64_t start);

    bool AtEnd() override;
    const CowOperation* Get() override;
    void Next() override;

    void Prev() override;
    bool AtBegin() override;

  private:
    std::shared_ptr<std::vector<CowOperation>> ops_;
    std::vector<int>::iterator block_iter_;
    std::shared_ptr<std::vector<int>> cow_op_index_vec_;
    uint64_t start_;
};

CowMergeOpIter::CowMergeOpIter(std::shared_ptr<std::vector<CowOperation>> ops,
                               std::shared_ptr<std::vector<int>> block_pos_index, uint64_t start) {
    ops_ = ops;
    start_ = start;
    cow_op_index_vec_ = block_pos_index;
    block_iter_ = cow_op_index_vec_->begin() + start;
}

bool CowMergeOpIter::AtBegin() {
    return block_iter_ == cow_op_index_vec_->begin();
}

void CowMergeOpIter::Prev() {
    CHECK(!AtBegin());
    block_iter_--;
}

bool CowMergeOpIter::AtEnd() {
    return block_iter_ == cow_op_index_vec_->end();
}

void CowMergeOpIter::Next() {
    CHECK(!AtEnd());
    block_iter_++;
}

const CowOperation* CowMergeOpIter::Get() {
    CHECK(!AtEnd());
    return &ops_->data()[*block_iter_];
}

CowRevMergeOpIter::CowRevMergeOpIter(std::shared_ptr<std::vector<CowOperation>> ops,
                                     std::shared_ptr<std::vector<int>> block_pos_index,
                                     uint64_t start) {
    ops_ = ops;
    start_ = start;
    cow_op_index_vec_ = block_pos_index;
    block_riter_ = cow_op_index_vec_->rbegin();
}

bool CowRevMergeOpIter::AtBegin() {
    return block_riter_ == cow_op_index_vec_->rbegin();
}

void CowRevMergeOpIter::Prev() {
    CHECK(!AtBegin());
    block_riter_--;
}

bool CowRevMergeOpIter::AtEnd() {
    return block_riter_ == cow_op_index_vec_->rend() - start_;
}

void CowRevMergeOpIter::Next() {
    CHECK(!AtEnd());
    block_riter_++;
}

const CowOperation* CowRevMergeOpIter::Get() {
    CHECK(!AtEnd());
    return &ops_->data()[*block_riter_];
}

std::unique_ptr<ICowOpIter> CowReader::GetOpIter(bool merge_progress) {
    return std::make_unique<CowOpIter>(ops_, merge_progress ? merge_op_start_ : 0);
}

std::unique_ptr<ICowOpIter> CowReader::GetRevOpIter(bool merge_progress) {
    return std::make_unique<CowRevOpIter>(ops_, merge_progress ? merge_op_start_ : 0);
}

std::unique_ptr<ICowOpIter> CowReader::GetRevMergeOpIter(bool ignore_progress) {
    // num_total_data_ops_ == 0 means PrepMergeOps() never called, merge order
    // is not initialized.
    if (num_total_data_ops_ == 0) {
        return std::make_unique<CowRevOpIter>(ops_, ops_->size());
    }
    return GetRevOpIter(!ignore_progress);
}

std::unique_ptr<ICowOpIter> CowReader::GetMergeOpIter(bool ignore_progress) {
    // num_total_data_ops_ == 0 means PrepMergeOps() never called, merge order
    // is not initialized.
    if (num_total_data_ops_ == 0) {
        return std::make_unique<CowOpIter>(ops_, ops_->size());
    }
    return GetOpIter(!ignore_progress);
}

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(op->source(), 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_ || offset + len > fd_size_ ||
        len >= fd_size_) {
        LOG(ERROR) << "invalid data offset: " << offset << ", " << len << " bytes";
        return false;
    }
    ssize_t rv = TEMP_FAILURE_RETRY(::pread(fd_.get(), buffer, len, offset));
    if (rv < 0) {
        PLOG(ERROR) << "read failed";
        return false;
    }
    *read = rv;
    return true;
}

class CowDataStream final : public IByteStream {
  public:
    CowDataStream(CowReader* reader, uint64_t offset, size_t data_length)
        : reader_(reader), offset_(offset), data_length_(data_length) {
        remaining_ = data_length_;
    }

    ssize_t Read(void* buffer, size_t length) override {
        size_t to_read = std::min(length, remaining_);
        if (!to_read) {
            return 0;
        }
        size_t read;
        if (!reader_->GetRawBytes(offset_, buffer, to_read, &read)) {
            return -1;
        }
        offset_ += read;
        remaining_ -= read;
        return read;
    }

    size_t Size() const override { return data_length_; }

  private:
    CowReader* reader_;
    uint64_t offset_;
    size_t data_length_;
    size_t remaining_;
};

uint8_t CowReader::GetCompressionType() {
    return header_.compression_algorithm;
}

ssize_t CowReader::ReadData(const CowOperation* op, void* buffer, size_t buffer_size,
                            size_t ignore_bytes) {
    std::unique_ptr<IDecompressor> decompressor;
    const size_t op_buf_size = CowOpCompressionSize(op, header_.block_size);
    if (!op_buf_size) {
        LOG(ERROR) << "Compression size is zero. op: " << *op;
        return -1;
    }
    switch (GetCompressionType()) {
        case kCowCompressNone:
            break;
        case kCowCompressGz:
            decompressor = IDecompressor::Gz();
            break;
        case kCowCompressZstd:
            if (op_buf_size != op->data_length) {
                decompressor = IDecompressor::Zstd();
            }
            break;
        case kCowCompressLz4:
            if (op_buf_size != op->data_length) {
                decompressor = IDecompressor::Lz4();
            }
            break;
        default:
            LOG(ERROR) << "Unknown compression type: " << GetCompressionType();
            return -1;
    }

    uint64_t offset;
    if (op->type() == kCowXorOp) {
        offset = xor_data_loc_->at(op->new_block);
    } else {
        offset = op->source();
    }
    if (!decompressor ||
        ((op->data_length == op_buf_size) && (header_.prefix.major_version == 3))) {
        CowDataStream stream(this, offset + ignore_bytes, op->data_length - ignore_bytes);
        return stream.ReadFully(buffer, buffer_size);
    }

    CowDataStream stream(this, offset, op->data_length);
    decompressor->set_stream(&stream);
    return decompressor->Decompress(buffer, buffer_size, op_buf_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
