// Copyright 2015 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/utility/safe_browsing/mac/udif.h"

#include <CoreFoundation/CoreFoundation.h>
#include <bzlib.h>
#include <libkern/OSByteOrder.h>
#include <uuid/uuid.h>

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

#include "base/apple/foundation_util.h"
#include "base/apple/scoped_cftyperef.h"
#include "base/compiler_specific.h"
#include "base/containers/buffer_iterator.h"
#include "base/containers/span.h"
#include "base/logging.h"
#include "base/memory/raw_ptr.h"
#include "base/notreached.h"
#include "base/numerics/ostream_operators.h"
#include "base/numerics/safe_math.h"
#include "base/strings/sys_string_conversions.h"
#include "chrome/utility/safe_browsing/mac/convert_big_endian.h"
#include "chrome/utility/safe_browsing/mac/read_stream.h"
#include "third_party/zlib/zlib.h"

namespace safe_browsing {
namespace dmg {

#pragma pack(push, 1)

// The following structures come from the analysis provided by Jonathan Levin
// at <http://newosxbook.com/DMG.html>.
//
// Note that all fields are stored in big endian.

struct UDIFChecksum {
  uint32_t type;
  uint32_t size;
  std::array<uint32_t, 32> data;
};

static void ConvertBigEndian(UDIFChecksum* checksum) {
  ConvertBigEndian(&checksum->type);
  ConvertBigEndian(&checksum->size);
  for (size_t i = 0; i < std::size(checksum->data); ++i) {
    ConvertBigEndian(&checksum->data[i]);
  }
}

// The trailer structure for a UDIF file.
struct UDIFResourceFile {
  static const uint32_t kSignature = 'koly';
  static const uint32_t kVersion = 4;

  uint32_t signature;
  uint32_t version;
  uint32_t header_size;  // Size of this structure.
  uint32_t flags;
  uint64_t running_data_fork_offset;
  uint64_t data_fork_offset;
  uint64_t data_fork_length;
  uint64_t rsrc_fork_offset;
  uint64_t rsrc_fork_length;
  uint32_t segment_number;
  uint32_t segment_count;
  uuid_t   segment_id;

  UDIFChecksum data_checksum;

  uint64_t plist_offset;  // Offset and length of the blkx plist.
  uint64_t plist_length;

  uint8_t  reserved1[64];

  uint64_t code_signature_offset;
  uint64_t code_signature_length;

  uint8_t  reserved2[40];

  UDIFChecksum main_checksum;

  uint32_t image_variant;
  uint64_t sector_count;

  uint32_t reserved3;
  uint32_t reserved4;
  uint32_t reserved5;
};

static void ConvertBigEndian(uuid_t* uuid) {
  // UUID is never consulted, so do not swap.
}

static void ConvertBigEndian(UDIFResourceFile* file) {
  ConvertBigEndian(&file->signature);
  ConvertBigEndian(&file->version);
  ConvertBigEndian(&file->flags);
  ConvertBigEndian(&file->header_size);
  ConvertBigEndian(&file->running_data_fork_offset);
  ConvertBigEndian(&file->data_fork_offset);
  ConvertBigEndian(&file->data_fork_length);
  ConvertBigEndian(&file->rsrc_fork_offset);
  ConvertBigEndian(&file->rsrc_fork_length);
  ConvertBigEndian(&file->segment_number);
  ConvertBigEndian(&file->segment_count);
  ConvertBigEndian(&file->segment_id);
  ConvertBigEndian(&file->data_checksum);
  ConvertBigEndian(&file->plist_offset);
  ConvertBigEndian(&file->plist_length);
  ConvertBigEndian(&file->code_signature_offset);
  ConvertBigEndian(&file->code_signature_length);
  ConvertBigEndian(&file->main_checksum);
  ConvertBigEndian(&file->image_variant);
  // `sector_count` is never consulted, so do not swap.
  // Note: If this is ever needed in the future, one must make a copy when byte
  // swapping to avoid unaligned access.

  // Reserved fields are skipped.
}

struct UDIFBlockChunk {
  enum class Type : uint32_t {
    ZERO_FILL     = 0x00000000,
    UNCOMPRESSED  = 0x00000001,
    IGNORED       = 0x00000002,
    COMPRESS_ADC  = 0x80000004,
    COMPRESS_ZLIB = 0x80000005,
    COMPRESSS_BZ2 = 0x80000006,
    COMMENT       = 0x7ffffffe,
    LAST_BLOCK    = 0xffffffff,
  };

  Type type;
  uint32_t comment;
  uint64_t start_sector;  // Logical chunk offset and length, in sectors.
  uint64_t sector_count;
  uint64_t compressed_offset;  // Compressed offset and length, in bytes.
  uint64_t compressed_length;
};

static void ConvertBigEndian(UDIFBlockChunk* chunk) {
  ConvertBigEndian(reinterpret_cast<uint32_t*>(&chunk->type));
  ConvertBigEndian(&chunk->comment);
  ConvertBigEndian(&chunk->start_sector);
  ConvertBigEndian(&chunk->sector_count);
  ConvertBigEndian(&chunk->compressed_offset);
  ConvertBigEndian(&chunk->compressed_length);
}

struct UDIFBlockData {
  static const uint32_t kSignature = 'mish';
  static const uint32_t kVersion = 1;

  uint32_t signature;
  uint32_t version;
  uint64_t start_sector;  // Logical block offset and length, in sectors.
  uint64_t sector_count;

  uint64_t data_offset;
  uint32_t buffers_needed;
  uint32_t block_descriptors;

  uint32_t reserved1;
  uint32_t reserved2;
  uint32_t reserved3;
  uint32_t reserved4;
  uint32_t reserved5;
  uint32_t reserved6;

  UDIFChecksum checksum;

  uint32_t chunk_count;
};

static void ConvertBigEndian(UDIFBlockData* block) {
  ConvertBigEndian(&block->signature);
  ConvertBigEndian(&block->version);
  ConvertBigEndian(&block->start_sector);
  ConvertBigEndian(&block->sector_count);
  ConvertBigEndian(&block->data_offset);
  ConvertBigEndian(&block->buffers_needed);
  ConvertBigEndian(&block->block_descriptors);
  // Reserved fields are skipped.
  ConvertBigEndian(&block->checksum);
  ConvertBigEndian(&block->chunk_count);
}

#pragma pack(pop)

// UDIFBlock takes a raw, big-endian block data pointer and stores, in host
// endian, the data for both the block and the chunk.
class UDIFBlock {
 public:
  UDIFBlock() : block_() {}

  UDIFBlock(const UDIFBlock&) = delete;
  UDIFBlock& operator=(const UDIFBlock&) = delete;

  bool ParseBlockData(base::span<const uint8_t> block_data,
                      uint16_t sector_size) {
    base::BufferIterator iterator(block_data);
    const UDIFBlockData* block_header = iterator.Object<UDIFBlockData>();
    if (!block_header) {
      DLOG(ERROR) << "UDIF block data is smaller than expected";
      return false;
    }

    block_ = *block_header;
    ConvertBigEndian(&block_);

    // Make sure the number of sectors doesn't overflow.
    auto block_size = base::CheckedNumeric<size_t>(sector_count()) *
                      sector_size;
    if (!block_size.IsValid()) {
      DLOG(ERROR) << "UDIF block size overflows";
      return false;
    }

    // Make sure the block data contains the reported number of chunks.
    auto block_and_chunks_size =
        (base::CheckedNumeric<size_t>(sizeof(UDIFBlockChunk)) *
         block_.chunk_count) +
        sizeof(block_);
    if (!block_and_chunks_size.IsValid() ||
        block_data.size() < block_and_chunks_size.ValueOrDie()) {
      DLOG(ERROR) << "UDIF block does not contain reported number of chunks, "
                  << block_and_chunks_size.ValueOrDie() << " bytes expected, "
                  << "got " << block_data.size();
      return false;
    }

    // Make sure that the chunk data isn't larger than the block reports.
    base::CheckedNumeric<size_t> chunk_sectors(0);
    for (uint32_t i = 0; i < block_.chunk_count; ++i) {
      const UDIFBlockChunk* raw_chunk = iterator.Object<UDIFBlockChunk>();
      // Total size check above should ensure that the chunk always exists
      CHECK(raw_chunk);
      chunks_.push_back(*raw_chunk);

      UDIFBlockChunk* chunk = &chunks_[i];
      ConvertBigEndian(chunk);

      chunk_sectors += chunk->sector_count;
      if (!chunk_sectors.IsValid() ||
          chunk_sectors.ValueOrDie() > sector_count()) {
        DLOG(ERROR) << "Total chunk sectors larger than reported block sectors";
        return false;
      }

      auto chunk_end_offset =
          base::CheckedNumeric<size_t>(chunk->compressed_offset) +
          chunk->compressed_length;
      if (!chunk_end_offset.IsValid() ||
          chunk->compressed_length > block_size.ValueOrDie()) {
        DLOG(ERROR) << "UDIF chunk data length " << i << " overflows";
        return false;
      }
    }

    return true;
  }

  uint32_t signature() const { return block_.signature; }
  uint32_t version() const { return block_.version; }
  uint64_t start_sector() const { return block_.start_sector; }
  uint64_t sector_count() const { return block_.sector_count; }
  uint64_t chunk_count() const { return chunks_.size(); }

  const UDIFBlockChunk* chunk(uint32_t i) const {
    if (i >= chunk_count())
      return nullptr;
    return &chunks_[i];
  }

 private:
  UDIFBlockData block_;
  std::vector<UDIFBlockChunk> chunks_;
};

namespace {

const size_t kSectorSize = 512;

class UDIFBlockChunkReadStream;

// A UDIFPartitionReadStream virtualizes a partition's non-contiguous blocks
// into a single stream.
class UDIFPartitionReadStream : public ReadStream {
 public:
  UDIFPartitionReadStream(ReadStream* stream,
                          uint16_t block_size,
                          const UDIFBlock* partition_block);

  UDIFPartitionReadStream(const UDIFPartitionReadStream&) = delete;
  UDIFPartitionReadStream& operator=(const UDIFPartitionReadStream&) = delete;

  ~UDIFPartitionReadStream() override;

  bool Read(base::span<uint8_t> buf, size_t* bytes_read) override;
  // Seek only supports SEEK_SET and SEEK_CUR.
  off_t Seek(off_t offset, int whence) override;

 private:
  const raw_ptr<ReadStream> stream_;  // The UDIF stream.
  const uint16_t block_size_;  // The UDIF block size.
  const raw_ptr<const UDIFBlock> block_;  // The block for this partition.
  uint64_t current_chunk_;  // The current chunk number.
  // The current chunk stream.
  std::unique_ptr<UDIFBlockChunkReadStream> chunk_stream_;
};

// A ReadStream for a single block chunk, which transparently handles
// decompression.
class UDIFBlockChunkReadStream : public ReadStream {
 public:
  UDIFBlockChunkReadStream(ReadStream* stream,
                           uint16_t block_size,
                           const UDIFBlockChunk* chunk);

  UDIFBlockChunkReadStream(const UDIFBlockChunkReadStream&) = delete;
  UDIFBlockChunkReadStream& operator=(const UDIFBlockChunkReadStream&) = delete;

  ~UDIFBlockChunkReadStream() override;

  bool Read(base::span<uint8_t> buf, size_t* bytes_read) override;
  // Seek only supports SEEK_SET.
  off_t Seek(off_t offset, int whence) override;

  bool IsAtEnd() { return offset_ >= length_in_bytes_; }

  const UDIFBlockChunk* chunk() const { return chunk_; }
  size_t length_in_bytes() const { return length_in_bytes_; }

 private:
  bool CopyOutZeros(base::span<uint8_t> buf, size_t* bytes_read);
  bool CopyOutUncompressed(base::span<uint8_t> buf, size_t* bytes_read);
  bool CopyOutDecompressed(base::span<uint8_t> buf, size_t* bytes_read);
  bool HandleADC(base::span<uint8_t> buf, size_t* bytes_read);
  bool HandleZLib(base::span<uint8_t> buf, size_t* bytes_read);
  bool HandleBZ2(base::span<uint8_t> buf, size_t* bytes_read);

  // Reads from |stream_| |chunk_->compressed_length| bytes, starting at
  // |chunk_->compressed_offset|. Returns (possibly empty) vector containing
  // data, or nullopt on error.
  std::optional<std::vector<uint8_t>> ReadCompressedData();

  const raw_ptr<ReadStream> stream_;           // The UDIF stream.
  const raw_ptr<const UDIFBlockChunk> chunk_;  // The chunk to be read.
  size_t length_in_bytes_;  // The decompressed length in bytes.
  size_t offset_;  // The offset into the decompressed buffer.
  std::vector<uint8_t> decompress_buffer_;  // Decompressed data buffer.
  bool did_decompress_;  // Whether or not the chunk has been decompressed.
};

}  // namespace

UDIFParser::UDIFParser(ReadStream* stream)
    : stream_(stream),
      partition_names_(),
      blocks_(),
      block_size_(kSectorSize) {}

UDIFParser::~UDIFParser() = default;

bool UDIFParser::Parse() {
  if (!ParseBlkx())
    return false;

  return true;
}

const std::vector<uint8_t>& UDIFParser::GetCodeSignature() {
  return signature_blob_;
}

size_t UDIFParser::GetNumberOfPartitions() {
  return blocks_.size();
}

std::string UDIFParser::GetPartitionName(size_t part_number) {
  DCHECK_LT(part_number, partition_names_.size());
  return partition_names_[part_number];
}

std::string UDIFParser::GetPartitionType(size_t part_number) {
  // The partition type is embedded in the Name field, as such:
  // "Partition-Name (Partition-Type : Partition-ID)".
  std::string name = GetPartitionName(part_number);
  size_t open = name.rfind('(');
  size_t separator = name.rfind(':');
  if (open == std::string::npos || separator == std::string::npos)
    return std::string();

  // Name does not end in ')' or no space after ':'.
  if (*(name.end() - 1) != ')' ||
      (name.size() - separator < 2 || name[separator + 1] != ' ')) {
    return std::string();
  }

  --separator;
  ++open;
  if (separator <= open)
    return std::string();
  return name.substr(open, separator - open);
}

size_t UDIFParser::GetPartitionSize(size_t part_number) {
  DCHECK_LT(part_number, blocks_.size());
  auto size =
      base::CheckedNumeric<size_t>(blocks_[part_number]->sector_count()) *
      block_size_;
  return size.ValueOrDie();
}

std::unique_ptr<ReadStream> UDIFParser::GetPartitionReadStream(
    size_t part_number) {
  DCHECK_LT(part_number, blocks_.size());
  return std::make_unique<UDIFPartitionReadStream>(stream_, block_size_,
                                                   blocks_[part_number].get());
}

bool UDIFParser::ParseBlkx() {
  UDIFResourceFile trailer;
  off_t trailer_start = stream_->Seek(-sizeof(trailer), SEEK_END);
  if (trailer_start == -1)
    return false;

  if (!stream_->ReadType(trailer)) {
    DLOG(ERROR) << "Failed to read UDIFResourceFile";
    return false;
  }
  ConvertBigEndian(&trailer);

  if (trailer.signature != trailer.kSignature) {
    DLOG(ERROR) << "blkx signature does not match, is 0x"
                << std::hex << trailer.signature;
    return false;
  }
  if (trailer.version != trailer.kVersion) {
    DLOG(ERROR) << "blkx version does not match, is " << trailer.version;
    return false;
  }

  auto plist_end = base::CheckedNumeric<size_t>(trailer.plist_offset) +
                   trailer.plist_length;
  if (!plist_end.IsValid() ||
      plist_end.ValueOrDie() > base::checked_cast<size_t>(trailer_start)) {
    DLOG(ERROR) << "blkx plist extends past UDIF trailer";
    return false;
  }

  std::vector<uint8_t> plist_bytes(trailer.plist_length, 0);

  if (stream_->Seek(trailer.plist_offset, SEEK_SET) == -1)
    return false;

  if (trailer.plist_length == 0 || !stream_->ReadExact(plist_bytes)) {
    DLOG(ERROR) << "Failed to read blkx plist data";
    return false;
  }

  base::apple::ScopedCFTypeRef<CFDataRef> plist_data(
      CFDataCreateWithBytesNoCopy(kCFAllocatorDefault, plist_bytes.data(),
                                  plist_bytes.size(), kCFAllocatorNull));
  if (!plist_data) {
    DLOG(ERROR) << "Failed to create data from bytes";
    return false;
  }

  CFErrorRef error = nullptr;
  base::apple::ScopedCFTypeRef<CFPropertyListRef> plist(
      CFPropertyListCreateWithData(kCFAllocatorDefault, plist_data.get(),
                                   kCFPropertyListImmutable, nullptr, &error));

  CFDictionaryRef plist_dict =
      base::apple::CFCast<CFDictionaryRef>(plist.get());
  base::apple::ScopedCFTypeRef<CFErrorRef> error_ref(error);
  if (error) {
    base::apple::ScopedCFTypeRef<CFStringRef> error_string(
        CFErrorCopyDescription(error));
    DLOG(ERROR) << "Failed to parse XML plist: "
                << base::SysCFStringRefToUTF8(error_string.get());
    return false;
  }

  if (!plist_dict) {
    DLOG(ERROR) << "Plist is not a dictionary";
    return false;
  }

  auto* resource_fork = base::apple::GetValueFromDictionary<CFDictionaryRef>(
      plist_dict, CFSTR("resource-fork"));
  if (!resource_fork) {
    DLOG(ERROR) << "No resource-fork entry in plist";
    return false;
  }

  auto* blkx = base::apple::GetValueFromDictionary<CFArrayRef>(resource_fork,
                                                               CFSTR("blkx"));
  if (!blkx) {
    DLOG(ERROR) << "No blkx entry in resource-fork";
    return false;
  }

  for (CFIndex i = 0; i < CFArrayGetCount(blkx); ++i) {
    auto* block_dictionary =
        base::apple::CFCast<CFDictionaryRef>(CFArrayGetValueAtIndex(blkx, i));
    if (!block_dictionary) {
      DLOG(ERROR) << "Skipping block " << i
                  << " because it is not a CFDictionary";
      continue;
    }

    auto* data = base::apple::GetValueFromDictionary<CFDataRef>(
        block_dictionary, CFSTR("Data"));
    if (!data) {
      DLOG(ERROR) << "Skipping block " << i
                  << " because it has no Data section";
      continue;
    }

    // Copy the block table out of the plist.
    auto block = std::make_unique<UDIFBlock>();
    // SAFETY: CFDataGetBytePtr is provided by Apple and documented to
    // return CFDataGetLength bytes.
    if (!block->ParseBlockData(UNSAFE_BUFFERS(
            base::span(CFDataGetBytePtr(data),
                       base::checked_cast<size_t>(CFDataGetLength(data))),
            block_size_))) {
      DLOG(ERROR) << "Failed to parse UDIF block data";
      return false;
    }

    if (block->signature() != UDIFBlockData::kSignature) {
      DLOG(ERROR) << "Skipping block " << i << " because its signature does not"
                  << " match, is 0x" << std::hex << block->signature();
      continue;
    }
    if (block->version() != UDIFBlockData::kVersion) {
      DLOG(ERROR) << "Skipping block " << i << "because its version does not "
                  << "match, is " << block->version();
      continue;
    }

    CFStringRef partition_name_cf = base::apple::CFCast<CFStringRef>(
        CFDictionaryGetValue(block_dictionary, CFSTR("Name")));
    if (!partition_name_cf) {
      DLOG(ERROR) << "Skipping block " << i << " because it has no name";
      continue;
    }
    std::string partition_name = base::SysCFStringRefToUTF8(partition_name_cf);

    if (DLOG_IS_ON(INFO) && VLOG_IS_ON(1)) {
      DVLOG(1) << "Name: " << partition_name;
      DVLOG(1) << "StartSector = " << block->start_sector()
               << ", SectorCount = " << block->sector_count()
               << ", ChunkCount = " << block->chunk_count();
      for (uint32_t j = 0; j < block->chunk_count(); ++j) {
        const UDIFBlockChunk* chunk = block->chunk(j);
        DVLOG(1) << "Chunk#" << j
                 << " type = " << std::hex << static_cast<uint32_t>(chunk->type)
                 << ", StartSector = " << std::dec << chunk->start_sector
                 << ", SectorCount = " << chunk->sector_count
                 << ", CompressOffset = " << chunk->compressed_offset
                 << ", CompressLen = " << chunk->compressed_length;
      }
    }

    blocks_.push_back(std::move(block));
    partition_names_.push_back(partition_name);
  }

  // The offsets in the trailer could be garbage in DMGs that aren't signed.
  // Need a sanity check that the DMG has legit values for these fields.
  if (trailer.code_signature_length != 0 && trailer_start > 0) {
    auto code_signature_end =
        base::CheckedNumeric<size_t>(trailer.code_signature_offset) +
        trailer.code_signature_length;
    if (code_signature_end.IsValid() &&
        code_signature_end.ValueOrDie() <=
            base::checked_cast<size_t>(trailer_start)) {
      signature_blob_.resize(trailer.code_signature_length);

      off_t code_signature_start =
          stream_->Seek(trailer.code_signature_offset, SEEK_SET);
      if (code_signature_start == -1)
        return false;

      size_t bytes_read = 0;

      if (!stream_->Read(signature_blob_, &bytes_read)) {
        DLOG(ERROR) << "Failed to read raw signature bytes";
        return false;
      }

      if (bytes_read != trailer.code_signature_length) {
        DLOG(ERROR) << "Read unexpected number of raw signature bytes";
        return false;
      }
    }
  }

  return true;
}

namespace {

UDIFPartitionReadStream::UDIFPartitionReadStream(
    ReadStream* stream,
    uint16_t block_size,
    const UDIFBlock* partition_block)
    : stream_(stream),
      block_size_(block_size),
      block_(partition_block),
      current_chunk_(0),
      chunk_stream_() {
}

UDIFPartitionReadStream::~UDIFPartitionReadStream() = default;

bool UDIFPartitionReadStream::Read(base::span<uint8_t> buf,
                                   size_t* bytes_read) {
  size_t buffer_space_remaining = buf.size();
  *bytes_read = 0;

  for (uint32_t i = current_chunk_; i < block_->chunk_count(); ++i) {
    const UDIFBlockChunk* chunk = block_->chunk(i);

    // If this is the last block chunk, then the read is complete.
    if (chunk->type == UDIFBlockChunk::Type::LAST_BLOCK) {
      break;
    }

    // If the buffer is full, then the read is complete.
    if (buffer_space_remaining == 0)
      break;

    // A chunk stream may exist if the last read from this chunk was partial,
    // or if the stream was Seek()ed.
    if (!chunk_stream_) {
      chunk_stream_ = std::make_unique<UDIFBlockChunkReadStream>(
          stream_, block_size_, chunk);
    }
    DCHECK_EQ(chunk, chunk_stream_->chunk());

    size_t chunk_bytes_read = 0;
    if (!chunk_stream_->Read(buf.last(buffer_space_remaining),
                             &chunk_bytes_read)) {
      DLOG(ERROR) << "Failed to read " << buffer_space_remaining << " bytes "
                  << "from chunk " << i;
      return false;
    }
    *bytes_read += chunk_bytes_read;
    buffer_space_remaining -= chunk_bytes_read;

    if (chunk_stream_->IsAtEnd()) {
      chunk_stream_.reset();
      ++current_chunk_;
    }
  }

  return true;
}

off_t UDIFPartitionReadStream::Seek(off_t offset, int whence) {
  // Translate SEEK_END to SEEK_SET. SEEK_CUR is not currently supported.
  if (whence == SEEK_END) {
    base::CheckedNumeric<off_t> safe_offset(block_->sector_count());
    safe_offset *= block_size_;
    safe_offset += offset;
    if (!safe_offset.IsValid()) {
      DLOG(ERROR) << "Seek offset overflows";
      return -1;
    }
    offset = safe_offset.ValueOrDie();
  } else if (whence != SEEK_SET) {
    DCHECK_EQ(SEEK_SET, whence);
  }

  uint64_t sector = offset / block_size_;

  // Find the chunk for this sector.
  uint32_t chunk_number = 0;
  const UDIFBlockChunk* chunk = nullptr;
  for (uint32_t i = 0; i < block_->chunk_count(); ++i) {
    const UDIFBlockChunk* chunk_it = block_->chunk(i);
    // This assumes that all the chunks are ordered by sector.
    if (i != 0) {
      DLOG_IF(ERROR,
              chunk_it->start_sector < block_->chunk(i - 1)->start_sector)
          << "Chunks are not ordered by sector at chunk " << i
          << " , previous start_sector = "
          << block_->chunk(i - 1)->start_sector << ", current = "
          << chunk_it->start_sector;
    }
    if (sector >= chunk_it->start_sector) {
      chunk = chunk_it;
      chunk_number = i;
    } else {
      break;
    }
  }
  if (!chunk) {
    DLOG(ERROR) << "Failed to Seek to partition offset " << offset;
    return -1;
  }

  // Compute the offset into the chunk.
  uint64_t offset_in_sector = offset % block_size_;
  uint64_t start_sector = sector - chunk->start_sector;
  base::CheckedNumeric<uint64_t> decompress_read_offset(start_sector);
  decompress_read_offset *= block_size_;
  decompress_read_offset += offset_in_sector;

  if (!decompress_read_offset.IsValid()) {
    DLOG(ERROR) << "Partition decompress read offset overflows";
    return -1;
  }

  if (!chunk_stream_ || chunk != chunk_stream_->chunk()) {
    chunk_stream_ =
        std::make_unique<UDIFBlockChunkReadStream>(stream_, block_size_, chunk);
  }
  current_chunk_ = chunk_number;
  if (chunk_stream_->Seek(
          base::ValueOrDieForType<off_t>(decompress_read_offset), SEEK_SET) ==
      -1)
    return -1;

  return offset;
}

UDIFBlockChunkReadStream::UDIFBlockChunkReadStream(ReadStream* stream,
                                                   uint16_t block_size,
                                                   const UDIFBlockChunk* chunk)
    : stream_(stream),
      chunk_(chunk),
      length_in_bytes_(chunk->sector_count * block_size),
      offset_(0),
      decompress_buffer_(),
      did_decompress_(false) {
  // Make sure the multiplication above did not overflow.
  CHECK(length_in_bytes_ == 0 || length_in_bytes_ >= block_size);
}

UDIFBlockChunkReadStream::~UDIFBlockChunkReadStream() = default;

bool UDIFBlockChunkReadStream::Read(base::span<uint8_t> buf,
                                    size_t* bytes_read) {
  switch (chunk_->type) {
    case UDIFBlockChunk::Type::ZERO_FILL:
    case UDIFBlockChunk::Type::IGNORED:
      return CopyOutZeros(buf, bytes_read);
    case UDIFBlockChunk::Type::UNCOMPRESSED:
      return CopyOutUncompressed(buf, bytes_read);
    case UDIFBlockChunk::Type::COMPRESS_ADC:
      return HandleADC(buf, bytes_read);
    case UDIFBlockChunk::Type::COMPRESS_ZLIB:
      return HandleZLib(buf, bytes_read);
    case UDIFBlockChunk::Type::COMPRESSS_BZ2:
      return HandleBZ2(buf, bytes_read);
    case UDIFBlockChunk::Type::COMMENT:
      NOTREACHED();
    case UDIFBlockChunk::Type::LAST_BLOCK:
      *bytes_read = 0;
      return true;
  }
  return false;
}

off_t UDIFBlockChunkReadStream::Seek(off_t offset, int whence) {
  DCHECK_EQ(SEEK_SET, whence);
  if (static_cast<uint64_t>(offset) >= length_in_bytes_)
    return -1;
  offset_ = offset;
  return offset_;
}

bool UDIFBlockChunkReadStream::CopyOutZeros(base::span<uint8_t> buf,
                                            size_t* bytes_read) {
  *bytes_read = std::min(buf.size(), length_in_bytes_ - offset_);
  UNSAFE_TODO(bzero(buf.data(), *bytes_read));
  offset_ += *bytes_read;
  return true;
}

bool UDIFBlockChunkReadStream::CopyOutUncompressed(base::span<uint8_t> buf,
                                                   size_t* bytes_read) {
  *bytes_read = std::min(buf.size(), length_in_bytes_ - offset_);

  if (*bytes_read == 0) {
    return true;
  }

  uint64_t offset = chunk_->compressed_offset + offset_;
  if (stream_->Seek(offset, SEEK_SET) == -1) {
    return false;
  }

  bool rv = stream_->Read(buf.first(*bytes_read), bytes_read);
  if (rv) {
    offset_ += *bytes_read;
  } else {
    DLOG(ERROR) << "Failed to read uncompressed chunk data";
  }
  return rv;
}

bool UDIFBlockChunkReadStream::CopyOutDecompressed(base::span<uint8_t> buf,
                                                   size_t* bytes_read) {
  DCHECK(did_decompress_);
  *bytes_read = std::min(buf.size(), decompress_buffer_.size() - offset_);
  base::span<uint8_t> src_data =
      base::span(decompress_buffer_).subspan(offset_, *bytes_read);
  buf.copy_prefix_from(src_data);
  offset_ += *bytes_read;
  return true;
}

bool UDIFBlockChunkReadStream::HandleADC(base::span<uint8_t> buf,
                                         size_t* bytes_read) {
  // TODO(rsesek): Implement ADC handling.
  NOTIMPLEMENTED();
  return false;
}

bool UDIFBlockChunkReadStream::HandleZLib(base::span<uint8_t> buf,
                                          size_t* bytes_read) {
  if (!did_decompress_) {
    auto compressed_data_or_error = ReadCompressedData();
    if (!compressed_data_or_error.has_value()) {
      return false;
    }
    std::vector<uint8_t>& compressed_data = compressed_data_or_error.value();

    z_stream zlib = {};
    if (inflateInit(&zlib) != Z_OK) {
      DLOG(ERROR) << "Failed to initialize zlib";
      return false;
    }

    decompress_buffer_.resize(length_in_bytes_);
    zlib.next_in = compressed_data.data();
    zlib.avail_in = compressed_data.size();
    zlib.next_out = decompress_buffer_.data();
    zlib.avail_out = decompress_buffer_.size();

    int rv = inflate(&zlib, Z_FINISH);
    inflateEnd(&zlib);

    if (rv != Z_STREAM_END) {
      DLOG(ERROR) << "Failed to decompress zlib data, error = " << rv;
      return false;
    }

    did_decompress_ = true;
  }

  return CopyOutDecompressed(buf, bytes_read);
}

bool UDIFBlockChunkReadStream::HandleBZ2(base::span<uint8_t> buf,
                                         size_t* bytes_read) {
  if (!did_decompress_) {
    auto compressed_data_or_error = ReadCompressedData();
    if (!compressed_data_or_error.has_value()) {
      return false;
    }
    std::vector<uint8_t>& compressed_data = compressed_data_or_error.value();

    bz_stream bz = {};
    if (BZ2_bzDecompressInit(&bz, 0, 0) != BZ_OK) {
      DLOG(ERROR) << "Failed to initialize bzlib";
      return false;
    }

    decompress_buffer_.resize(length_in_bytes_);
    bz.next_in = reinterpret_cast<char*>(compressed_data.data());
    bz.avail_in = compressed_data.size();
    bz.next_out = reinterpret_cast<char*>(decompress_buffer_.data());
    bz.avail_out = decompress_buffer_.size();

    int rv = BZ2_bzDecompress(&bz);
    BZ2_bzDecompressEnd(&bz);

    if (rv != BZ_STREAM_END) {
      DLOG(ERROR) << "Failed to decompress BZ2 data, error = " << rv;
      return false;
    }

    did_decompress_ = true;
  }

  return CopyOutDecompressed(buf, bytes_read);
}

std::optional<std::vector<uint8_t>>
UDIFBlockChunkReadStream::ReadCompressedData() {
  std::vector<uint8_t> data;
  data.resize(chunk_->compressed_length);

  if (stream_->Seek(chunk_->compressed_offset, SEEK_SET) == -1) {
    return std::nullopt;
  }

  if (!stream_->ReadExact(data)) {
    return std::nullopt;
  }
  return data;
}

}  // namespace

}  // namespace dmg
}  // namespace safe_browsing