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/instrumenters/instrumenter_with_agent.h"
16 :
17 : #include "base/file_util.h"
18 : #include "base/logging.h"
19 : #include "syzygy/common/application.h"
20 : #include "syzygy/pe/pe_utils.h"
21 :
22 : namespace instrument {
23 : namespace instrumenters {
24 :
25 : namespace {
26 :
27 E : bool GetImageFormat(const base::FilePath& path, pe::ImageFormat* image_format) {
28 E : DCHECK(image_format != NULL);
29 :
30 : // Determine the type of the input.
31 E : pe::FileType file_type = pe::kUnknownFileType;
32 E : if (!pe::GuessFileType(path, &file_type)) {
33 E : LOG(ERROR) << "Failed to determine file type of \""
34 : << path.value() << "\".";
35 E : return false;
36 : }
37 :
38 E : if (file_type == pe::kCoffFileType) {
39 E : *image_format = pe::COFF_IMAGE;
40 E : return true;
41 : }
42 :
43 E : if (file_type == pe::kPeFileType) {
44 E : *image_format = pe::PE_IMAGE;
45 E : return true;
46 : }
47 :
48 i : LOG(ERROR) << "File is not a PE or COFF image: " << path.value();
49 i : return false;
50 E : }
51 :
52 : } // namespace
53 :
54 E : bool InstrumenterWithAgent::ParseCommandLine(const CommandLine* command_line) {
55 E : DCHECK(command_line != NULL);
56 :
57 : // TODO(chrisha): Simplify the input/output image parsing once external
58 : // tools have been updated.
59 :
60 : // Parse the input image.
61 E : if (command_line->HasSwitch("input-dll")) {
62 E : LOG(WARNING) << "DEPRECATED: Using --input-dll.";
63 : input_image_path_ = common::AppImplBase::AbsolutePath(
64 E : command_line->GetSwitchValuePath("input-dll"));
65 E : } else {
66 : input_image_path_ = common::AppImplBase::AbsolutePath(
67 E : command_line->GetSwitchValuePath("input-image"));
68 : }
69 :
70 : // Parse the output image.
71 E : if (command_line->HasSwitch("output-dll")) {
72 E : LOG(WARNING) << "DEPRECATED: Using --output-dll.";
73 : output_image_path_ = common::AppImplBase::AbsolutePath(
74 E : command_line->GetSwitchValuePath("output-dll"));
75 E : } else {
76 : output_image_path_ = common::AppImplBase::AbsolutePath(
77 E : command_line->GetSwitchValuePath("output-image"));
78 : }
79 :
80 : // Ensure that both input and output have been specified.
81 E : if (input_image_path_.empty() || output_image_path_.empty()) {
82 E : LOG(ERROR) << "You must provide input and output file names.";
83 E : return false;
84 : }
85 :
86 : // Parse the remaining command line arguments.
87 : input_pdb_path_ = common::AppImplBase::AbsolutePath(
88 E : command_line->GetSwitchValuePath("input-pdb"));
89 : output_pdb_path_ = common::AppImplBase::AbsolutePath(
90 E : command_line->GetSwitchValuePath("output-pdb"));
91 E : allow_overwrite_ = command_line->HasSwitch("overwrite");
92 E : debug_friendly_ = command_line->HasSwitch("debug-friendly");
93 E : old_decomposer_ = command_line->HasSwitch("old-decomposer");
94 E : no_augment_pdb_ = command_line->HasSwitch("no-augment-pdb");
95 E : no_strip_strings_ = command_line->HasSwitch("no-strip-strings");
96 :
97 E : if (!agent_dll_.empty()) {
98 E : LOG(INFO) << "Default agent DLL for " << InstrumentationMode() << " mode "
99 : << "is \"" << agent_dll_ << "\".";
100 : }
101 :
102 : // Parse the custom agent if one is specified.
103 E : if (command_line->HasSwitch("agent")) {
104 E : std::string new_agent_dll = command_line->GetSwitchValueASCII("agent");
105 E : if (new_agent_dll != agent_dll_) {
106 E : agent_dll_ = new_agent_dll;
107 E : LOG(INFO) << "Using custom agent DLL \"" << agent_dll_ << "\".";
108 : }
109 E : }
110 :
111 E : if (agent_dll_.empty()) {
112 i : LOG(ERROR) << "No agent DLL has been specified.";
113 i : return false;
114 : }
115 :
116 E : if (!ParseAdditionalCommandLineArguments(command_line)) {
117 E : LOG(ERROR) << "Unable to parse the additional arguments from the command "
118 : << "line.";
119 E : return false;
120 : }
121 :
122 E : return true;
123 E : }
124 :
125 E : bool InstrumenterWithAgent::Instrument() {
126 E : if (!CreateRelinker())
127 E : return false;
128 :
129 : // Initialize the relinker. This does the decomposition, etc.
130 E : if (!relinker_->Init()) {
131 E : LOG(ERROR) << "Failed to initialize relinker.";
132 E : return false;
133 : }
134 :
135 : // Let the instrumenter implementation set up the relinker and anything else
136 : // that is required.
137 E : if (!InstrumentImpl())
138 i : return false;
139 :
140 : // Do the actual instrumentation by running the relinker.
141 E : if (!relinker_->Relink()) {
142 E : LOG(ERROR) << "Unable to relink input image.";
143 E : return false;
144 : }
145 :
146 E : return true;
147 E : }
148 :
149 : bool InstrumenterWithAgent::ImageFormatIsSupported(
150 E : pe::ImageFormat image_format) {
151 : // By default we only support PE images.
152 E : if (image_format == pe::PE_IMAGE)
153 E : return true;
154 i : return false;
155 E : }
156 :
157 E : pe::PETransformPolicy* InstrumenterWithAgent::GetPETransformPolicy() {
158 E : DCHECK_EQ(pe::PE_IMAGE, image_format_);
159 E : DCHECK(policy_object_.get() == NULL);
160 E : policy_object_.reset(new pe::PETransformPolicy());
161 E : return static_cast<pe::PETransformPolicy*>(policy_object_.get());
162 E : }
163 :
164 i : pe::CoffTransformPolicy* InstrumenterWithAgent::GetCoffTransformPolicy() {
165 i : DCHECK_EQ(pe::COFF_IMAGE, image_format_);
166 i : DCHECK(policy_object_.get() == NULL);
167 i : policy_object_.reset(new pe::CoffTransformPolicy());
168 i : return static_cast<pe::CoffTransformPolicy*>(policy_object_.get());
169 i : }
170 :
171 E : pe::PERelinker* InstrumenterWithAgent::GetPERelinker() {
172 E : DCHECK_EQ(pe::PE_IMAGE, image_format_);
173 E : DCHECK(relinker_object_.get() == NULL);
174 E : relinker_object_.reset(new pe::PERelinker(GetPETransformPolicy()));
175 E : return static_cast<pe::PERelinker*>(relinker_object_.get());
176 E : }
177 :
178 i : pe::CoffRelinker* InstrumenterWithAgent::GetCoffRelinker() {
179 i : DCHECK_EQ(pe::COFF_IMAGE, image_format_);
180 i : relinker_object_.reset(new pe::CoffRelinker(GetCoffTransformPolicy()));
181 i : return static_cast<pe::CoffRelinker*>(relinker_object_.get());
182 i : }
183 :
184 E : bool InstrumenterWithAgent::CreateRelinker() {
185 : // Get the image format by quickly inspecting the image. This logs verbosely
186 : // on failure.
187 E : if (!GetImageFormat(input_image_path_, &image_format_))
188 E : return false;
189 :
190 : // Check if the format is supported and bail if it isn't.
191 E : if (!ImageFormatIsSupported(image_format_)) {
192 i : LOG(ERROR) << "Instrumenter \"" << InstrumentationMode()
193 : << "\" does not support input image format.";
194 i : return false;
195 : }
196 :
197 : // Create and setup an image format specific relinker.
198 E : if (image_format_ == pe::COFF_IMAGE) {
199 E : pe::CoffRelinker* relinker = GetCoffRelinker();
200 E : DCHECK(relinker != NULL);
201 E : relinker_ = relinker;
202 E : relinker->set_input_path(input_image_path_);
203 E : relinker->set_output_path(output_image_path_);
204 E : relinker->set_allow_overwrite(allow_overwrite_);
205 E : } else {
206 E : pe::PERelinker* relinker = GetPERelinker();
207 E : DCHECK(relinker != NULL);
208 E : relinker_ = relinker;
209 E : relinker->set_input_path(input_image_path_);
210 E : relinker->set_input_pdb_path(input_pdb_path_);
211 E : relinker->set_output_path(output_image_path_);
212 E : relinker->set_output_pdb_path(output_pdb_path_);
213 E : relinker->set_allow_overwrite(allow_overwrite_);
214 E : relinker->set_augment_pdb(!no_augment_pdb_);
215 E : relinker->set_use_old_decomposer(old_decomposer_);
216 E : relinker->set_strip_strings(!no_strip_strings_);
217 : }
218 :
219 E : return true;
220 E : }
221 :
222 : } // namespace instrumenters
223 : } // namespace instrument
|