Coverage for /Syzygy/zap_timestamp/zap_timestamp.cc

CoverageLines executed / instrumented / missingexe / inst / missLanguageGroup
78.7%3634610.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    :      if (!LookupBlockOffset(*image_layout, ref_it->first, ref_it->second.size,
 160  E :                             &src_block, &src_offset)) {
 161  E :        continue;
 162    :      }
 163    :  
 164  E :      BlockGraph::Block* dst_block = NULL;
 165  E :      BlockGraph::Offset dst_offset = 0;
 166    :      if (!LookupBlockOffset(*image_layout, ref_it->second.address, 1, &dst_block,
 167  E :                             &dst_offset)) {
 168  E :        continue;
 169    :      }
 170    :  
 171    :      // Make the final reference.
 172    :      BlockGraph::Reference ref(ref_it->second.type, ref_it->second.size,
 173  E :                                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* 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    :    if (!file_addr_space->Insert(PatchAddressSpace::Range(file_addr, size),
 203  E :                                 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* 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    :    const IMAGE_DATA_DIRECTORY& data_dir_info =
 230  E :        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    :    if (!file_addr_space->Insert(
 258    :            PatchAddressSpace::Range(timestamp_addr, sizeof(DWORD)),
 259  E :            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* 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* 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    :    size_t section_contrib_end_pos = dbi_header->gp_modi_size + sizeof(uint32) +
 426  E :                                     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);  // Skip the signature.
 434  E :    uint8* section_contrib_end = dbi_data + dbi_header->section_contribution_size;
 435  E :    while (dbi_data < section_contrib_end) {
 436    :      pdb::DbiSectionContrib* section_contrib =
 437  E :          reinterpret_cast<pdb::DbiSectionContrib*>(dbi_data);
 438  E :      section_contrib->pad1 = 0;
 439  E :      section_contrib->pad2 = 0;
 440  E :      dbi_data += sizeof(*section_contrib);
 441  E :    }
 442    :  
 443  E :    return true;
 444  E :  }
 445    :  
 446  E :  bool NormalizeSymbolRecordStream(PdbByteStream* stream) {
 447  E :    DCHECK(stream != NULL);
 448    :  
 449  E :    uint8* data = stream->data();
 450  E :    uint8* data_end = data + stream->length();
 451    :  
 452  E :    while (data < data_end) {
 453    :      // Get the size of the symbol record and skip past it.
 454  E :      uint16* size = reinterpret_cast<uint16*>(data);
 455  E :      data += sizeof(*size);
 456    :  
 457    :      // The size of the symbol record, plus its uint16 length, must be a multiple
 458    :      // of 4. Each symbol record consists of the length followed by a symbol
 459    :      // type (also a short), so the size needs to be at least of length 2.
 460    :      // See http://code.google.com/p/syzygy/wiki/PdbFileFormat for a discussion
 461    :      // of the format of this stream.
 462  E :      DCHECK_LE(2u, *size);
 463  E :      DCHECK_EQ(0u, ((*size + sizeof(*size)) % 4));
 464    :  
 465    :      // Up to the last 3 bytes are padding, as the record gets rounded up to
 466    :      // a multiple of 4 in size.
 467    :      static const size_t kMaxPadding = 3;
 468  E :      uint8* end = data + *size;
 469  E :      uint8* tail = end - kMaxPadding;
 470    :  
 471    :      // Skip past the symbol record.
 472  E :      data = end;
 473    :  
 474    :      // Find the null terminator for the record.
 475  E :      for (; tail + 1 < end && *tail != 0; ++tail) {
 476    :        // Intentionally empty.
 477  E :      }
 478    :  
 479    :      // Pad out the rest of the record with nulls (these are usually full of
 480    :      // junk bytes).
 481  E :      for (; tail < end; ++tail)
 482  E :        *tail = 0;
 483  E :    }
 484    :  
 485  E :    return true;
 486  E :  }
 487    :  
 488    :  }  // namespace
 489    :  
 490    :  ZapTimestamp::ZapTimestamp()
 491    :      : image_layout_(&block_graph_),
 492    :        dos_header_block_(NULL),
 493    :        write_image_(true),
 494    :        write_pdb_(true),
 495  E :        overwrite_(false) {
 496    :    // The timestamp can't just be set to zero as that represents a special
 497    :    // value in the PE file. We set it to some arbitrary fixed date in the past.
 498    :    // This is Jan 1, 2010, 0:00:00 GMT. This date shouldn't be too much in
 499    :    // the past, otherwise Windows might trigger a warning saying that the
 500    :    // instrumented image has known incompatibility issues when someone tries to
 501    :    // run it.
 502  E :    timestamp_data_ = 1262304000;
 503    :  
 504    :    // Initialize the age to 1.
 505  E :    pdb_age_data_ = 1;
 506  E :  }
 507    :  
 508  E :  bool ZapTimestamp::Init() {
 509  E :    if (!ValidatePeAndPdbFiles())
 510  E :      return false;
 511    :  
 512  E :    if (!ValidateOutputPaths())
 513  E :      return false;
 514    :  
 515  E :    if (!DecomposePeFile())
 516  i :      return false;
 517    :  
 518  E :    if (!MarkPeFileRanges())
 519  i :      return false;
 520    :  
 521  E :    if (!input_pdb_.empty()) {
 522  E :      if (!CalculatePdbGuid())
 523  i :        return false;
 524    :  
 525  E :      if (!LoadAndUpdatePdbFile())
 526  i :        return false;
 527    :    }
 528    :  
 529  E :    return true;
 530  E :  }
 531    :  
 532  E :  bool ZapTimestamp::Zap() {
 533  E :    if (write_image_) {
 534  E :      if (!WritePeFile())
 535  i :        return false;
 536  E :      OutputSummaryStats(input_image_);
 537    :    }
 538    :  
 539  E :    if (!input_pdb_.empty() && write_pdb_) {
 540  E :      if (!WritePdbFile())
 541  i :        return false;
 542  E :      OutputSummaryStats(input_pdb_);
 543    :    }
 544    :  
 545  E :    return true;
 546  E :  }
 547    :  
 548  E :  bool ZapTimestamp::ValidatePeAndPdbFiles() {
 549  E :    LOG(INFO) << "Analyzing PE file: " << input_image_.value();
 550    :  
 551  E :    if (!base::PathExists(input_image_) || base::DirectoryExists(input_image_)) {
 552  E :      LOG(ERROR) << "PE file not found: " << input_image_.value();
 553  E :      return false;
 554    :    }
 555    :  
 556  E :    if (!pe_file_.Init(input_image_)) {
 557  i :      LOG(ERROR) << "Failed to read PE file: " << input_image_.value();
 558  i :      return false;
 559    :    }
 560    :  
 561  E :    if (input_pdb_.empty()) {
 562    :      // If the image has no code view entry (ie: no matching PDB file)
 563    :      // then accept this fact and leave the PDB path empty.
 564  E :      pe::PdbInfo pe_pdb_info;
 565  E :      if (!pe_pdb_info.Init(input_image_))
 566  E :        return true;
 567    :  
 568    :      // Find the matching PDB file.
 569  E :      if (!pe::FindPdbForModule(input_image_, &input_pdb_)) {
 570  i :        LOG(ERROR) << "Error while searching for PDB file.";
 571  i :        return false;
 572    :      }
 573  E :      if (input_pdb_.empty()) {
 574  E :        LOG(ERROR) << "PDB file not found for PE file: " << input_image_.value();
 575  E :        return false;
 576    :      }
 577  E :      DCHECK(base::PathExists(input_pdb_));
 578  E :    } else {
 579  E :      if (!base::PathExists(input_pdb_) || base::DirectoryExists(input_pdb_)) {
 580  i :        LOG(ERROR) << "PDB file not found: " << input_pdb_.value();
 581    :      }
 582    :    }
 583    :  
 584    :    // Ensure that the PDB and the PE file are consistent with each other.
 585  E :    if (!pe::PeAndPdbAreMatched(input_image_, input_pdb_))
 586  i :      return false;  // This logs verbosely.
 587    :  
 588  E :    LOG(INFO) << "Found matching PDB file: " << input_pdb_.value();
 589    :  
 590  E :    return true;
 591  E :  }
 592    :  
 593  E :  bool ZapTimestamp::ValidateOutputPaths() {
 594  E :    if (output_image_.empty())
 595  E :      output_image_ = input_image_;
 596    :  
 597  E :    if (input_pdb_.empty()) {
 598  E :      if (!output_pdb_.empty()) {
 599  i :        LOG(INFO) << "Ignoring output-pdb path: " << output_pdb_.value();
 600  i :        output_pdb_.clear();
 601    :      }
 602  E :    } else {
 603  E :      if (output_pdb_.empty()) {
 604  E :        if (input_image_.BaseName() == output_image_.BaseName()) {
 605    :          // The input and output have the same basename. Use the input PDB
 606    :          // basename, but place it alongside the output image.
 607  E :          output_pdb_ = output_image_.DirName().Append(input_pdb_.BaseName());
 608  E :        } else {
 609    :          // The basenames don't match. Simpy append ".pdb" to the output
 610    :          // image.
 611  E :          output_pdb_ = output_image_.AddExtension(L"pdb");
 612    :        }
 613    :      }
 614    :    }
 615    :  
 616    :    // If overwriting isn't allowed then double check everything is kosher.
 617  E :    if (!overwrite_) {
 618    :      if (write_image_ && (base::PathExists(output_image_) ||
 619    :                           core::CompareFilePaths(input_image_, output_image_) ==
 620  E :                               core::kEquivalentFilePaths)) {
 621  E :        LOG(ERROR) << "Output image file exists. Must enable overwrite.";
 622  E :        return false;
 623    :      }
 624    :      if (write_pdb_ && !output_pdb_.empty() &&
 625    :          (base::PathExists(output_pdb_) ||
 626    :           core::CompareFilePaths(input_pdb_, output_pdb_) ==
 627  E :               core::kEquivalentFilePaths)) {
 628  E :        LOG(ERROR) << "Output PDB file exists. Must enable overwrite.";
 629  E :        return false;
 630    :      }
 631    :    }
 632    :  
 633  E :    return true;
 634  E :  }
 635    :  
 636  E :  bool ZapTimestamp::DecomposePeFile() {
 637    :    // Decompose the image. This is a very high level decomposition only
 638    :    // chunking out the PE structures and references from/to PE blocks.
 639  E :    if (!MiniDecompose(pe_file_, &image_layout_, &dos_header_block_))
 640  i :      return false;
 641    :  
 642  E :    return true;
 643  E :  }
 644    :  
 645  E :  bool ZapTimestamp::MarkPeFileRanges() {
 646  E :    DCHECK(dos_header_block_ != NULL);
 647  E :    LOG(INFO) << "Finding PE fields that need updating.";
 648    :  
 649  E :    DosHeader dos_header;
 650  E :    if (!dos_header.Init(0, dos_header_block_)) {
 651  i :      LOG(ERROR) << "Failed to cast IMAGE_DOS_HEADER.";
 652  i :      return false;
 653    :    }
 654    :  
 655  E :    NtHeaders nt_headers;
 656  E :    if (!dos_header.Dereference(dos_header->e_lfanew, &nt_headers)) {
 657  i :      LOG(ERROR) << "Failed to dereference IMAGE_NT_HEADERS.";
 658  i :      return false;
 659    :    }
 660    :  
 661    :    // Mark the export data directory timestamp.
 662    :    if (!MarkDataDirectoryTimestamps<IMAGE_EXPORT_DIRECTORY>(
 663    :            pe_file_, nt_headers, IMAGE_DIRECTORY_ENTRY_EXPORT,
 664    :            "Export Directory", reinterpret_cast<const uint8*>(&timestamp_data_),
 665  E :            &pe_file_addr_space_)) {
 666    :      // This logs verbosely on failure.
 667  i :      return false;
 668    :    }
 669    :  
 670    :    // Mark the resource data directory timestamp.
 671    :    if (!MarkDataDirectoryTimestamps<IMAGE_RESOURCE_DIRECTORY>(
 672    :            pe_file_, nt_headers, IMAGE_DIRECTORY_ENTRY_RESOURCE,
 673    :            "Resource Directory",
 674    :            reinterpret_cast<const uint8*>(&timestamp_data_),
 675  E :            &pe_file_addr_space_)) {
 676    :      // This logs verbosely on failure.
 677  i :      return false;
 678    :    }
 679    :  
 680    :    // Find the debug directory.
 681  E :    ImageDebugDirectory debug_dir;
 682    :    const IMAGE_DATA_DIRECTORY& debug_dir_info =
 683  E :        nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG];
 684  E :    if (nt_headers.HasReference(debug_dir_info.VirtualAddress))
 685  E :      nt_headers.Dereference(debug_dir_info.VirtualAddress, &debug_dir);
 686    :  
 687    :    // Within that, find the codeview debug entry. We also update every other
 688    :    // debug timestamp.
 689  E :    CvInfoPdb cv_info_pdb;
 690  E :    RelativeAddress rel_addr;
 691  E :    if (debug_dir.block()) {
 692  E :      for (size_t i = 0; i < debug_dir.ElementCount(); ++i) {
 693    :        rel_addr = debug_dir.block()->addr() +
 694  E :                   debug_dir.OffsetOf(debug_dir[i].TimeDateStamp);
 695  E :        std::string name = base::StringPrintf("Debug Directory %d Timestamp", i);
 696    :        if (!MarkData(pe_file_, rel_addr, sizeof(timestamp_data_),
 697    :                      reinterpret_cast<const uint8*>(&timestamp_data_), name,
 698  E :                      &pe_file_addr_space_)) {
 699  i :          LOG(ERROR) << "Failed to mark TimeDateStamp of debug directory " << i
 700    :                     << ".";
 701  i :          return false;
 702    :        }
 703    :  
 704  E :        if (debug_dir[i].Type == IMAGE_DEBUG_TYPE_CODEVIEW) {
 705  E :          if (cv_info_pdb.block() != NULL) {
 706  i :            LOG(ERROR) << "Found multiple CodeView debug directories.";
 707  i :            return false;
 708    :          }
 709    :          if (!debug_dir.Dereference(debug_dir[i].PointerToRawData,
 710  E :                                     &cv_info_pdb)) {
 711  i :            LOG(ERROR) << "Failed to dereference CodeView debug directory.";
 712  i :            return false;
 713    :          }
 714    :        }
 715  E :      }
 716    :    }
 717    :  
 718    :    // We should have found a code view debug directory pointing to the PDB file.
 719  E :    if (!input_pdb_.empty()) {
 720  E :      if (cv_info_pdb.block() == NULL) {
 721  i :        LOG(ERROR) << "Failed to find CodeView debug directory.";
 722  i :        return false;
 723    :      }
 724    :  
 725    :      // Get the file offset of the PDB age and mark it.
 726    :      rel_addr = cv_info_pdb.block()->addr() +
 727  E :                 cv_info_pdb.OffsetOf(cv_info_pdb->pdb_age);
 728    :      if (!MarkData(pe_file_, rel_addr, sizeof(pdb_age_data_),
 729    :                    reinterpret_cast<const uint8*>(&pdb_age_data_), "PDB Age",
 730  E :                    &pe_file_addr_space_)) {
 731  i :        LOG(ERROR) << "Failed to mark PDB age.";
 732  i :        return false;
 733    :      }
 734    :  
 735    :      // Get the file offset of the PDB guid and mark it.
 736    :      rel_addr = cv_info_pdb.block()->addr() +
 737  E :                 cv_info_pdb.OffsetOf(cv_info_pdb->signature);
 738    :      if (!MarkData(pe_file_, rel_addr, sizeof(pdb_guid_data_),
 739    :                    reinterpret_cast<const uint8*>(&pdb_guid_data_), "PDB GUID",
 740  E :                    &pe_file_addr_space_)) {
 741  i :        LOG(ERROR) << "Failed to mark PDB GUID.";
 742  i :        return false;
 743    :      }
 744    :    }
 745    :  
 746    :    // Get the file offset of the PE checksum and mark it.
 747    :    rel_addr = nt_headers.block()->addr() +
 748  E :               nt_headers.OffsetOf(nt_headers->OptionalHeader.CheckSum);
 749    :    if (!MarkData(pe_file_, rel_addr, sizeof(DWORD), NULL, "PE Checksum",
 750  E :                  &pe_file_addr_space_)) {
 751  i :      LOG(ERROR) << "Failed to mark PE checksum.";
 752  i :      return false;
 753    :    }
 754    :  
 755    :    // Get the file offset of the PE timestamp and mark it.
 756    :    rel_addr = nt_headers.block()->addr() +
 757  E :               nt_headers.OffsetOf(nt_headers->FileHeader.TimeDateStamp);
 758    :    if (!MarkData(pe_file_, rel_addr, sizeof(timestamp_data_),
 759    :                  reinterpret_cast<uint8*>(&timestamp_data_), "PE Timestamp",
 760  E :                  &pe_file_addr_space_)) {
 761  i :      LOG(ERROR) << "Failed to mark PE timestamp.";
 762  i :      return false;
 763    :    }
 764    :  
 765  E :    return true;
 766  E :  }
 767    :  
 768  E :  bool ZapTimestamp::CalculatePdbGuid() {
 769  E :    DCHECK(!input_pdb_.empty());
 770    :  
 771  E :    LOG(INFO) << "Calculating PDB GUID from PE file contents.";
 772    :  
 773  E :    base::ScopedFILE pe_file(base::OpenFile(input_image_, "rb"));
 774  E :    if (pe_file.get() == NULL) {
 775  i :      LOG(ERROR) << "Failed to open PE file for reading: "
 776    :                 << input_image_.value();
 777  i :      return false;
 778    :    }
 779    :  
 780    :    // Get the length of the entire file.
 781  E :    if (::fseek(pe_file.get(), 0, SEEK_END) != 0) {
 782  i :      LOG(ERROR) << "Failed to fseek to end of file.";
 783  i :      return false;
 784    :    }
 785  E :    FileOffsetAddress end(::ftell(pe_file.get()));
 786    :  
 787    :    // Seek back to the beginning.
 788  E :    if (::fseek(pe_file.get(), 0, SEEK_SET) != 0) {
 789  i :      LOG(ERROR) << "Failed to fseek to beginning of file.";
 790  i :      return false;
 791    :    }
 792    :  
 793    :    // Initialize the MD5 structure.
 794  E :    base::MD5Context md5_context = {0};
 795  E :    base::MD5Init(&md5_context);
 796    :  
 797    :    // We seek through the bits of the file that will be changed, and skip those.
 798    :    // The rest of the file (the static parts) are fed through an MD5 hash and
 799    :    // used to generated a unique and stable GUID.
 800  E :    FileOffsetAddress cur(0);
 801  E :    PatchAddressSpace::const_iterator range_it = pe_file_addr_space_.begin();
 802  E :    for (; range_it != pe_file_addr_space_.end(); ++range_it) {
 803    :      // Consume any data before this range.
 804  E :      if (cur < range_it->first.start()) {
 805  E :        size_t bytes_to_hash = range_it->first.start() - cur;
 806  E :        if (!Md5Consume(bytes_to_hash, pe_file.get(), &md5_context))
 807  i :          return false;  // This logs verbosely for us.
 808    :      }
 809    :  
 810  E :      if (::fseek(pe_file.get(), range_it->first.size(), SEEK_CUR)) {
 811  i :        LOG(ERROR) << "Failed to fseek past marked range.";
 812    :      }
 813    :  
 814  E :      cur = range_it->first.end();
 815  E :    }
 816    :  
 817    :    // Consume any left-over data.
 818  E :    if (cur < end) {
 819  E :      if (!Md5Consume(end - cur, pe_file.get(), &md5_context))
 820  i :        return false;  // This logs verbosely for us.
 821    :    }
 822    :  
 823  E :    DCHECK_EQ(end.value(), static_cast<uint32>(::ftell(pe_file.get())));
 824    :  
 825    :    static_assert(sizeof(base::MD5Digest) == sizeof(pdb_guid_data_),
 826    :                  "MD5Digest and GUID size mismatch.");
 827    :    base::MD5Final(reinterpret_cast<base::MD5Digest*>(&pdb_guid_data_),
 828  E :                   &md5_context);
 829  E :    LOG(INFO) << "Final GUID is "
 830    :              << base::MD5DigestToBase16(
 831    :                     *reinterpret_cast<base::MD5Digest*>(&pdb_guid_data_))
 832    :              << ".";
 833    :  
 834  E :    return true;
 835  E :  }
 836    :  
 837  E :  bool ZapTimestamp::LoadAndUpdatePdbFile() {
 838  E :    DCHECK(!input_pdb_.empty());
 839  E :    DCHECK(pdb_file_.get() == NULL);
 840    :  
 841  E :    pdb_file_.reset(new PdbFile());
 842  E :    PdbReader pdb_reader;
 843  E :    if (!pdb_reader.Read(input_pdb_, pdb_file_.get())) {
 844  i :      LOG(ERROR) << "Failed to read PDB file: " << input_pdb_.value();
 845  i :      return false;
 846    :    }
 847    :  
 848    :    // We turf the old directory stream as a fresh PDB does not have one. It's
 849    :    // also meaningless after we rewrite a PDB as the old blocks it refers to
 850    :    // will no longer exist.
 851  E :    pdb_file_->ReplaceStream(pdb::kPdbOldDirectoryStream, NULL);
 852    :  
 853    :    scoped_refptr<PdbStream> header_reader =
 854  E :        GetWritableStream(pdb::kPdbHeaderInfoStream, pdb_file_.get());
 855  E :    if (header_reader.get() == NULL) {
 856  i :      LOG(ERROR) << "No header info stream in PDB file: " << input_pdb_.value();
 857  i :      return false;
 858    :    }
 859    :  
 860    :    scoped_refptr<WritablePdbStream> header_writer =
 861  E :        header_reader->GetWritableStream();
 862  E :    DCHECK(header_writer.get() != NULL);
 863    :  
 864    :    // Update the timestamp, the age and the signature.
 865  E :    LOG(INFO) << "Updating PDB header.";
 866  E :    header_writer->set_pos(offsetof(pdb::PdbInfoHeader70, timestamp));
 867  E :    header_writer->Write(static_cast<uint32>(timestamp_data_));
 868  E :    header_writer->Write(static_cast<uint32>(pdb_age_data_));
 869  E :    header_writer->Write(pdb_guid_data_);
 870    :  
 871    :    // Normalize the DBI stream in place.
 872  E :    scoped_refptr<PdbByteStream> dbi_stream(new PdbByteStream());
 873  E :    CHECK(dbi_stream->Init(pdb_file_->GetStream(pdb::kDbiStream).get()));
 874  E :    pdb_file_->ReplaceStream(pdb::kDbiStream, dbi_stream.get());
 875  E :    if (!NormalizeDbiStream(pdb_age_data_, dbi_stream.get())) {
 876  i :      LOG(ERROR) << "Failed to normalize DBI stream.";
 877  i :      return false;
 878    :    }
 879    :  
 880  E :    uint8* dbi_data = dbi_stream->data();
 881  E :    pdb::DbiHeader* dbi_header = reinterpret_cast<pdb::DbiHeader*>(dbi_data);
 882    :  
 883    :    // Normalize the symbol record stream in place.
 884  E :    scoped_refptr<PdbByteStream> symrec_stream(new PdbByteStream());
 885    :    CHECK(symrec_stream->Init(
 886  E :        pdb_file_->GetStream(dbi_header->symbol_record_stream).get()));
 887    :    pdb_file_->ReplaceStream(dbi_header->symbol_record_stream,
 888  E :                             symrec_stream.get());
 889  E :    if (!NormalizeSymbolRecordStream(symrec_stream.get())) {
 890  i :      LOG(ERROR) << "Failed to normalize symbol record stream.";
 891  i :      return false;
 892    :    }
 893    :  
 894    :    // Normalize the public symbol info stream. There's a DWORD of padding at
 895    :    // offset 24 that we want to zero.
 896    :    scoped_refptr<PdbStream> pubsym_reader =
 897  E :        GetWritableStream(dbi_header->public_symbol_info_stream, pdb_file_.get());
 898    :    scoped_refptr<WritablePdbStream> pubsym_writer =
 899  E :        pubsym_reader->GetWritableStream();
 900  E :    DCHECK(pubsym_writer.get() != NULL);
 901  E :    pubsym_writer->set_pos(24);
 902  E :    pubsym_writer->Write(static_cast<uint32>(0));
 903    :  
 904  E :    return true;
 905  E :  }
 906    :  
 907  E :  bool ZapTimestamp::WritePeFile() {
 908    :    if (core::CompareFilePaths(input_image_, output_image_) !=
 909  E :        core::kEquivalentFilePaths) {
 910    :      if (::CopyFileW(input_image_.value().c_str(), output_image_.value().c_str(),
 911  E :                      FALSE) == FALSE) {
 912  i :        LOG(ERROR) << "Failed to write output image: %s" << output_image_.value();
 913  i :        return false;
 914    :      }
 915    :    }
 916    :  
 917  E :    if (!UpdateFileInPlace(output_image_, pe_file_addr_space_))
 918  i :      return false;
 919    :  
 920  E :    LOG(INFO) << "Updating checksum for PE file: " << output_image_.value();
 921  E :    if (!pe::PEFileWriter::UpdateFileChecksum(output_image_)) {
 922  i :      LOG(ERROR) << "Failed to update checksum for PE file: "
 923    :                 << output_image_.value();
 924  i :      return false;
 925    :    }
 926    :  
 927  E :    return true;
 928  E :  }
 929    :  
 930  E :  bool ZapTimestamp::WritePdbFile() {
 931  E :    DCHECK(!input_pdb_.empty());
 932    :  
 933    :    // We actually completely rewrite the PDB file to a temporary location, and
 934    :    // then move it over top of the existing one. This is because pdb_file_
 935    :    // actually has an open file handle to the original PDB.
 936    :  
 937    :    // We create a temporary directory alongside the final destination so as
 938    :    // not to cross volume boundaries.
 939  E :    base::FilePath output_dir = output_pdb_.DirName();
 940  E :    base::ScopedTempDir temp_dir;
 941  E :    if (!temp_dir.CreateUniqueTempDirUnderPath(output_dir)) {
 942  i :      LOG(ERROR) << "Failed to create temporary directory in \""
 943    :                 << output_dir.value() << "\".";
 944  i :      return false;
 945    :    }
 946    :  
 947    :    // Generate the path to the rewritten PDB.
 948  E :    base::FilePath temp_path = temp_dir.path().Append(input_pdb_.BaseName());
 949    :  
 950  E :    PdbWriter pdb_writer;
 951  E :    LOG(INFO) << "Creating temporary PDB file: " << temp_path.value();
 952  E :    if (!pdb_writer.Write(temp_path, *pdb_file_.get())) {
 953  i :      LOG(ERROR) << "Failed to write new PDB: " << temp_path.value();
 954  i :      return false;
 955    :    }
 956    :  
 957    :    // Free up the PDB file. This will close the open file handle to the original
 958    :    // PDB file.
 959  E :    pdb_file_.reset(NULL);
 960    :  
 961    :    // Copy over top of the original file.
 962  E :    LOG(INFO) << "Temporary PDB file replacing destination PDB: "
 963    :              << output_pdb_.value();
 964    :    base::File::Error error;
 965  E :    if (!base::ReplaceFileW(temp_path, output_pdb_, &error)) {
 966  i :      LOG(ERROR) << "Unable to replace PDB file.";
 967  i :      return false;
 968    :    }
 969    :  
 970  E :    return true;
 971  E :  }
 972    :  
 973    :  }  // namespace zap_timestamp

Coverage information generated Thu Jan 14 17:40:38 2016.