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    :  #include "syzygy/pe/transforms/add_debug_directory_entry_transform.h"
  16    :  
  17    :  #include "syzygy/block_graph/typed_block.h"
  18    :  #include "syzygy/pe/pe_utils.h"
  19    :  
  20    :  namespace pe {
  21    :  namespace transforms {
  22    :  
  23    :  using block_graph::TypedBlock;
  24    :  
  25    :  typedef TypedBlock<IMAGE_DOS_HEADER> DosHeader;
  26    :  typedef TypedBlock<IMAGE_NT_HEADERS> NtHeaders;
  27    :  typedef TypedBlock<IMAGE_DEBUG_DIRECTORY> ImageDebugDirectory;
  28    :  
  29    :  const char AddDebugDirectoryEntryTransform::kTransformName[] =
  30    :      "AddDebugDirectoryEntryTransform";
  31    :  
  32    :  bool AddDebugDirectoryEntryTransform::TransformBlockGraph(
  33    :      const TransformPolicyInterface* policy,
  34    :      BlockGraph* block_graph,
  35  E :      BlockGraph::Block* dos_header_block) {
  36  E :    DCHECK(block_graph != NULL);
  37  E :    DCHECK(dos_header_block != NULL);
  38    :  
  39  E :    added_ = false;
  40  E :    block_ = NULL;
  41  E :    offset_ = -1;
  42    :  
  43  E :    DosHeader dos_header;
  44  E :    NtHeaders nt_headers;
  45    :    if (!dos_header.Init(0, dos_header_block) ||
  46  E :        !dos_header.Dereference(dos_header->e_lfanew, &nt_headers)) {
  47  i :      LOG(ERROR) << "Unable to dereference PE image headers.";
  48  i :      return false;
  49    :    }
  50    :  
  51    :    // Don't have a debug directory? Then make one with a single entry.
  52    :    // In general, keeping around a reference to data inside a TypedBlock is not
  53    :    // safe, as if the data is resized the reference will no longer be valid.
  54    :    // However, I do not modify the underlying block for the lifetime of this
  55    :    // function, hence reusing this reference is safe.
  56    :    IMAGE_DATA_DIRECTORY& debug_dir_info =
  57  E :        nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG];
  58  E :    if (!nt_headers.HasReference(debug_dir_info.VirtualAddress)) {
  59  E :      debug_dir_info.Size = sizeof(IMAGE_DEBUG_DIRECTORY);
  60    :  
  61    :      BlockGraph::Section* section = block_graph->FindOrAddSection(
  62  E :          kReadOnlyDataSectionName, kReadOnlyDataCharacteristics);
  63  E :      DCHECK(section != NULL);
  64    :  
  65    :      BlockGraph::Block* debug_dir_block = block_graph->AddBlock(
  66  E :          BlockGraph::DATA_BLOCK, debug_dir_info.Size, "Debug Directory");
  67  E :      DCHECK(debug_dir_block != NULL);
  68  E :      debug_dir_block->set_section(section->id());
  69  E :      debug_dir_block->AllocateData(debug_dir_info.Size);
  70    :  
  71    :      nt_headers.SetReference(BlockGraph::RELATIVE_REF,
  72    :                              debug_dir_info.VirtualAddress,
  73    :                              debug_dir_block,
  74  E :                              0, 0);
  75    :  
  76  E :      added_ = true;
  77    :    }
  78    :  
  79    :    // Get the debug directory, and remember it for post-transform.
  80  E :    ImageDebugDirectory debug_dir;
  81  E :    if (!nt_headers.Dereference(debug_dir_info.VirtualAddress, &debug_dir)) {
  82  i :        LOG(ERROR) << "Unable to dereference ImageDebugDirectory.";
  83  i :        return false;
  84    :    }
  85  E :    block_ = debug_dir.block();
  86    :  
  87    :    // Did we already add an entry? Initialize it and be done with it. This can
  88    :    // happen if there was no debug directory to begin with.
  89  E :    if (added_) {
  90  E :      offset_ = 0;
  91  E :      debug_dir->Type = type_;
  92  E :      return true;
  93    :    }
  94    :  
  95    :    // If we get here we've got a non-empty debug data directory with entries
  96    :    // that we did not make. We either have to find an existing entry or create
  97    :    // a new one.
  98    :  
  99    :    // If we're not explicitly adding another entry, look for an existing one
 100    :    // with the matching type.
 101  E :    if (!always_add_) {
 102  E :      for (size_t i = 0; i < debug_dir.ElementCount(); ++i) {
 103  E :        if (debug_dir[i].Type == type_) {
 104  E :          offset_ = debug_dir.OffsetOf(debug_dir[i]);
 105  E :          break;
 106    :        }
 107  i :      }
 108    :    }
 109    :  
 110    :    // If we found an existing entry we're done.
 111  E :    if (offset_ != -1)
 112  E :      return true;
 113    :  
 114    :    // Make the new entry and initialize it. We only set the type as the rest of
 115    :    // it is already going to be initialized with zeros.
 116  i :    added_ = true;
 117  i :    size_t entry_index = debug_dir.ElementCount();
 118  i :    size_t entry_size = sizeof(IMAGE_DEBUG_DIRECTORY);
 119    :    offset_ = debug_dir.offset() +
 120  i :        entry_index * entry_size;
 121  i :    debug_dir.block()->InsertData(offset_, entry_size, true);
 122  i :    DCHECK_EQ(entry_index + 1, debug_dir.ElementCount());
 123  i :    debug_dir[entry_index].Type = type_;
 124    :  
 125    :    // Update the debug directory info struct.
 126  i :    debug_dir_info.Size += entry_size;
 127    :  
 128  i :    return true;
 129  E :  }
 130    :  
 131    :  }  // namespace transforms
 132    :  }  // namespace pe