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