chrome/utility/safe_browsing/mac/udif.cc

   1 // Copyright 2015 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "chrome/utility/safe_browsing/mac/udif.h"
   6
   7 #include <bzlib.h>
   8 #include <CoreFoundation/CoreFoundation.h>
   9 #include <libkern/OSByteOrder.h>
  10 #include <uuid/uuid.h>
  11
  12 #include <algorithm>
  13
  14 #include "base/logging.h"
  15 #include "base/mac/foundation_util.h"
  16 #include "base/mac/scoped_cftyperef.h"
  17 #include "base/numerics/safe_math.h"
  18 #include "base/strings/sys_string_conversions.h"
  19 #include "chrome/utility/safe_browsing/mac/convert_big_endian.h"
  20 #include "chrome/utility/safe_browsing/mac/read_stream.h"
  21 #include "third_party/zlib/zlib.h"
  22
  23 namespace safe_browsing {
  24 namespace dmg {
  25
  26 #pragma pack(push, 1)
  27
  28 // The following structures come from the analysis provided by Jonathan Levin
  29 // at <http://newosxbook.com/DMG.html>.
  30 //
  31 // Note that all fields are stored in big endian.
  32
  33 struct UDIFChecksum {
  34   uint32_t type;
  35   uint32_t size;
  36   uint32_t data[32];
  37 };
  38
  39 static void ConvertBigEndian(UDIFChecksum* checksum) {
  40   ConvertBigEndian(&checksum->type);
  41   ConvertBigEndian(&checksum->size);
  42   for (size_t i = 0; i < arraysize(checksum->data); ++i) {
  43     ConvertBigEndian(&checksum->data[i]);
  44   }
  45 }
  46
  47 // The trailer structure for a UDIF file.
  48 struct UDIFResourceFile {
  49   static const uint32_t kSignature = 'koly';
  50   static const uint32_t kVersion = 4;
  51
  52   uint32_t signature;
  53   uint32_t version;
  54   uint32_t header_size;  // Size of this structure.
  55   uint32_t flags;
  56   uint64_t running_data_fork_offset;
  57   uint64_t data_fork_offset;
  58   uint64_t data_fork_length;
  59   uint64_t rsrc_fork_offset;
  60   uint64_t rsrc_fork_length;
  61   uint32_t segment_number;
  62   uint32_t segment_count;
  63   uuid_t   segment_id;
  64
  65   UDIFChecksum data_checksum;
  66
  67   uint64_t plist_offset;  // Offset and length of the blkx plist.
  68   uint64_t plist_length;
  69
  70   uint8_t  reserved1[120];
  71
  72   UDIFChecksum master_checksum;
  73
  74   uint32_t image_variant;
  75   uint64_t sector_count;
  76
  77   uint32_t reserved2;
  78   uint32_t reserved3;
  79   uint32_t reserved4;
  80 };
  81
  82 static void ConvertBigEndian(uuid_t* uuid) {
  83   // UUID is never consulted, so do not swap.
  84 }
  85
  86 static void ConvertBigEndian(UDIFResourceFile* file) {
  87   ConvertBigEndian(&file->signature);
  88   ConvertBigEndian(&file->version);
  89   ConvertBigEndian(&file->flags);
  90   ConvertBigEndian(&file->header_size);
  91   ConvertBigEndian(&file->running_data_fork_offset);
  92   ConvertBigEndian(&file->data_fork_offset);
  93   ConvertBigEndian(&file->data_fork_length);
  94   ConvertBigEndian(&file->rsrc_fork_offset);
  95   ConvertBigEndian(&file->rsrc_fork_length);
  96   ConvertBigEndian(&file->segment_number);
  97   ConvertBigEndian(&file->segment_count);
  98   ConvertBigEndian(&file->segment_id);
  99   ConvertBigEndian(&file->data_checksum);
 100   ConvertBigEndian(&file->plist_offset);
 101   ConvertBigEndian(&file->plist_length);
 102   ConvertBigEndian(&file->master_checksum);
 103   ConvertBigEndian(&file->image_variant);
 104   ConvertBigEndian(&file->sector_count);
 105   // Reserved fields are skipped.
 106 }
 107
 108 struct UDIFBlockChunk {
 109   enum class Type : uint32_t {
 110     ZERO_FILL     = 0x00000000,
 111     UNCOMPRESSED  = 0x00000001,
 112     IGNORED       = 0x00000002,
 113     COMPRESS_ADC  = 0x80000004,
 114     COMPRESS_ZLIB = 0x80000005,
 115     COMPRESSS_BZ2 = 0x80000006,
 116     COMMENT       = 0x7ffffffe,
 117     LAST_BLOCK    = 0xffffffff,
 118   };
 119
 120   Type type;
 121   uint32_t comment;
 122   uint64_t start_sector;  // Logical chunk offset and length, in sectors.
 123   uint64_t sector_count;
 124   uint64_t compressed_offset;  // Compressed offset and length, in bytes.
 125   uint64_t compressed_length;
 126 };
 127
 128 static void ConvertBigEndian(UDIFBlockChunk* chunk) {
 129   ConvertBigEndian(reinterpret_cast<uint32_t*>(&chunk->type));
 130   ConvertBigEndian(&chunk->comment);
 131   ConvertBigEndian(&chunk->start_sector);
 132   ConvertBigEndian(&chunk->sector_count);
 133   ConvertBigEndian(&chunk->compressed_offset);
 134   ConvertBigEndian(&chunk->compressed_length);
 135 }
 136
 137 struct UDIFBlockData {
 138   static const uint32_t kSignature = 'mish';
 139   static const uint32_t kVersion = 1;
 140
 141   uint32_t signature;
 142   uint32_t version;
 143   uint64_t start_sector;  // Logical block offset and length, in sectors.
 144   uint64_t sector_count;
 145
 146   uint64_t data_offset;
 147   uint32_t buffers_needed;
 148   uint32_t block_descriptors;
 149
 150   uint32_t reserved1;
 151   uint32_t reserved2;
 152   uint32_t reserved3;
 153   uint32_t reserved4;
 154   uint32_t reserved5;
 155   uint32_t reserved6;
 156
 157   UDIFChecksum checksum;
 158
 159   uint32_t chunk_count;
 160   UDIFBlockChunk chunks[0];
 161 };
 162
 163 static void ConvertBigEndian(UDIFBlockData* block) {
 164   ConvertBigEndian(&block->signature);
 165   ConvertBigEndian(&block->version);
 166   ConvertBigEndian(&block->start_sector);
 167   ConvertBigEndian(&block->sector_count);
 168   ConvertBigEndian(&block->data_offset);
 169   ConvertBigEndian(&block->buffers_needed);
 170   ConvertBigEndian(&block->block_descriptors);
 171   // Reserved fields are skipped.
 172   ConvertBigEndian(&block->checksum);
 173   ConvertBigEndian(&block->chunk_count);
 174   // Note: This deliberately does not swap the chunks themselves.
 175 }
 176
 177 // UDIFBlock takes a raw, big-endian block data pointer and stores, in host
 178 // endian, the data for both the block and the chunk.
 179 class UDIFBlock {
 180  public:
 181   explicit UDIFBlock(const UDIFBlockData* block_data) : block(*block_data) {
 182     ConvertBigEndian(&block);
 183     for (uint32_t i = 0; i < block.chunk_count; ++i) {
 184       chunks.push_back(block_data->chunks[i]);
 185       ConvertBigEndian(&chunks[i]);
 186     }
 187   }
 188
 189   uint32_t signature() const { return block.signature; }
 190   uint32_t version() const { return block.version; }
 191   uint64_t start_sector() const { return block.start_sector; }
 192   uint64_t sector_count() const { return block.sector_count; }
 193   uint64_t chunk_count() const { return chunks.size(); }
 194
 195   const UDIFBlockChunk* chunk(uint32_t i) const {
 196     if (i >= chunk_count())
 197       return nullptr;
 198     return &chunks[i];
 199   }
 200
 201  private:
 202   UDIFBlockData block;
 203   std::vector<UDIFBlockChunk> chunks;
 204
 205   DISALLOW_COPY_AND_ASSIGN(UDIFBlock);
 206 };
 207
 208 #pragma pack(pop)
 209
 210 namespace {
 211
 212 const size_t kSectorSize = 512;
 213
 214 class UDIFBlockChunkReadStream;
 215
 216 // A UDIFPartitionReadStream virtualizes a partition's non-contiguous blocks
 217 // into a single stream.
 218 class UDIFPartitionReadStream : public ReadStream {
 219  public:
 220   UDIFPartitionReadStream(ReadStream* stream,
 221                           uint16_t block_size,
 222                           const UDIFBlock* partition_block);
 223   ~UDIFPartitionReadStream() override;
 224
 225   bool Read(uint8_t* buffer, size_t buffer_size, size_t* bytes_read) override;
 226   // Seek only supports SEEK_SET and SEEK_CUR.
 227   off_t Seek(off_t offset, int whence) override;
 228
 229  private:
 230   ReadStream* const stream_;  // The UDIF stream.
 231   const uint16_t block_size_;  // The UDIF block size.
 232   const UDIFBlock* const block_;  // The block for this partition.
 233   uint64_t current_chunk_;  // The current chunk number.
 234   // The current chunk stream.
 235   scoped_ptr<UDIFBlockChunkReadStream> chunk_stream_;
 236
 237   DISALLOW_COPY_AND_ASSIGN(UDIFPartitionReadStream);
 238 };
 239
 240 // A ReadStream for a single block chunk, which transparently handles
 241 // decompression.
 242 class UDIFBlockChunkReadStream : public ReadStream {
 243  public:
 244   UDIFBlockChunkReadStream(ReadStream* stream,
 245                            uint16_t block_size,
 246                            const UDIFBlockChunk* chunk);
 247   ~UDIFBlockChunkReadStream() override;
 248
 249   bool Read(uint8_t* buffer, size_t buffer_size, size_t* bytes_read) override;
 250   // Seek only supports SEEK_SET.
 251   off_t Seek(off_t offset, int whence) override;
 252
 253   bool IsAtEnd() { return offset_ >= length_in_bytes_; }
 254
 255   const UDIFBlockChunk* chunk() const { return chunk_; }
 256   size_t length_in_bytes() const { return length_in_bytes_; }
 257
 258  private:
 259   bool CopyOutZeros(uint8_t* buffer, size_t buffer_size, size_t* bytes_read);
 260   bool CopyOutUncompressed(
 261       uint8_t* buffer, size_t buffer_size, size_t* bytes_read);
 262   bool CopyOutDecompressed(
 263       uint8_t* buffer, size_t buffer_size, size_t* bytes_read);
 264   bool HandleADC(uint8_t* buffer, size_t buffer_size, size_t* bytes_read);
 265   bool HandleZLib(uint8_t* buffer, size_t buffer_size, size_t* bytes_read);
 266   bool HandleBZ2(uint8_t* buffer, size_t buffer_size, size_t* bytes_read);
 267
 268   // Reads from |stream_| |chunk_->compressed_length| bytes, starting at
 269   // |chunk_->compressed_offset| into |out_data|.
 270   bool ReadCompressedData(std::vector<uint8_t>* out_data);
 271
 272   ReadStream* const stream_;  // The UDIF stream.
 273   const UDIFBlockChunk* const chunk_;  // The chunk to be read.
 274   size_t length_in_bytes_;  // The decompressed length in bytes.
 275   size_t offset_;  // The offset into the decompressed buffer.
 276   std::vector<uint8_t> decompress_buffer_;  // Decompressed data buffer.
 277   bool did_decompress_;  // Whether or not the chunk has been decompressed.
 278
 279   DISALLOW_COPY_AND_ASSIGN(UDIFBlockChunkReadStream);
 280 };
 281
 282 }  // namespace
 283
 284 UDIFParser::UDIFParser(ReadStream* stream)
 285     : stream_(stream),
 286       partition_names_(),
 287       blocks_(),
 288       block_size_(kSectorSize) {
 289 }
 290
 291 UDIFParser::~UDIFParser() {}
 292
 293 bool UDIFParser::Parse() {
 294   if (!ParseBlkx())
 295     return false;
 296
 297   return true;
 298 }
 299
 300 size_t UDIFParser::GetNumberOfPartitions() {
 301   return blocks_.size();
 302 }
 303
 304 std::string UDIFParser::GetPartitionName(size_t part_number) {
 305   DCHECK_LT(part_number, partition_names_.size());
 306   return partition_names_[part_number];
 307 }
 308
 309 std::string UDIFParser::GetPartitionType(size_t part_number) {
 310   // The partition type is embedded in the Name field, as such:
 311   // "Partition-Name (Partition-Type : Partition-ID)".
 312   std::string name = GetPartitionName(part_number);
 313   size_t open = name.rfind('(');
 314   size_t separator = name.rfind(':');
 315   if (open == std::string::npos || separator == std::string::npos)
 316     return std::string();
 317
 318   // Name does not end in ')' or no space after ':'.
 319   if (*(name.end() - 1) != ')' ||
 320       (name.size() - separator < 2 || name[separator + 1] != ' ')) {
 321     return std::string();
 322   }
 323
 324   --separator;
 325   ++open;
 326   if (separator <= open)
 327     return std::string();
 328   return name.substr(open, separator - open);
 329 }
 330
 331 size_t UDIFParser::GetPartitionSize(size_t part_number) {
 332   DCHECK_LT(part_number, blocks_.size());
 333   auto size =
 334       base::CheckedNumeric<size_t>(blocks_[part_number]->sector_count()) *
 335       block_size_;
 336   return size.ValueOrDie();
 337 }
 338
 339 scoped_ptr<ReadStream> UDIFParser::GetPartitionReadStream(size_t part_number) {
 340   DCHECK_LT(part_number, blocks_.size());
 341   return make_scoped_ptr(
 342       new UDIFPartitionReadStream(stream_, block_size_, blocks_[part_number]));
 343 }
 344
 345 bool UDIFParser::ParseBlkx() {
 346   UDIFResourceFile trailer;
 347   if (stream_->Seek(-sizeof(trailer), SEEK_END) == -1)
 348     return false;
 349
 350   if (!stream_->ReadType(&trailer)) {
 351     DLOG(ERROR) << "Failed to read UDIFResourceFile";
 352     return false;
 353   }
 354   ConvertBigEndian(&trailer);
 355
 356   if (trailer.signature != trailer.kSignature) {
 357     DLOG(ERROR) << "blkx signature does not match, is 0x"
 358                 << std::hex << trailer.signature;
 359     return false;
 360   }
 361   if (trailer.version != trailer.kVersion) {
 362     DLOG(ERROR) << "blkx version does not match, is " << trailer.version;
 363     return false;
 364   }
 365
 366   std::vector<uint8_t> plist_bytes(trailer.plist_length, 0);
 367
 368   if (stream_->Seek(trailer.plist_offset, SEEK_SET) == -1)
 369     return false;
 370
 371   if (!stream_->ReadExact(&plist_bytes[0], trailer.plist_length)) {
 372     DLOG(ERROR) << "Failed to read blkx plist data";
 373     return false;
 374   }
 375
 376   base::ScopedCFTypeRef<CFDataRef> plist_data(
 377       CFDataCreateWithBytesNoCopy(kCFAllocatorDefault,
 378           &plist_bytes[0], plist_bytes.size(), kCFAllocatorNull));
 379   if (!plist_data) {
 380     DLOG(ERROR) << "Failed to create data from bytes";
 381     return false;
 382   }
 383
 384   CFErrorRef error = nullptr;
 385   base::ScopedCFTypeRef<CFDictionaryRef> plist(
 386       base::mac::CFCast<CFDictionaryRef>(
 387            CFPropertyListCreateWithData(kCFAllocatorDefault,
 388                                         plist_data,
 389                                         kCFPropertyListImmutable,
 390                                         nullptr,
 391                                         &error)),
 392       base::scoped_policy::RETAIN);
 393   base::ScopedCFTypeRef<CFErrorRef> error_ref(error);
 394   if (error) {
 395     base::ScopedCFTypeRef<CFStringRef> error_string(
 396         CFErrorCopyDescription(error));
 397     DLOG(ERROR) << "Failed to parse XML plist: "
 398                 << base::SysCFStringRefToUTF8(error_string);
 399     return false;
 400   }
 401
 402   if (!plist) {
 403     DLOG(ERROR) << "Plist is not a dictionary";
 404     return false;
 405   }
 406
 407   auto resource_fork = base::mac::GetValueFromDictionary<CFDictionaryRef>(
 408       plist.get(), CFSTR("resource-fork"));
 409   if (!resource_fork) {
 410     DLOG(ERROR) << "No resource-fork entry in plist";
 411     return false;
 412   }
 413
 414   auto blkx = base::mac::GetValueFromDictionary<CFArrayRef>(resource_fork,
 415                                                             CFSTR("blkx"));
 416   if (!blkx) {
 417     DLOG(ERROR) << "No blkx entry in resource-fork";
 418     return false;
 419   }
 420
 421   for (CFIndex i = 0; i < CFArrayGetCount(blkx); ++i) {
 422     auto block_dictionary =
 423         base::mac::CFCast<CFDictionaryRef>(CFArrayGetValueAtIndex(blkx, i));
 424     auto data = base::mac::GetValueFromDictionary<CFDataRef>(block_dictionary,
 425                                                              CFSTR("Data"));
 426     if (!data) {
 427       DLOG(ERROR) << "Skipping block " << i
 428                   << " because it has no Data section";
 429       continue;
 430     }
 431
 432     // Copy the block table out of the plist.
 433     auto block_data =
 434         reinterpret_cast<const UDIFBlockData*>(CFDataGetBytePtr(data));
 435     scoped_ptr<UDIFBlock> block(new UDIFBlock(block_data));
 436
 437     if (block->signature() != UDIFBlockData::kSignature) {
 438       DLOG(ERROR) << "Skipping block " << i << " because its signature does not"
 439                   << " match, is 0x" << std::hex << block->signature();
 440       continue;
 441     }
 442     if (block->version() != UDIFBlockData::kVersion) {
 443       DLOG(ERROR) << "Skipping block " << i << "because its version does not "
 444                   << "match, is " << block->version();
 445       continue;
 446     }
 447
 448     CFStringRef partition_name_cf = base::mac::CFCast<CFStringRef>(
 449         CFDictionaryGetValue(block_dictionary, CFSTR("Name")));
 450     if (!partition_name_cf) {
 451       DLOG(ERROR) << "Skipping block " << i << " because it has no name";
 452       continue;
 453     }
 454     std::string partition_name = base::SysCFStringRefToUTF8(partition_name_cf);
 455
 456     if (DLOG_IS_ON(INFO) && VLOG_IS_ON(1)) {
 457       DVLOG(1) << "Name: " << partition_name;
 458       DVLOG(1) << "StartSector = " << block->start_sector()
 459                << ", SectorCount = " << block->sector_count()
 460                << ", ChunkCount = " << block->chunk_count();
 461       for (uint32_t j = 0; j < block->chunk_count(); ++j) {
 462         const UDIFBlockChunk* chunk = block->chunk(j);
 463         DVLOG(1) << "Chunk#" << j
 464                  << " type = " << std::hex << static_cast<uint32_t>(chunk->type)
 465                  << ", StartSector = " << std::dec << chunk->start_sector
 466                  << ", SectorCount = " << chunk->sector_count
 467                  << ", CompressOffset = " << chunk->compressed_offset
 468                  << ", CompressLen = " << chunk->compressed_length;
 469       }
 470     }
 471
 472     blocks_.push_back(block.Pass());
 473     partition_names_.push_back(partition_name);
 474   }
 475
 476   return true;
 477 }
 478
 479 bool UDIFParser::ReadBlockChunk(const UDIFBlockChunk* chunk,
 480                                 std::vector<uint8_t>* decompressed_data) {
 481   UDIFBlockChunkReadStream chunk_read_stream(stream_, block_size_, chunk);
 482   decompressed_data->resize(chunk_read_stream.length_in_bytes());
 483   return chunk_read_stream.ReadExact(&(*decompressed_data)[0],
 484                                      decompressed_data->size());
 485 }
 486
 487 namespace {
 488
 489 UDIFPartitionReadStream::UDIFPartitionReadStream(
 490     ReadStream* stream,
 491     uint16_t block_size,
 492     const UDIFBlock* partition_block)
 493     : stream_(stream),
 494       block_size_(block_size),
 495       block_(partition_block),
 496       current_chunk_(0),
 497       chunk_stream_() {
 498 }
 499
 500 UDIFPartitionReadStream::~UDIFPartitionReadStream() {}
 501
 502 bool UDIFPartitionReadStream::Read(uint8_t* buffer,
 503                                    size_t buffer_size,
 504                                    size_t* bytes_read) {
 505   size_t buffer_space_remaining = buffer_size;
 506   *bytes_read = 0;
 507
 508   for (uint32_t i = current_chunk_; i < block_->chunk_count(); ++i) {
 509     const UDIFBlockChunk* chunk = block_->chunk(i);
 510
 511     // If this is the last block chunk, then the read is complete.
 512     if (chunk->type == UDIFBlockChunk::Type::LAST_BLOCK) {
 513       break;
 514     }
 515
 516     // If the buffer is full, then the read is complete.
 517     if (buffer_space_remaining == 0)
 518       break;
 519
 520     // A chunk stream may exist if the last read from this chunk was partial,
 521     // or if the stream was Seek()ed.
 522     if (!chunk_stream_) {
 523       chunk_stream_.reset(
 524           new UDIFBlockChunkReadStream(stream_, block_size_, chunk));
 525     }
 526     DCHECK_EQ(chunk, chunk_stream_->chunk());
 527
 528     size_t chunk_bytes_read = 0;
 529     if (!chunk_stream_->Read(&buffer[buffer_size - buffer_space_remaining],
 530                              buffer_space_remaining,
 531                              &chunk_bytes_read)) {
 532       DLOG(ERROR) << "Failed to read " << buffer_space_remaining << " bytes "
 533                   << "from chunk " << i;
 534       return false;
 535     }
 536     *bytes_read += chunk_bytes_read;
 537     buffer_space_remaining -= chunk_bytes_read;
 538
 539     if (chunk_stream_->IsAtEnd()) {
 540       chunk_stream_.reset();
 541       ++current_chunk_;
 542     }
 543   }
 544
 545   return true;
 546 }
 547
 548 off_t UDIFPartitionReadStream::Seek(off_t offset, int whence) {
 549   // Translate SEEK_END to SEEK_SET. SEEK_CUR is not currently supported.
 550   if (whence == SEEK_END) {
 551     base::CheckedNumeric<off_t> safe_offset(block_->sector_count());
 552     safe_offset *= block_size_;
 553     safe_offset += offset;
 554     if (!safe_offset.IsValid()) {
 555       DLOG(ERROR) << "Seek offset overflows";
 556       return -1;
 557     }
 558     offset = safe_offset.ValueOrDie();
 559   } else if (whence != SEEK_SET) {
 560     DCHECK_EQ(SEEK_SET, whence);
 561   }
 562
 563   uint64_t sector = offset / block_size_;
 564
 565   // Find the chunk for this sector.
 566   uint32_t chunk_number = 0;
 567   const UDIFBlockChunk* chunk = nullptr;
 568   for (uint32_t i = 0; i < block_->chunk_count(); ++i) {
 569     const UDIFBlockChunk* chunk_it = block_->chunk(i);
 570     // This assumes that all the chunks are ordered by sector.
 571     if (i != 0) {
 572       DLOG_IF(ERROR,
 573               chunk_it->start_sector < block_->chunk(i - 1)->start_sector)
 574           << "Chunks are not ordered by sector at chunk " << i
 575           << " , previous start_sector = "
 576           << block_->chunk(i - 1)->start_sector << ", current = "
 577           << chunk_it->start_sector;
 578     }
 579     if (sector >= chunk_it->start_sector) {
 580       chunk = chunk_it;
 581       chunk_number = i;
 582     } else {
 583       break;
 584     }
 585   }
 586   if (!chunk) {
 587     DLOG(ERROR) << "Failed to Seek to partition offset " << offset;
 588     return -1;
 589   }
 590
 591   // Compute the offset into the chunk.
 592   uint64_t offset_in_sector = offset % block_size_;
 593   uint64_t start_sector = sector - chunk->start_sector;
 594   base::CheckedNumeric<uint64_t> decompress_read_offset(start_sector);
 595   decompress_read_offset *= block_size_;
 596   decompress_read_offset += offset_in_sector;
 597
 598   if (!decompress_read_offset.IsValid()) {
 599     DLOG(ERROR) << "Partition decompress read offset overflows";
 600     return -1;
 601   }
 602
 603   if (!chunk_stream_ || chunk != chunk_stream_->chunk()) {
 604     chunk_stream_.reset(
 605         new UDIFBlockChunkReadStream(stream_, block_size_, chunk));
 606   }
 607   current_chunk_ = chunk_number;
 608   chunk_stream_->Seek(decompress_read_offset.ValueOrDie(), SEEK_SET);
 609
 610   return offset;
 611 }
 612
 613 UDIFBlockChunkReadStream::UDIFBlockChunkReadStream(ReadStream* stream,
 614                                                    uint16_t block_size,
 615                                                    const UDIFBlockChunk* chunk)
 616     : stream_(stream),
 617       chunk_(chunk),
 618       length_in_bytes_(chunk->sector_count * block_size),
 619       offset_(0),
 620       decompress_buffer_(),
 621       did_decompress_(false) {
 622   // Make sure the multiplication above did not overflow.
 623   CHECK(length_in_bytes_ == 0 || length_in_bytes_ >= block_size);
 624 }
 625
 626 UDIFBlockChunkReadStream::~UDIFBlockChunkReadStream() {
 627 }
 628
 629 bool UDIFBlockChunkReadStream::Read(uint8_t* buffer,
 630                                     size_t buffer_size,
 631                                     size_t* bytes_read) {
 632   switch (chunk_->type) {
 633     case UDIFBlockChunk::Type::ZERO_FILL:
 634     case UDIFBlockChunk::Type::IGNORED:
 635       return CopyOutZeros(buffer, buffer_size, bytes_read);
 636     case UDIFBlockChunk::Type::UNCOMPRESSED:
 637       return CopyOutUncompressed(buffer, buffer_size, bytes_read);
 638     case UDIFBlockChunk::Type::COMPRESS_ADC:
 639       return HandleADC(buffer, buffer_size, bytes_read);
 640     case UDIFBlockChunk::Type::COMPRESS_ZLIB:
 641       return HandleZLib(buffer, buffer_size, bytes_read);
 642     case UDIFBlockChunk::Type::COMPRESSS_BZ2:
 643       return HandleBZ2(buffer, buffer_size, bytes_read);
 644     case UDIFBlockChunk::Type::COMMENT:
 645       NOTREACHED();
 646       break;
 647     case UDIFBlockChunk::Type::LAST_BLOCK:
 648       *bytes_read = 0;
 649       return true;
 650   }
 651   return false;
 652 }
 653
 654 off_t UDIFBlockChunkReadStream::Seek(off_t offset, int whence) {
 655   DCHECK_EQ(SEEK_SET, whence);
 656   DCHECK_LT(static_cast<uint64_t>(offset), length_in_bytes_);
 657   offset_ = offset;
 658   return offset_;
 659 }
 660
 661 bool UDIFBlockChunkReadStream::CopyOutZeros(uint8_t* buffer,
 662                                             size_t buffer_size,
 663                                             size_t* bytes_read) {
 664   *bytes_read = std::min(buffer_size, length_in_bytes_ - offset_);
 665   bzero(buffer, *bytes_read);
 666   offset_ += *bytes_read;
 667   return true;
 668 }
 669
 670 bool UDIFBlockChunkReadStream::CopyOutUncompressed(uint8_t* buffer,
 671                                                    size_t buffer_size,
 672                                                    size_t* bytes_read) {
 673   *bytes_read = std::min(buffer_size, length_in_bytes_ - offset_);
 674
 675   if (*bytes_read == 0)
 676     return true;
 677
 678   uint64_t offset = chunk_->compressed_offset + offset_;
 679   if (stream_->Seek(offset, SEEK_SET) == -1)
 680     return false;
 681
 682   bool rv = stream_->Read(buffer, *bytes_read, bytes_read);
 683   if (rv)
 684     offset_ += *bytes_read;
 685   else
 686     DLOG(ERROR) << "Failed to read uncompressed chunk data";
 687   return rv;
 688 }
 689
 690 bool UDIFBlockChunkReadStream::CopyOutDecompressed(uint8_t* buffer,
 691                                                    size_t buffer_size,
 692                                                    size_t* bytes_read) {
 693   DCHECK(did_decompress_);
 694   *bytes_read = std::min(buffer_size, decompress_buffer_.size() - offset_);
 695   memcpy(buffer, &decompress_buffer_[offset_], *bytes_read);
 696   offset_ += *bytes_read;
 697   return true;
 698 }
 699
 700 bool UDIFBlockChunkReadStream::HandleADC(uint8_t* buffer,
 701                                          size_t buffer_size,
 702                                          size_t* bytes_read) {
 703   // TODO(rsesek): Implement ADC handling.
 704   NOTIMPLEMENTED();
 705   return false;
 706 }
 707
 708 bool UDIFBlockChunkReadStream::HandleZLib(uint8_t* buffer,
 709                                           size_t buffer_size,
 710                                           size_t* bytes_read) {
 711   if (!did_decompress_) {
 712     std::vector<uint8_t> compressed_data(chunk_->compressed_length, 0);
 713     if (!ReadCompressedData(&compressed_data))
 714       return false;
 715
 716     z_stream zlib = {};
 717     if (inflateInit(&zlib) != Z_OK) {
 718       DLOG(ERROR) << "Failed to initialize zlib";
 719       return false;
 720     }
 721
 722     decompress_buffer_.resize(length_in_bytes_);
 723     zlib.next_in = &compressed_data[0];
 724     zlib.avail_in = compressed_data.size();
 725     zlib.next_out = &decompress_buffer_[0];
 726     zlib.avail_out = decompress_buffer_.size();
 727
 728     int rv = inflate(&zlib, Z_FINISH);
 729     inflateEnd(&zlib);
 730
 731     if (rv != Z_STREAM_END) {
 732       DLOG(ERROR) << "Failed to decompress zlib data, error = " << rv;
 733       return false;
 734     }
 735
 736     did_decompress_ = true;
 737   }
 738
 739   return CopyOutDecompressed(buffer, buffer_size, bytes_read);
 740 }
 741
 742 bool UDIFBlockChunkReadStream::HandleBZ2(uint8_t* buffer,
 743                                           size_t buffer_size,
 744                                           size_t* bytes_read) {
 745   if (!did_decompress_) {
 746     std::vector<uint8_t> compressed_data(chunk_->compressed_length, 0);
 747     if (!ReadCompressedData(&compressed_data))
 748       return false;
 749
 750     bz_stream bz = {};
 751     if (BZ2_bzDecompressInit(&bz, 0, 0) != BZ_OK) {
 752       DLOG(ERROR) << "Failed to initialize bzlib";
 753       return false;
 754     }
 755
 756     decompress_buffer_.resize(length_in_bytes_);
 757     bz.next_in = reinterpret_cast<char*>(&compressed_data[0]);
 758     bz.avail_in = compressed_data.size();
 759     bz.next_out = reinterpret_cast<char*>(&decompress_buffer_[0]);
 760     bz.avail_out = decompress_buffer_.size();
 761
 762     int rv = BZ2_bzDecompress(&bz);
 763     BZ2_bzDecompressEnd(&bz);
 764
 765     if (rv != BZ_STREAM_END) {
 766       DLOG(ERROR) << "Failed to decompress BZ2 data, error = " << rv;
 767       return false;
 768     }
 769
 770     did_decompress_ = true;
 771   }
 772
 773   return CopyOutDecompressed(buffer, buffer_size, bytes_read);
 774 }
 775
 776 bool UDIFBlockChunkReadStream::ReadCompressedData(
 777     std::vector<uint8_t>* out_data) {
 778   DCHECK_EQ(chunk_->compressed_length, out_data->size());
 779
 780   if (stream_->Seek(chunk_->compressed_offset, SEEK_SET) == -1)
 781     return false;
 782
 783   if (!stream_->ReadExact(&(*out_data)[0], out_data->size())) {
 784     DLOG(ERROR) << "Failed to read chunk compressed data at "
 785                 << chunk_->compressed_offset;
 786     return false;
 787   }
 788   return true;
 789 }
 790
 791 }  // namespace
 792
 793 }  // namespace dmg
 794 }  // namespace safe_browsing