Coverage for /Syzygy/zap_timestamp/zap_timestamp.cc

CoverageLines executed / instrumented / missingexe / inst / missLanguageGroup
78.9%3674650.C++source

Line-by-line coverage:

   1    :  // Copyright 2012 Google Inc. All Rights Reserved.
   2    :  //
   3    :  // Licensed under the Apache License, Version 2.0 (the "License");
   4    :  // you may not use this file except in compliance with the License.
   5    :  // You may obtain a copy of the License at
   6    :  //
   7    :  //     http://www.apache.org/licenses/LICENSE-2.0
   8    :  //
   9    :  // Unless required by applicable law or agreed to in writing, software
  10    :  // distributed under the License is distributed on an "AS IS" BASIS,
  11    :  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12    :  // See the License for the specific language governing permissions and
  13    :  // limitations under the License.
  14    :  //
  15    :  // ZapTimestamps uses PEFile/ImageLayout/BlockGraph to represent a PE file in
  16    :  // memory, and TypedBlock to navigate through the PE structures of the file.
  17    :  // We don't use a full decomposition of the image, but rather only decompose
  18    :  // the PE headers and structures. As such, ZapTimetamps can be seen as a
  19    :  // lightweight decomposer. It would be better do this directly using the
  20    :  // internal intermediate representation formats of PEFileParser, but this
  21    :  // functionality would require some refactoring.
  22    :  //
  23    :  // Changes that are required to be made to the PE file are represented by an
  24    :  // address space, mapping replacement data to file offsets. This address-space
  25    :  // can then be simply 'stamped' on to the PE file to be modified.
  26    :  //
  27    :  // The matching PDB file is completely rewritten to guarantee that it is
  28    :  // canonical (as long as the underlying PdbWriter doesn't change). We load all
  29    :  // of the streams into memory, reach in and make local modifications, and
  30    :  // rewrite the entire file to disk.
  31    :  
  32    :  #include "syzygy/zap_timestamp/zap_timestamp.h"
  33    :  
  34    :  #include "base/at_exit.h"
  35    :  #include "base/bind.h"
  36    :  #include "base/command_line.h"
  37    :  #include "base/logging.h"
  38    :  #include "base/md5.h"
  39    :  #include "base/files/file_util.h"
  40    :  #include "base/files/scoped_temp_dir.h"
  41    :  #include "base/strings/stringprintf.h"
  42    :  #include "syzygy/block_graph/typed_block.h"
  43    :  #include "syzygy/core/file_util.h"
  44    :  #include "syzygy/pdb/pdb_byte_stream.h"
  45    :  #include "syzygy/pdb/pdb_constants.h"
  46    :  #include "syzygy/pdb/pdb_reader.h"
  47    :  #include "syzygy/pdb/pdb_util.h"
  48    :  #include "syzygy/pdb/pdb_writer.h"
  49    :  #include "syzygy/pe/find.h"
  50    :  #include "syzygy/pe/pdb_info.h"
  51    :  #include "syzygy/pe/pe_data.h"
  52    :  #include "syzygy/pe/pe_file_parser.h"
  53    :  #include "syzygy/pe/pe_file_writer.h"
  54    :  
  55    :  namespace zap_timestamp {
  56    :  
  57    :  namespace {
  58    :  
  59    :  using block_graph::BlockGraph;
  60    :  using block_graph::ConstTypedBlock;
  61    :  using core::FileOffsetAddress;
  62    :  using core::RelativeAddress;
  63    :  using pdb::PdbByteStream;
  64    :  using pdb::PdbFile;
  65    :  using pdb::PdbReader;
  66    :  using pdb::PdbStream;
  67    :  using pdb::PdbWriter;
  68    :  using pdb::WritablePdbStream;
  69    :  using pe::ImageLayout;
  70    :  using pe::PEFile;
  71    :  using pe::PEFileParser;
  72    :  
  73    :  typedef ConstTypedBlock<IMAGE_DEBUG_DIRECTORY> ImageDebugDirectory;
  74    :  typedef ConstTypedBlock<IMAGE_DOS_HEADER> DosHeader;
  75    :  typedef ConstTypedBlock<IMAGE_NT_HEADERS> NtHeaders;
  76    :  typedef ConstTypedBlock<pe::CvInfoPdb70> CvInfoPdb;
  77    :  typedef ZapTimestamp::PatchAddressSpace PatchAddressSpace;
  78    :  typedef ZapTimestamp::PatchData PatchData;
  79    :  
  80    :  // An intermediate reference type used to track references generated by
  81    :  // PEFileParser.
  82    :  struct IntermediateReference {
  83    :    BlockGraph::ReferenceType type;
  84    :    BlockGraph::Size size;
  85    :    RelativeAddress address;
  86    :  };
  87    :  
  88    :  // A map of intermediate references. This tracks references created by the
  89    :  // PEFileParser.
  90    :  typedef std::map<RelativeAddress, IntermediateReference>
  91    :      IntermediateReferenceMap;
  92    :  
  93    :  // Adds a reference to the given intermediate reference map. Used as a callback
  94    :  // to PEFileParser.
  95    :  bool AddReference(IntermediateReferenceMap* references,
  96    :                    RelativeAddress source,
  97    :                    BlockGraph::ReferenceType type,
  98    :                    BlockGraph::Size size,
  99  E :                    RelativeAddress destination) {
 100  E :    DCHECK(references != NULL);
 101    :  
 102  E :    IntermediateReference ref = {type, size, destination};
 103  E :    return references->insert(std::make_pair(source, ref)).second;
 104  E :  }
 105    :  
 106    :  // Returns the block and offset into the block associated with the given
 107    :  // address and size. Sets block pointer to NULL and offset to 0 if no block
 108    :  // was found for the address and size. Returns true if a block is found, false
 109    :  // otherwise.
 110    :  bool LookupBlockOffset(const ImageLayout& image_layout,
 111    :                         RelativeAddress address,
 112    :                         size_t size,
 113    :                         BlockGraph::Block** block,
 114  E :                         BlockGraph::Offset* offset) {
 115  E :    DCHECK(block != NULL);
 116  E :    DCHECK(offset != NULL);
 117    :  
 118  E :    *block = image_layout.blocks.GetContainingBlock(address, size);
 119  E :    if (*block == NULL) {
 120  E :      *offset = 0;
 121  E :      return false;
 122    :    }
 123    :  
 124  E :    *offset = address - (*block)->addr();
 125  E :    return true;
 126  E :  }
 127    :  
 128    :  // Performs a decomposition of the given PE file, only parsing out the PE
 129    :  // data blocks and references between them.
 130    :  bool MiniDecompose(const PEFile& pe_file,
 131    :                     ImageLayout* image_layout,
 132  E :                     BlockGraph::Block** dos_header_block) {
 133  E :    DCHECK(image_layout != NULL);
 134  E :    DCHECK(dos_header_block != NULL);
 135    :  
 136  E :    IntermediateReferenceMap references;
 137    :  
 138    :    PEFileParser::AddReferenceCallback add_reference =
 139  E :        base::Bind(&AddReference, base::Unretained(&references));
 140  E :    PEFileParser pe_file_parser(pe_file, &image_layout->blocks, add_reference);
 141    :  
 142  E :    PEFileParser::PEHeader pe_header;
 143  E :    if (!pe_file_parser.ParseImage(&pe_header)) {
 144  i :      LOG(ERROR) << "Failed to parse PE file: " << pe_file.path().value();
 145  i :      return false;
 146    :    }
 147    :  
 148  E :    if (!pe::CopyHeaderToImageLayout(pe_header.nt_headers, image_layout)) {
 149  i :      LOG(ERROR) << "Failed to copy NT headers to image layout.";
 150  i :      return false;
 151    :    }
 152    :  
 153    :    // Finalize the intermediate references. We only finalize those that are
 154    :    // within the closed set of blocks.
 155  E :    IntermediateReferenceMap::const_iterator ref_it = references.begin();
 156  E :    for (; ref_it != references.end(); ++ref_it) {
 157  E :      BlockGraph::Block* src_block = NULL;
 158  E :      BlockGraph::Offset src_offset = 0;
 159  E :      if (!LookupBlockOffset(*image_layout, ref_it->first, ref_it->second.size,
 160    :                             &src_block, &src_offset)) {
 161  E :        continue;
 162    :      }
 163    :  
 164  E :      BlockGraph::Block* dst_block = NULL;
 165  E :      BlockGraph::Offset dst_offset = 0;
 166  E :      if (!LookupBlockOffset(*image_layout, ref_it->second.address, 1, &dst_block,
 167    :                             &dst_offset)) {
 168  E :        continue;
 169    :      }
 170    :  
 171    :      // Make the final reference.
 172  E :      BlockGraph::Reference ref(ref_it->second.type, ref_it->second.size,
 173    :                                dst_block, dst_offset, dst_offset);
 174  E :      CHECK(src_block->SetReference(src_offset, ref));
 175  E :    }
 176    :  
 177  E :    *dos_header_block = pe_header.dos_header;
 178    :  
 179  E :    return true;
 180  E :  }
 181    :  
 182    :  // Marks the range of data at @p rel_addr and of size @p size as needing to be
 183    :  // changed. It will be replaced with the data in @p data, and marked with the
 184    :  // description @p name (for debugging purposes). The change is recorded in the
 185    :  // provided PatchAddressSpace @p file_addr_space in terms of file offsets.
 186    :  // This performs the necessary address space translations via @p pe_file and
 187    :  // ensures that the change does not conflict with any other required changes.
 188    :  bool MarkData(const PEFile& pe_file,
 189    :                RelativeAddress rel_addr,
 190    :                size_t size,
 191    :                const uint8_t* data,
 192    :                const base::StringPiece& name,
 193  E :                PatchAddressSpace* file_addr_space) {
 194  E :    DCHECK(file_addr_space);
 195    :  
 196  E :    FileOffsetAddress file_addr;
 197  E :    if (!pe_file.Translate(rel_addr, &file_addr)) {
 198  i :      LOG(ERROR) << "Failed to translate " << rel_addr << " to file offset.";
 199  i :      return false;
 200    :    }
 201    :  
 202  E :    if (!file_addr_space->Insert(PatchAddressSpace::Range(file_addr, size),
 203    :                                 PatchData(data, name))) {
 204  i :      LOG(ERROR) << "Failed to insert file range at " << file_addr
 205    :                 << " of length " << size << ".";
 206  i :      return false;
 207    :    }
 208    :  
 209  E :    return true;
 210  E :  }
 211    :  
 212    :  // Given a data directory of type T containing a member variable named
 213    :  // TimeDateStamp, this will mark the timestamp for changing to the value
 214    :  // provided in @p timestamp_data. The change will be recorded in the provided
 215    :  // PatchAddressSpace @p file_addr_space.
 216    :  template <typename T>
 217    :  bool MarkDataDirectoryTimestamps(const PEFile& pe_file,
 218    :                                   NtHeaders& nt_headers,
 219    :                                   size_t data_dir_index,
 220    :                                   const char* data_dir_name,
 221    :                                   const uint8_t* timestamp_data,
 222  E :                                   PatchAddressSpace* file_addr_space) {
 223  E :    DCHECK_GT(arraysize(nt_headers->OptionalHeader.DataDirectory),
 224    :              data_dir_index);
 225  E :    DCHECK(timestamp_data != NULL);
 226  E :    DCHECK(file_addr_space != NULL);
 227    :  
 228    :    // It is not an error if the debug directory doesn't exist.
 229  E :    const IMAGE_DATA_DIRECTORY& data_dir_info =
 230    :        nt_headers->OptionalHeader.DataDirectory[data_dir_index];
 231  E :    if (!nt_headers.HasReference(data_dir_info.VirtualAddress)) {
 232  i :      DCHECK_EQ(0u, data_dir_info.VirtualAddress);
 233  i :      LOG(INFO) << "PE file contains no data directory " << data_dir_index << ".";
 234  i :      return true;
 235    :    }
 236    :  
 237  E :    ConstTypedBlock<T> data_dir;
 238  E :    if (!nt_headers.Dereference(data_dir_info.VirtualAddress, &data_dir)) {
 239  i :      LOG(ERROR) << "Failed to dereference data directory " << data_dir_index
 240    :                 << ".";
 241  i :      return false;
 242    :    }
 243    :  
 244  E :    FileOffsetAddress data_dir_addr;
 245  E :    if (!pe_file.Translate(data_dir.block()->addr(), &data_dir_addr)) {
 246  i :      LOG(ERROR) << "Failed to locate data directory " << data_dir_index << ".";
 247  i :      return false;
 248    :    }
 249    :  
 250  E :    if (data_dir->TimeDateStamp == 0)
 251  E :      return true;
 252    :  
 253    :    FileOffsetAddress timestamp_addr =
 254  E :        data_dir_addr + data_dir.OffsetOf(data_dir->TimeDateStamp);
 255    :  
 256  E :    std::string name = base::StringPrintf("%s Timestamp", data_dir_name);
 257  E :    if (!file_addr_space->Insert(
 258    :            PatchAddressSpace::Range(timestamp_addr, sizeof(DWORD)),
 259    :            PatchData(timestamp_data, name))) {
 260  i :      LOG(ERROR) << "Failed to mark timestamp of data directory "
 261    :                 << data_dir_index << ".";
 262  i :      return false;
 263    :    }
 264    :  
 265  E :    return true;
 266  E :  }
 267    :  
 268  E :  bool Md5Consume(size_t bytes, FILE* file, base::MD5Context* context) {
 269  E :    char buffer[4096] = {0};
 270    :  
 271  E :    size_t cur = 0;
 272  E :    while (cur < bytes) {
 273  E :      size_t bytes_to_read = std::min(bytes - cur, sizeof(buffer));
 274  E :      size_t bytes_read = ::fread(buffer, 1, bytes_to_read, file);
 275  E :      if (bytes_read != bytes_to_read) {
 276  i :        LOG(ERROR) << "Error reading from file (got " << bytes_read
 277    :                   << ", expected " << bytes_to_read << ").";
 278  i :        return false;
 279    :      }
 280    :  
 281  E :      base::MD5Update(context, base::StringPiece(buffer, bytes_read));
 282  E :      cur += bytes_read;
 283  E :    }
 284  E :    DCHECK_EQ(cur, bytes);
 285    :  
 286  E :    return true;
 287  E :  }
 288    :  
 289    :  bool UpdateFileInPlace(const base::FilePath& path,
 290  E :                         const PatchAddressSpace& updates) {
 291  E :    LOG(INFO) << "Patching file: " << path.value();
 292    :  
 293  E :    base::ScopedFILE file(base::OpenFile(path, "rb+"));
 294  E :    if (file.get() == NULL) {
 295  i :      LOG(ERROR) << "Unable to open file for updating: " << path.value();
 296  i :      return false;
 297    :    }
 298    :  
 299  E :    PatchAddressSpace::const_iterator it = updates.begin();
 300  E :    for (; it != updates.end(); ++it) {
 301    :      // No data? Then nothing to update. This happens for the PE checksum, which
 302    :      // has a NULL data pointer. We update it later on in another pass.
 303  E :      if (it->second.data == NULL)
 304  E :        continue;
 305    :  
 306  E :      LOG(INFO) << "  Patching " << it->second.name << ", " << it->first.size()
 307    :                << " bytes at " << it->first.start();
 308    :  
 309    :      // Seek to the position to be updated.
 310  E :      if (::fseek(file.get(), it->first.start().value(), SEEK_SET) != 0) {
 311  i :        LOG(ERROR) << "Failed to seek to " << it->first.start()
 312    :                   << " of file: " << path.value();
 313  i :        return false;
 314    :      }
 315    :  
 316    :      // Write the updated data.
 317    :      size_t bytes_written =
 318  E :          ::fwrite(it->second.data, 1, it->first.size(), file.get());
 319  E :      if (bytes_written != it->first.size()) {
 320  i :        LOG(ERROR) << "Failed to write " << it->first.size() << " bytes to "
 321    :                   << "position " << it->first.start()
 322    :                   << " of file: " << path.value();
 323    :      }
 324  E :    }
 325    :  
 326  E :    LOG(INFO) << "Finished patching file: " << path.value();
 327  E :    file.reset();
 328    :  
 329  E :    return true;
 330  E :  }
 331    :  
 332    :  // Ensures that the stream with the given ID is writable, returning a scoped
 333    :  // pointer to it.
 334  E :  scoped_refptr<PdbStream> GetWritableStream(size_t index, PdbFile* pdb_file) {
 335  E :    DCHECK(pdb_file != NULL);
 336  E :    DCHECK_GT(pdb_file->StreamCount(), index);
 337    :  
 338  E :    scoped_refptr<PdbStream> reader = pdb_file->GetStream(index);
 339    :  
 340    :    // Try and get the writer. If it's not available, then replace the stream
 341    :    // with a byte stream, which is in-place writable.
 342  E :    scoped_refptr<WritablePdbStream> writer = reader->GetWritableStream();
 343  E :    if (writer.get() == NULL) {
 344  E :      scoped_refptr<PdbByteStream> byte_stream(new PdbByteStream());
 345  E :      byte_stream->Init(reader.get());
 346  E :      pdb_file->ReplaceStream(index, byte_stream.get());
 347  E :      reader = byte_stream;
 348  E :    }
 349    :  
 350  E :    return reader;
 351  E :  }
 352    :  
 353  E :  void OutputSummaryStats(base::FilePath& path) {
 354  E :    base::ScopedFILE file(base::OpenFile(path, "rb"));
 355  E :    if (file.get() == NULL) {
 356  i :      LOG(ERROR) << "Unable to open file for reading: " << path.value();
 357  i :      return;
 358    :    }
 359  E :    ::fseek(file.get(), 0, SEEK_END);
 360  E :    size_t file_size = ::ftell(file.get());
 361  E :    ::fseek(file.get(), 0, SEEK_SET);
 362    :  
 363    :    base::MD5Context md5_context;
 364  E :    base::MD5Init(&md5_context);
 365  E :    if (!Md5Consume(file_size, file.get(), &md5_context))
 366  i :      return;
 367    :  
 368    :    base::MD5Digest md5_digest;
 369  E :    base::MD5Final(&md5_digest, &md5_context);
 370  E :    std::string md5_string = base::MD5DigestToBase16(md5_digest);
 371    :  
 372  E :    LOG(INFO) << "Path: " << path.value();
 373  E :    LOG(INFO) << "  Size  : " << file_size;
 374  E :    LOG(INFO) << "  Digest: " << md5_string;
 375  E :  }
 376    :  
 377  E :  bool NormalizeDbiStream(DWORD pdb_age_data, PdbByteStream* dbi_stream) {
 378  E :    DCHECK(dbi_stream != NULL);
 379    :  
 380  E :    LOG(INFO) << "Updating PDB DBI stream.";
 381    :  
 382  E :    uint8_t* dbi_data = dbi_stream->data();
 383  E :    if (dbi_stream->length() < sizeof(pdb::DbiHeader)) {
 384  i :      LOG(ERROR) << "DBI stream too short.";
 385  i :      return false;
 386    :    }
 387  E :    pdb::DbiHeader* dbi_header = reinterpret_cast<pdb::DbiHeader*>(dbi_data);
 388    :  
 389    :    // Update the age in the DbiHeader as well. This needs to match pdb_age
 390    :    // in the PDB header.
 391  E :    dbi_header->age = pdb_age_data;
 392  E :    dbi_data += sizeof(*dbi_header);
 393    :  
 394    :    // Ensure that the module information is addressable.
 395  E :    if (dbi_stream->length() < dbi_header->gp_modi_size) {
 396  i :      LOG(ERROR) << "Invalid DBI header gp_modi_size.";
 397  i :      return false;
 398    :    }
 399    :  
 400    :    // Run over the module information.
 401    :    // TODO(chrisha): Use BufferWriter to do this. We need to update it to handle
 402    :    //     type casts and bounds checking.
 403  E :    uint8_t* module_info_end = dbi_data + dbi_header->gp_modi_size;
 404  E :    while (dbi_data < module_info_end) {
 405    :      pdb::DbiModuleInfoBase* module_info =
 406  E :          reinterpret_cast<pdb::DbiModuleInfoBase*>(dbi_data);
 407  E :      module_info->offsets = 0;
 408  E :      dbi_data += sizeof(*module_info);
 409    :  
 410    :      // Skip two NULL terminated strings after the module info.
 411  E :      while (*dbi_data != 0)
 412  E :        ++dbi_data;
 413  E :      ++dbi_data;
 414  E :      while (*dbi_data != 0)
 415  E :        ++dbi_data;
 416  E :      ++dbi_data;
 417    :  
 418    :      // Skip until we're at a multiple of 4 position.
 419  E :      size_t offset = dbi_data - dbi_stream->data();
 420  E :      offset = ((offset + 3) / 4) * 4;
 421  E :      dbi_data = dbi_stream->data() + offset;
 422  E :    }
 423    :  
 424    :    // Ensure that the section contributions are addressable.
 425  E :    size_t section_contrib_end_pos = dbi_header->gp_modi_size + sizeof(uint32_t) +
 426    :                                     dbi_header->section_contribution_size;
 427  E :    if (dbi_stream->length() < section_contrib_end_pos) {
 428  i :      LOG(ERROR) << "Invalid DBI header gp_modi_size.";
 429  i :      return false;
 430    :    }
 431    :  
 432    :    // Run over the section contributions.
 433  E :    dbi_data += sizeof(uint32_t);  // Skip the signature.
 434    :    uint8_t* section_contrib_end =
 435  E :        dbi_data + dbi_header->section_contribution_size;
 436  E :    while (dbi_data < section_contrib_end) {
 437    :      pdb::DbiSectionContrib* section_contrib =
 438  E :          reinterpret_cast<pdb::DbiSectionContrib*>(dbi_data);
 439  E :      section_contrib->pad1 = 0;
 440  E :      section_contrib->pad2 = 0;
 441  E :      dbi_data += sizeof(*section_contrib);
 442  E :    }
 443    :  
 444  E :    return true;
 445  E :  }
 446    :  
 447  E :  bool NormalizeSymbolRecordStream(PdbByteStream* stream) {
 448  E :    DCHECK(stream != NULL);
 449    :  
 450  E :    uint8_t* data = stream->data();
 451  E :    uint8_t* data_end = data + stream->length();
 452    :  
 453  E :    while (data < data_end) {
 454    :      // Get the size of the symbol record and skip past it.
 455  E :      uint16_t* size = reinterpret_cast<uint16_t*>(data);
 456  E :      data += sizeof(*size);
 457    :  
 458    :      // The size of the symbol record, plus its uint16_t length, must be a
 459    :      // multiple
 460    :      // of 4. Each symbol record consists of the length followed by a symbol
 461    :      // type (also a short), so the size needs to be at least of length 2.
 462    :      // See http://code.google.com/p/syzygy/wiki/PdbFileFormat for a discussion
 463    :      // of the format of this stream.
 464  E :      DCHECK_LE(2u, *size);
 465  E :      DCHECK_EQ(0u, ((*size + sizeof(*size)) % 4));
 466    :  
 467    :      // Up to the last 3 bytes are padding, as the record gets rounded up to
 468    :      // a multiple of 4 in size.
 469    :      static const size_t kMaxPadding = 3;
 470  E :      uint8_t* end = data + *size;
 471  E :      uint8_t* tail = end - kMaxPadding;
 472    :  
 473    :      // Skip past the symbol record.
 474  E :      data = end;
 475    :  
 476    :      // Find the null terminator for the record.
 477  E :      for (; tail + 1 < end && *tail != 0; ++tail) {
 478    :        // Intentionally empty.
 479  E :      }
 480    :  
 481    :      // Pad out the rest of the record with nulls (these are usually full of
 482    :      // junk bytes).
 483  E :      for (; tail < end; ++tail)
 484  E :        *tail = 0;
 485  E :    }
 486    :  
 487  E :    return true;
 488  E :  }
 489    :  
 490    :  }  // namespace
 491    :  
 492    :  ZapTimestamp::ZapTimestamp()
 493  E :      : image_layout_(&block_graph_),
 494  E :        dos_header_block_(NULL),
 495  E :        write_image_(true),
 496  E :        write_pdb_(true),
 497  E :        overwrite_(false) {
 498    :    // The timestamp can't just be set to zero as that represents a special
 499    :    // value in the PE file. We set it to some arbitrary fixed date in the past.
 500    :    // This is Jan 1, 2010, 0:00:00 GMT. This date shouldn't be too much in
 501    :    // the past, otherwise Windows might trigger a warning saying that the
 502    :    // instrumented image has known incompatibility issues when someone tries to
 503    :    // run it.
 504  E :    timestamp_data_ = 1262304000;
 505    :  
 506    :    // Initialize the age to 1.
 507  E :    pdb_age_data_ = 1;
 508  E :  }
 509    :  
 510  E :  bool ZapTimestamp::Init() {
 511  E :    if (!ValidatePeAndPdbFiles())
 512  E :      return false;
 513    :  
 514  E :    if (!ValidateOutputPaths())
 515  E :      return false;
 516    :  
 517  E :    if (!DecomposePeFile())
 518  i :      return false;
 519    :  
 520  E :    if (!MarkPeFileRanges())
 521  i :      return false;
 522    :  
 523  E :    if (!input_pdb_.empty()) {
 524  E :      if (!CalculatePdbGuid())
 525  i :        return false;
 526    :  
 527  E :      if (!LoadAndUpdatePdbFile())
 528  i :        return false;
 529    :    }
 530    :  
 531  E :    return true;
 532  E :  }
 533    :  
 534  E :  bool ZapTimestamp::Zap() {
 535  E :    if (write_image_) {
 536  E :      if (!WritePeFile())
 537  i :        return false;
 538  E :      OutputSummaryStats(input_image_);
 539    :    }
 540    :  
 541  E :    if (!input_pdb_.empty() && write_pdb_) {
 542  E :      if (!WritePdbFile())
 543  i :        return false;
 544  E :      OutputSummaryStats(input_pdb_);
 545    :    }
 546    :  
 547  E :    return true;
 548  E :  }
 549    :  
 550  E :  bool ZapTimestamp::ValidatePeAndPdbFiles() {
 551  E :    LOG(INFO) << "Analyzing PE file: " << input_image_.value();
 552    :  
 553  E :    if (!base::PathExists(input_image_) || base::DirectoryExists(input_image_)) {
 554  E :      LOG(ERROR) << "PE file not found: " << input_image_.value();
 555  E :      return false;
 556    :    }
 557    :  
 558  E :    if (!pe_file_.Init(input_image_)) {
 559  i :      LOG(ERROR) << "Failed to read PE file: " << input_image_.value();
 560  i :      return false;
 561    :    }
 562    :  
 563  E :    if (input_pdb_.empty()) {
 564    :      // If the image has no code view entry (ie: no matching PDB file)
 565    :      // then accept this fact and leave the PDB path empty.
 566  E :      pe::PdbInfo pe_pdb_info;
 567  E :      if (!pe_pdb_info.Init(input_image_))
 568  E :        return true;
 569    :  
 570    :      // Find the matching PDB file.
 571  E :      if (!pe::FindPdbForModule(input_image_, &input_pdb_)) {
 572  i :        LOG(ERROR) << "Error while searching for PDB file.";
 573  i :        return false;
 574    :      }
 575  E :      if (input_pdb_.empty()) {
 576  E :        LOG(ERROR) << "PDB file not found for PE file: " << input_image_.value();
 577  E :        return false;
 578    :      }
 579  E :      DCHECK(base::PathExists(input_pdb_));
 580  E :    } else {
 581  E :      if (!base::PathExists(input_pdb_) || base::DirectoryExists(input_pdb_)) {
 582  i :        LOG(ERROR) << "PDB file not found: " << input_pdb_.value();
 583    :      }
 584    :    }
 585    :  
 586    :    // Ensure that the PDB and the PE file are consistent with each other.
 587  E :    if (!pe::PeAndPdbAreMatched(input_image_, input_pdb_))
 588  i :      return false;  // This logs verbosely.
 589    :  
 590  E :    LOG(INFO) << "Found matching PDB file: " << input_pdb_.value();
 591    :  
 592  E :    return true;
 593  E :  }
 594    :  
 595  E :  bool ZapTimestamp::ValidateOutputPaths() {
 596  E :    if (output_image_.empty())
 597  E :      output_image_ = input_image_;
 598    :  
 599  E :    if (input_pdb_.empty()) {
 600  E :      if (!output_pdb_.empty()) {
 601  i :        LOG(INFO) << "Ignoring output-pdb path: " << output_pdb_.value();
 602  i :        output_pdb_.clear();
 603    :      }
 604  E :    } else {
 605  E :      if (output_pdb_.empty()) {
 606  E :        if (input_image_.BaseName() == output_image_.BaseName()) {
 607    :          // The input and output have the same basename. Use the input PDB
 608    :          // basename, but place it alongside the output image.
 609  E :          output_pdb_ = output_image_.DirName().Append(input_pdb_.BaseName());
 610  E :        } else {
 611    :          // The basenames don't match. Simpy append ".pdb" to the output
 612    :          // image.
 613  E :          output_pdb_ = output_image_.AddExtension(L"pdb");
 614    :        }
 615    :      }
 616    :    }
 617    :  
 618    :    // If overwriting isn't allowed then double check everything is kosher.
 619  E :    if (!overwrite_) {
 620  E :      if (write_image_ && (base::PathExists(output_image_) ||
 621    :                           core::CompareFilePaths(input_image_, output_image_) ==
 622    :                               core::kEquivalentFilePaths)) {
 623  E :        LOG(ERROR) << "Output image file exists. Must enable overwrite.";
 624  E :        return false;
 625    :      }
 626  E :      if (write_pdb_ && !output_pdb_.empty() &&
 627    :          (base::PathExists(output_pdb_) ||
 628    :           core::CompareFilePaths(input_pdb_, output_pdb_) ==
 629    :               core::kEquivalentFilePaths)) {
 630  E :        LOG(ERROR) << "Output PDB file exists. Must enable overwrite.";
 631  E :        return false;
 632    :      }
 633    :    }
 634    :  
 635  E :    return true;
 636  E :  }
 637    :  
 638  E :  bool ZapTimestamp::DecomposePeFile() {
 639    :    // Decompose the image. This is a very high level decomposition only
 640    :    // chunking out the PE structures and references from/to PE blocks.
 641  E :    if (!MiniDecompose(pe_file_, &image_layout_, &dos_header_block_))
 642  i :      return false;
 643    :  
 644  E :    return true;
 645  E :  }
 646    :  
 647  E :  bool ZapTimestamp::MarkPeFileRanges() {
 648  E :    DCHECK(dos_header_block_ != NULL);
 649  E :    LOG(INFO) << "Finding PE fields that need updating.";
 650    :  
 651  E :    DosHeader dos_header;
 652  E :    if (!dos_header.Init(0, dos_header_block_)) {
 653  i :      LOG(ERROR) << "Failed to cast IMAGE_DOS_HEADER.";
 654  i :      return false;
 655    :    }
 656    :  
 657  E :    NtHeaders nt_headers;
 658  E :    if (!dos_header.Dereference(dos_header->e_lfanew, &nt_headers)) {
 659  i :      LOG(ERROR) << "Failed to dereference IMAGE_NT_HEADERS.";
 660  i :      return false;
 661    :    }
 662    :  
 663    :    // Mark the export data directory timestamp.
 664  E :    if (!MarkDataDirectoryTimestamps<IMAGE_EXPORT_DIRECTORY>(
 665    :            pe_file_, nt_headers, IMAGE_DIRECTORY_ENTRY_EXPORT,
 666    :            "Export Directory",
 667    :            reinterpret_cast<const uint8_t*>(&timestamp_data_),
 668    :            &pe_file_addr_space_)) {
 669    :      // This logs verbosely on failure.
 670  i :      return false;
 671    :    }
 672    :  
 673    :    // Mark the resource data directory timestamp.
 674  E :    if (!MarkDataDirectoryTimestamps<IMAGE_RESOURCE_DIRECTORY>(
 675    :            pe_file_, nt_headers, IMAGE_DIRECTORY_ENTRY_RESOURCE,
 676    :            "Resource Directory",
 677    :            reinterpret_cast<const uint8_t*>(&timestamp_data_),
 678    :            &pe_file_addr_space_)) {
 679    :      // This logs verbosely on failure.
 680  i :      return false;
 681    :    }
 682    :  
 683    :    // Find the debug directory.
 684  E :    ImageDebugDirectory debug_dir;
 685  E :    const IMAGE_DATA_DIRECTORY& debug_dir_info =
 686    :        nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG];
 687  E :    if (nt_headers.HasReference(debug_dir_info.VirtualAddress))
 688  E :      nt_headers.Dereference(debug_dir_info.VirtualAddress, &debug_dir);
 689    :  
 690    :    // Within that, find the codeview debug entry. We also update every other
 691    :    // debug timestamp.
 692  E :    CvInfoPdb cv_info_pdb;
 693  E :    RelativeAddress rel_addr;
 694  E :    if (debug_dir.block()) {
 695  E :      for (size_t i = 0; i < debug_dir.ElementCount(); ++i) {
 696  E :        rel_addr = debug_dir.block()->addr() +
 697    :                   debug_dir.OffsetOf(debug_dir[i].TimeDateStamp);
 698  E :        std::string name = base::StringPrintf("Debug Directory %d Timestamp", i);
 699  E :        if (!MarkData(pe_file_, rel_addr, sizeof(timestamp_data_),
 700    :                      reinterpret_cast<const uint8_t*>(&timestamp_data_), name,
 701    :                      &pe_file_addr_space_)) {
 702  i :          LOG(ERROR) << "Failed to mark TimeDateStamp of debug directory " << i
 703    :                     << ".";
 704  i :          return false;
 705    :        }
 706    :  
 707  E :        if (debug_dir[i].Type == IMAGE_DEBUG_TYPE_CODEVIEW) {
 708  E :          if (cv_info_pdb.block() != NULL) {
 709  i :            LOG(ERROR) << "Found multiple CodeView debug directories.";
 710  i :            return false;
 711    :          }
 712  E :          if (!debug_dir.Dereference(debug_dir[i].PointerToRawData,
 713    :                                     &cv_info_pdb)) {
 714  i :            LOG(ERROR) << "Failed to dereference CodeView debug directory.";
 715  i :            return false;
 716    :          }
 717    :        }
 718  E :      }
 719    :    }
 720    :  
 721    :    // We should have found a code view debug directory pointing to the PDB file.
 722  E :    if (!input_pdb_.empty()) {
 723  E :      if (cv_info_pdb.block() == NULL) {
 724  i :        LOG(ERROR) << "Failed to find CodeView debug directory.";
 725  i :        return false;
 726    :      }
 727    :  
 728    :      // Get the file offset of the PDB age and mark it.
 729  E :      rel_addr = cv_info_pdb.block()->addr() +
 730    :                 cv_info_pdb.OffsetOf(cv_info_pdb->pdb_age);
 731  E :      if (!MarkData(pe_file_, rel_addr, sizeof(pdb_age_data_),
 732    :                    reinterpret_cast<const uint8_t*>(&pdb_age_data_), "PDB Age",
 733    :                    &pe_file_addr_space_)) {
 734  i :        LOG(ERROR) << "Failed to mark PDB age.";
 735  i :        return false;
 736    :      }
 737    :  
 738    :      // Get the file offset of the PDB guid and mark it.
 739  E :      rel_addr = cv_info_pdb.block()->addr() +
 740    :                 cv_info_pdb.OffsetOf(cv_info_pdb->signature);
 741  E :      if (!MarkData(pe_file_, rel_addr, sizeof(pdb_guid_data_),
 742    :                    reinterpret_cast<const uint8_t*>(&pdb_guid_data_), "PDB GUID",
 743    :                    &pe_file_addr_space_)) {
 744  i :        LOG(ERROR) << "Failed to mark PDB GUID.";
 745  i :        return false;
 746    :      }
 747    :    }
 748    :  
 749    :    // Get the file offset of the PE checksum and mark it.
 750  E :    rel_addr = nt_headers.block()->addr() +
 751    :               nt_headers.OffsetOf(nt_headers->OptionalHeader.CheckSum);
 752  E :    if (!MarkData(pe_file_, rel_addr, sizeof(DWORD), NULL, "PE Checksum",
 753    :                  &pe_file_addr_space_)) {
 754  i :      LOG(ERROR) << "Failed to mark PE checksum.";
 755  i :      return false;
 756    :    }
 757    :  
 758    :    // Get the file offset of the PE timestamp and mark it.
 759  E :    rel_addr = nt_headers.block()->addr() +
 760    :               nt_headers.OffsetOf(nt_headers->FileHeader.TimeDateStamp);
 761  E :    if (!MarkData(pe_file_, rel_addr, sizeof(timestamp_data_),
 762    :                  reinterpret_cast<uint8_t*>(&timestamp_data_), "PE Timestamp",
 763    :                  &pe_file_addr_space_)) {
 764  i :      LOG(ERROR) << "Failed to mark PE timestamp.";
 765  i :      return false;
 766    :    }
 767    :  
 768  E :    return true;
 769  E :  }
 770    :  
 771  E :  bool ZapTimestamp::CalculatePdbGuid() {
 772  E :    DCHECK(!input_pdb_.empty());
 773    :  
 774  E :    LOG(INFO) << "Calculating PDB GUID from PE file contents.";
 775    :  
 776  E :    base::ScopedFILE pe_file(base::OpenFile(input_image_, "rb"));
 777  E :    if (pe_file.get() == NULL) {
 778  i :      LOG(ERROR) << "Failed to open PE file for reading: "
 779    :                 << input_image_.value();
 780  i :      return false;
 781    :    }
 782    :  
 783    :    // Get the length of the entire file.
 784  E :    if (::fseek(pe_file.get(), 0, SEEK_END) != 0) {
 785  i :      LOG(ERROR) << "Failed to fseek to end of file.";
 786  i :      return false;
 787    :    }
 788  E :    FileOffsetAddress end(::ftell(pe_file.get()));
 789    :  
 790    :    // Seek back to the beginning.
 791  E :    if (::fseek(pe_file.get(), 0, SEEK_SET) != 0) {
 792  i :      LOG(ERROR) << "Failed to fseek to beginning of file.";
 793  i :      return false;
 794    :    }
 795    :  
 796    :    // Initialize the MD5 structure.
 797  E :    base::MD5Context md5_context = {0};
 798  E :    base::MD5Init(&md5_context);
 799    :  
 800    :    // We seek through the bits of the file that will be changed, and skip those.
 801    :    // The rest of the file (the static parts) are fed through an MD5 hash and
 802    :    // used to generated a unique and stable GUID.
 803  E :    FileOffsetAddress cur(0);
 804  E :    PatchAddressSpace::const_iterator range_it = pe_file_addr_space_.begin();
 805  E :    for (; range_it != pe_file_addr_space_.end(); ++range_it) {
 806    :      // Consume any data before this range.
 807  E :      if (cur < range_it->first.start()) {
 808  E :        size_t bytes_to_hash = range_it->first.start() - cur;
 809  E :        if (!Md5Consume(bytes_to_hash, pe_file.get(), &md5_context))
 810  i :          return false;  // This logs verbosely for us.
 811    :      }
 812    :  
 813  E :      if (::fseek(pe_file.get(), range_it->first.size(), SEEK_CUR)) {
 814  i :        LOG(ERROR) << "Failed to fseek past marked range.";
 815    :      }
 816    :  
 817  E :      cur = range_it->first.end();
 818  E :    }
 819    :  
 820    :    // Consume any left-over data.
 821  E :    if (cur < end) {
 822  E :      if (!Md5Consume(end - cur, pe_file.get(), &md5_context))
 823  i :        return false;  // This logs verbosely for us.
 824    :    }
 825    :  
 826  E :    DCHECK_EQ(end.value(), static_cast<uint32_t>(::ftell(pe_file.get())));
 827    :  
 828    :    static_assert(sizeof(base::MD5Digest) == sizeof(pdb_guid_data_),
 829    :                  "MD5Digest and GUID size mismatch.");
 830  E :    base::MD5Final(reinterpret_cast<base::MD5Digest*>(&pdb_guid_data_),
 831    :                   &md5_context);
 832  E :    LOG(INFO) << "Final GUID is "
 833    :              << base::MD5DigestToBase16(
 834    :                     *reinterpret_cast<base::MD5Digest*>(&pdb_guid_data_))
 835    :              << ".";
 836    :  
 837  E :    return true;
 838  E :  }
 839    :  
 840  E :  bool ZapTimestamp::LoadAndUpdatePdbFile() {
 841  E :    DCHECK(!input_pdb_.empty());
 842  E :    DCHECK(pdb_file_.get() == NULL);
 843    :  
 844  E :    pdb_file_.reset(new PdbFile());
 845  E :    PdbReader pdb_reader;
 846  E :    if (!pdb_reader.Read(input_pdb_, pdb_file_.get())) {
 847  i :      LOG(ERROR) << "Failed to read PDB file: " << input_pdb_.value();
 848  i :      return false;
 849    :    }
 850    :  
 851    :    // We turf the old directory stream as a fresh PDB does not have one. It's
 852    :    // also meaningless after we rewrite a PDB as the old blocks it refers to
 853    :    // will no longer exist.
 854  E :    pdb_file_->ReplaceStream(pdb::kPdbOldDirectoryStream, NULL);
 855    :  
 856    :    scoped_refptr<PdbStream> header_reader =
 857  E :        GetWritableStream(pdb::kPdbHeaderInfoStream, pdb_file_.get());
 858  E :    if (header_reader.get() == NULL) {
 859  i :      LOG(ERROR) << "No header info stream in PDB file: " << input_pdb_.value();
 860  i :      return false;
 861    :    }
 862    :  
 863    :    scoped_refptr<WritablePdbStream> header_writer =
 864  E :        header_reader->GetWritableStream();
 865  E :    DCHECK(header_writer.get() != NULL);
 866    :  
 867    :    // Update the timestamp, the age and the signature.
 868  E :    LOG(INFO) << "Updating PDB header.";
 869  E :    header_writer->set_pos(offsetof(pdb::PdbInfoHeader70, timestamp));
 870  E :    header_writer->Write(static_cast<uint32_t>(timestamp_data_));
 871  E :    header_writer->Write(static_cast<uint32_t>(pdb_age_data_));
 872  E :    header_writer->Write(pdb_guid_data_);
 873    :  
 874    :    // Normalize the DBI stream in place.
 875  E :    scoped_refptr<PdbByteStream> dbi_stream(new PdbByteStream());
 876  E :    CHECK(dbi_stream->Init(pdb_file_->GetStream(pdb::kDbiStream).get()));
 877  E :    pdb_file_->ReplaceStream(pdb::kDbiStream, dbi_stream.get());
 878  E :    if (!NormalizeDbiStream(pdb_age_data_, dbi_stream.get())) {
 879  i :      LOG(ERROR) << "Failed to normalize DBI stream.";
 880  i :      return false;
 881    :    }
 882    :  
 883  E :    uint8_t* dbi_data = dbi_stream->data();
 884  E :    pdb::DbiHeader* dbi_header = reinterpret_cast<pdb::DbiHeader*>(dbi_data);
 885    :  
 886    :    // Normalize the symbol record stream in place.
 887  E :    scoped_refptr<PdbByteStream> symrec_stream(new PdbByteStream());
 888  E :    CHECK(symrec_stream->Init(
 889    :        pdb_file_->GetStream(dbi_header->symbol_record_stream).get()));
 890  E :    pdb_file_->ReplaceStream(dbi_header->symbol_record_stream,
 891    :                             symrec_stream.get());
 892  E :    if (!NormalizeSymbolRecordStream(symrec_stream.get())) {
 893  i :      LOG(ERROR) << "Failed to normalize symbol record stream.";
 894  i :      return false;
 895    :    }
 896    :  
 897    :    // Normalize the public symbol info stream. There's a DWORD of padding at
 898    :    // offset 24 that we want to zero.
 899    :    scoped_refptr<PdbStream> pubsym_reader =
 900  E :        GetWritableStream(dbi_header->public_symbol_info_stream, pdb_file_.get());
 901    :    scoped_refptr<WritablePdbStream> pubsym_writer =
 902  E :        pubsym_reader->GetWritableStream();
 903  E :    DCHECK(pubsym_writer.get() != NULL);
 904  E :    pubsym_writer->set_pos(24);
 905  E :    pubsym_writer->Write(static_cast<uint32_t>(0));
 906    :  
 907  E :    return true;
 908  E :  }
 909    :  
 910  E :  bool ZapTimestamp::WritePeFile() {
 911  E :    if (core::CompareFilePaths(input_image_, output_image_) !=
 912    :        core::kEquivalentFilePaths) {
 913    :      if (::CopyFileW(input_image_.value().c_str(), output_image_.value().c_str(),
 914  E :                      FALSE) == FALSE) {
 915  i :        LOG(ERROR) << "Failed to write output image: %s" << output_image_.value();
 916  i :        return false;
 917    :      }
 918    :    }
 919    :  
 920  E :    if (!UpdateFileInPlace(output_image_, pe_file_addr_space_))
 921  i :      return false;
 922    :  
 923  E :    LOG(INFO) << "Updating checksum for PE file: " << output_image_.value();
 924  E :    if (!pe::PEFileWriter::UpdateFileChecksum(output_image_)) {
 925  i :      LOG(ERROR) << "Failed to update checksum for PE file: "
 926    :                 << output_image_.value();
 927  i :      return false;
 928    :    }
 929    :  
 930  E :    return true;
 931  E :  }
 932    :  
 933  E :  bool ZapTimestamp::WritePdbFile() {
 934  E :    DCHECK(!input_pdb_.empty());
 935    :  
 936    :    // We actually completely rewrite the PDB file to a temporary location, and
 937    :    // then move it over top of the existing one. This is because pdb_file_
 938    :    // actually has an open file handle to the original PDB.
 939    :  
 940    :    // We create a temporary directory alongside the final destination so as
 941    :    // not to cross volume boundaries.
 942  E :    base::FilePath output_dir = output_pdb_.DirName();
 943  E :    base::ScopedTempDir temp_dir;
 944  E :    if (!temp_dir.CreateUniqueTempDirUnderPath(output_dir)) {
 945  i :      LOG(ERROR) << "Failed to create temporary directory in \""
 946    :                 << output_dir.value() << "\".";
 947  i :      return false;
 948    :    }
 949    :  
 950    :    // Generate the path to the rewritten PDB.
 951  E :    base::FilePath temp_path = temp_dir.path().Append(input_pdb_.BaseName());
 952    :  
 953  E :    PdbWriter pdb_writer;
 954  E :    LOG(INFO) << "Creating temporary PDB file: " << temp_path.value();
 955  E :    if (!pdb_writer.Write(temp_path, *pdb_file_.get())) {
 956  i :      LOG(ERROR) << "Failed to write new PDB: " << temp_path.value();
 957  i :      return false;
 958    :    }
 959    :  
 960    :    // Free up the PDB file. This will close the open file handle to the original
 961    :    // PDB file.
 962  E :    pdb_file_.reset(NULL);
 963    :  
 964    :    // Copy over top of the original file.
 965  E :    LOG(INFO) << "Temporary PDB file replacing destination PDB: "
 966    :              << output_pdb_.value();
 967    :    base::File::Error error;
 968  E :    if (!base::ReplaceFileW(temp_path, output_pdb_, &error)) {
 969  i :      LOG(ERROR) << "Unable to replace PDB file.";
 970  i :      return false;
 971    :    }
 972    :  
 973  E :    return true;
 974  E :  }
 975    :  
 976    :  }  // namespace zap_timestamp

Coverage information generated Fri Jul 29 11:00:21 2016.