1    :  // Copyright 2013 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/pe_coff_relinker.h"
  16    :  
  17    :  #include "syzygy/block_graph/orderers/original_orderer.h"
  18    :  
  19    :  namespace pe {
  20    :  namespace {
  21    :  
  22    :  typedef block_graph::BlockGraphTransformInterface Transform;
  23    :  typedef block_graph::BlockGraphOrdererInterface Orderer;
  24    :  
  25    :  using block_graph::ApplyBlockGraphTransform;
  26    :  using block_graph::BlockGraph;
  27    :  using block_graph::OrderedBlockGraph;
  28    :  using block_graph::TransformPolicyInterface;
  29    :  using core::RelativeAddress;
  30    :  
  31    :  }  // namespace
  32    :  
  33    :  PECoffRelinker::PECoffRelinker(const TransformPolicyInterface* transform_policy)
  34  E :      : transform_policy_(transform_policy),
  35  E :        allow_overwrite_(false),
  36  E :        inited_(false),
  37  E :        input_image_layout_(&block_graph_),
  38  E :        headers_block_(NULL) {
  39  E :    DCHECK(transform_policy != NULL);
  40  E :  }
  41    :  
  42  E :  bool PECoffRelinker::AppendTransform(BlockGraphTransform* transform) {
  43  E :    DCHECK(transform != NULL);
  44  E :    transforms_.push_back(transform);
  45  E :    return true;
  46  E :  }
  47    :  
  48    :  bool PECoffRelinker::AppendTransforms(
  49  E :      const std::vector<BlockGraphTransform*>& transforms) {
  50  E :    transforms_.insert(transforms_.end(), transforms.begin(), transforms.end());
  51  E :    return true;
  52  E :  }
  53    :  
  54  E :  bool PECoffRelinker::AppendLayoutTransform(ImageLayoutTransform* transform) {
  55  E :    DCHECK(transform != NULL);
  56  E :    layout_transforms_.push_back(transform);
  57  E :    return true;
  58  E :  }
  59    :  
  60    :  bool PECoffRelinker::AppendLayoutTransforms(
  61  E :    const std::vector<ImageLayoutTransform*>& transforms) {
  62  E :    layout_transforms_.insert(layout_transforms_.end(),
  63    :                              transforms.begin(),
  64    :                              transforms.end());
  65  E :    return true;
  66  E :  }
  67    :  
  68  E :  bool PECoffRelinker::AppendOrderer(Orderer* orderer) {
  69  E :    DCHECK(orderer != NULL);
  70  E :    orderers_.push_back(orderer);
  71  E :    return true;
  72  E :  }
  73    :  
  74  E :  bool PECoffRelinker::AppendOrderers(const std::vector<Orderer*>& orderers) {
  75  E :    orderers_.insert(orderers_.end(), orderers.begin(), orderers.end());
  76  E :    return true;
  77  E :  }
  78    :  
  79  E :  bool PECoffRelinker::ApplyUserTransforms() {
  80  E :    LOG(INFO) << "Transforming block graph.";
  81  E :    if (!block_graph::ApplyBlockGraphTransforms(
  82    :             transforms_, transform_policy_, &block_graph_, headers_block_)) {
  83  E :      return false;
  84    :    }
  85  E :    return true;
  86  E :  }
  87    :  
  88    :  bool PECoffRelinker::ApplyUserLayoutTransforms(
  89    :      pe::ImageLayout* image_layout,
  90  E :      OrderedBlockGraph* ordered_graph) {
  91  E :    LOG(INFO) << "Transforming layout.";
  92  E :    if (!block_graph::ApplyImageLayoutTransforms(
  93    :        layout_transforms_, transform_policy_, image_layout, ordered_graph)) {
  94  i :      return false;
  95    :    }
  96  E :    return true;
  97  E :  }
  98    :  
  99  E :  bool PECoffRelinker::ApplyUserOrderers(OrderedBlockGraph* ordered_graph) {
 100  E :    LOG(INFO) << "Ordering block graph.";
 101    :  
 102  E :    if (orderers_.empty()) {
 103    :      // Default orderer.
 104  E :      LOG(INFO) << "No orderers specified, applying default orderer.";
 105  E :      block_graph::orderers::OriginalOrderer default_orderer;
 106  E :      if (!default_orderer.OrderBlockGraph(ordered_graph, headers_block_)) {
 107  i :        LOG(ERROR) << "Orderer failed: " << default_orderer.name() << ".";
 108  i :        return false;
 109    :      }
 110  E :    } else {
 111    :      // Supplied orderers.
 112  E :      if (!block_graph::ApplyBlockGraphOrderers(
 113    :               orderers_, ordered_graph, headers_block_)) {
 114  E :        return false;
 115    :      }
 116    :    }
 117    :  
 118  E :    return true;
 119  E :  }
 120    :  
 121    :  }  // namespace pe