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/instrument/instrumenter_with_agent.h"
  16    :  
  17    :  #include "base/file_util.h"
  18    :  #include "base/logging.h"
  19    :  #include "syzygy/common/application.h"
  20    :  
  21    :  namespace instrument {
  22    :  
  23  E :  bool InstrumenterWithAgent::ParseCommandLine(const CommandLine* command_line) {
  24  E :    DCHECK(command_line != NULL);
  25    :  
  26    :    // TODO(chrisha): Simplify the input/output image parsing once external
  27    :    //     tools have been updated.
  28    :  
  29    :    // Parse the input image.
  30  E :    if (command_line->HasSwitch("input-dll")) {
  31  E :      LOG(WARNING) << "DEPRECATED: Using --input-dll.";
  32    :      input_dll_path_ = common::AppImplBase::AbsolutePath(
  33  E :          command_line->GetSwitchValuePath("input-dll"));
  34  E :    } else {
  35    :      input_dll_path_ = common::AppImplBase::AbsolutePath(
  36  E :          command_line->GetSwitchValuePath("input-image"));
  37    :    }
  38    :  
  39    :    // Parse the output image.
  40  E :    if (command_line->HasSwitch("output-dll")) {
  41  E :      LOG(WARNING) << "DEPRECATED: Using --output-dll.";
  42    :      output_dll_path_ = common::AppImplBase::AbsolutePath(
  43  E :          command_line->GetSwitchValuePath("output-dll"));
  44  E :    } else {
  45    :      output_dll_path_ = common::AppImplBase::AbsolutePath(
  46  E :          command_line->GetSwitchValuePath("output-image"));
  47    :    }
  48    :  
  49    :    // Ensure that both input and output have been specified.
  50  E :    if (input_dll_path_.empty() || output_dll_path_.empty())
  51  E :      return false;
  52    :  
  53    :    // Parse the remaining command line arguments.
  54    :    input_pdb_path_ = common::AppImplBase::AbsolutePath(
  55  E :        command_line->GetSwitchValuePath("input-pdb"));
  56    :    output_pdb_path_ = common::AppImplBase::AbsolutePath(
  57  E :        command_line->GetSwitchValuePath("output-pdb"));
  58  E :    allow_overwrite_ = command_line->HasSwitch("overwrite");
  59  E :    new_decomposer_ = command_line->HasSwitch("new-decomposer");
  60  E :    no_augment_pdb_ = command_line->HasSwitch("no-augment-pdb");
  61  E :    no_parse_debug_info_ = command_line->HasSwitch("no-parse-debug-info");
  62  E :    no_strip_strings_ = command_line->HasSwitch("no-strip-strings");
  63    :  
  64  E :    return true;
  65  E :  }
  66    :  
  67  E :  bool InstrumenterWithAgent::Instrument() {
  68  E :    pe::PERelinker* relinker = GetRelinker();
  69  E :    DCHECK(relinker != NULL);
  70  E :    relinker->set_input_path(input_dll_path_);
  71  E :    relinker->set_input_pdb_path(input_pdb_path_);
  72  E :    relinker->set_output_path(output_dll_path_);
  73  E :    relinker->set_output_pdb_path(output_pdb_path_);
  74  E :    relinker->set_allow_overwrite(allow_overwrite_);
  75  E :    relinker->set_augment_pdb(!no_augment_pdb_);
  76  E :    relinker->set_parse_debug_info(!no_parse_debug_info_);
  77  E :    relinker->set_use_new_decomposer(new_decomposer_);
  78  E :    relinker->set_strip_strings(!no_strip_strings_);
  79    :  
  80    :    // Initialize the relinker. This does the decomposition, etc.
  81  E :    if (!relinker->Init()) {
  82  E :      LOG(ERROR) << "Failed to initialize relinker.";
  83  E :      return false;
  84    :    }
  85    :  
  86    :    // Do the actual instrumentation.
  87  E :    if (!InstrumentImpl())
  88  i :      return false;
  89    :  
  90    :    // We let the PERelinker use the implicit OriginalOrderer.
  91  E :    if (!relinker->Relink()) {
  92  E :      LOG(ERROR) << "Unable to relink input image.";
  93  E :      return false;
  94    :    }
  95    :  
  96  E :    return true;
  97  E :  }
  98    :  
  99  i :  pe::PERelinker* InstrumenterWithAgent::GetRelinker() {
 100  i :    if (relinker_.get() == NULL) {
 101  i :      relinker_.reset(new pe::PERelinker());
 102  i :      CHECK(relinker_.get() != NULL);
 103    :    }
 104  i :    return relinker_.get();
 105  i :  }
 106    :  
 107    :  }  // namespace instrument