1 : // Copyright 2012 Google Inc.
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 : // Defines the InstrumentApp class, which implements the command-line
16 : // "instrument" tool.
17 :
18 : #include "syzygy/instrument/instrument_app.h"
19 :
20 : #include <algorithm>
21 : #include <iostream>
22 :
23 : #include "base/string_util.h"
24 : #include "base/stringprintf.h"
25 : #include "syzygy/instrument/mutators/add_bb_ranges_stream.h"
26 : #include "syzygy/instrument/transforms/asan_transform.h"
27 : #include "syzygy/instrument/transforms/coverage_transform.h"
28 : #include "syzygy/instrument/transforms/entry_thunk_transform.h"
29 : #include "syzygy/instrument/transforms/thunk_import_references_transform.h"
30 : #include "syzygy/pe/decomposer.h"
31 : #include "syzygy/pe/pe_relinker.h"
32 :
33 : namespace instrument {
34 :
35 : namespace {
36 :
37 : using block_graph::BlockGraph;
38 : using pe::Decomposer;
39 : using pe::PEFile;
40 :
41 : static const char kUsageFormatStr[] =
42 : "Usage: %ls [options]\n"
43 : " Required arguments:\n"
44 : " --input-image=<path> The input image to instrument.\n"
45 : " --mode=ASAN|CALLTRACE|COVERAGE|PROFILER\n"
46 : " Specifies which instrumentation mode is to be\n"
47 : " used. If this is not specified it is equivalent\n"
48 : " to specifying --mode=CALLTRACE (this default\n"
49 : " behaviour is DEPRECATED.\n"
50 : " --output-image=<path>\n"
51 : " The instrumented output image.\n"
52 : " DEPRECATED options:\n"
53 : " --input-dll is aliased to --input-image.\n"
54 : " --output-dll is aliased to --output-image.\n"
55 : " --call-trace-client=RPC\n"
56 : " Equivalent to --mode=CALLTRACE.\n"
57 : " --call-trace-client=PROFILER\n"
58 : " Equivalent to --mode=PROFILER.\n"
59 : " --call-trace-client=<path>\n"
60 : " Equivalent to --mode=CALLTRACE --agent=<path>.\n"
61 : " General options (applicable in all modes):\n"
62 : " --agent=<path> If specified indicates exactly which DLL should\n"
63 : " be used in instrumenting the provided module.\n"
64 : " If not specified a default agent library will\n"
65 : " be used. This is ignored in ASAN mode.\n"
66 : " --debug-friendly Generate more debugger friendly output by\n"
67 : " making the thunks resolve to the original\n"
68 : " function's name. This is at the cost of the\n"
69 : " uniqueness of address->name resolution.\n"
70 : " --input-pdb=<path> The PDB for the DLL to instrument. If not\n"
71 : " explicitly provided will be searched for.\n"
72 : " --no-augment-pdb Indicates that the relinker should not augment\n"
73 : " the output PDB with additional metadata.\n"
74 : " --no-strip-strings Indicates that the relinker should not strip\n"
75 : " the strings when augmenting the PDB. They are\n"
76 : " stripped by default to keep PDB sizes down.\n"
77 : " --output-pdb=<path> The PDB for the instrumented DLL. If not\n"
78 : " provided will attempt to generate one.\n"
79 : " --overwrite Allow output files to be overwritten.\n"
80 : " CALLTRACE mode options:\n"
81 : " --instrument-imports Also instrument calls to imports.\n"
82 : " --module-entry-only If specified then the per-function entry hook\n"
83 : " will not be used and only module entry points\n"
84 : " will be hooked.\n"
85 : " --no-unsafe-refs Perform no instrumentation of references\n"
86 : " between code blocks that contain anything but\n"
87 : " C/C++.\n"
88 : " PROFILE mode options:\n"
89 : " --instrument-imports Also instrument calls to imports.\n"
90 : "\n";
91 :
92 : } // namespace
93 :
94 : const char InstrumentApp::kCallTraceClientDllCoverage[] = "coverage_client.dll";
95 : const char InstrumentApp::kCallTraceClientDllProfiler[] = "profile_client.dll";
96 : const char InstrumentApp::kCallTraceClientDllRpc[] = "call_trace_client.dll";
97 :
98 i : pe::PERelinker& InstrumentApp::GetRelinker() {
99 i : if (relinker_.get() == NULL) {
100 i : relinker_.reset(new pe::PERelinker());
101 i : CHECK(relinker_.get() != NULL);
102 : }
103 i : return *(relinker_.get());
104 i : }
105 :
106 E : void InstrumentApp::ParseDeprecatedMode(const CommandLine* cmd_line) {
107 E : DCHECK(cmd_line != NULL);
108 :
109 E : std::string client = cmd_line->GetSwitchValueASCII("call-trace-client");
110 :
111 E : if (client.empty()) {
112 E : LOG(INFO) << "DEPRECATED: No mode specified, using --mode=CALLTRACE.";
113 E : mode_ = kInstrumentCallTraceMode;
114 E : client_dll_ = kCallTraceClientDllRpc;
115 E : return;
116 : }
117 :
118 E : if (LowerCaseEqualsASCII(client, "profiler")) {
119 E : LOG(INFO) << "DEPRECATED: Using --mode=PROFILER.";
120 E : mode_ = kInstrumentProfilerMode;
121 E : client_dll_ = kCallTraceClientDllProfiler;
122 E : } else if (LowerCaseEqualsASCII(client, "rpc")) {
123 E : LOG(INFO) << "DEPRECATED: Using --mode=CALLTRACE.";
124 E : mode_ = kInstrumentCallTraceMode;
125 E : client_dll_ = kCallTraceClientDllRpc;
126 E : } else {
127 E : LOG(INFO) << "DEPRECATED: Using --mode=CALLTRACE --agent=" << client << ".";
128 E : mode_ = kInstrumentCallTraceMode;
129 E : client_dll_ = client;
130 : }
131 E : }
132 :
133 E : bool InstrumentApp::ParseCommandLine(const CommandLine* cmd_line) {
134 E : DCHECK(cmd_line != NULL);
135 :
136 E : if (cmd_line->HasSwitch("help"))
137 E : return Usage(cmd_line, "");
138 :
139 : // TODO(chrisha): Simplify the input/output image parsing once external
140 : // tools have been updated.
141 :
142 : // Parse the input image.
143 E : if (cmd_line->HasSwitch("input-dll")) {
144 E : LOG(WARNING) << "DEPRECATED: Using --input-dll.";
145 E : input_dll_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("input-dll"));
146 E : } else {
147 E : input_dll_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("input-image"));
148 : }
149 :
150 : // Parse the output image.
151 E : if (cmd_line->HasSwitch("output-dll")) {
152 E : LOG(WARNING) << "DEPRECATED: Using --output-dll.";
153 E : output_dll_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("output-dll"));
154 E : } else {
155 : output_dll_path_ = AbsolutePath(cmd_line->GetSwitchValuePath(
156 E : "output-image"));
157 : }
158 :
159 : // Ensure that both input and output have been specified.
160 E : if (input_dll_path_.empty() || output_dll_path_.empty())
161 E : return Usage(cmd_line, "You must provide input and output file names.");
162 :
163 : // Get the mode and the default client DLL.
164 E : if (!cmd_line->HasSwitch("mode")) {
165 : // TODO(chrisha): Remove this once build scripts and profiling tools have
166 : // been updated.
167 E : ParseDeprecatedMode(cmd_line);
168 E : } else {
169 E : std::string mode = cmd_line->GetSwitchValueASCII("mode");
170 E : if (LowerCaseEqualsASCII(mode, "asan")) {
171 E : mode_ = kInstrumentAsanMode;
172 E : } else if (LowerCaseEqualsASCII(mode, "calltrace")) {
173 E : mode_ = kInstrumentCallTraceMode;
174 E : client_dll_ = kCallTraceClientDllRpc;
175 E : } else if (LowerCaseEqualsASCII(mode, "coverage")) {
176 E : mode_ = kInstrumentCoverageMode;
177 E : client_dll_ = kCallTraceClientDllCoverage;
178 E : } else if (LowerCaseEqualsASCII(mode, "profiler")) {
179 E : mode_ = kInstrumentProfilerMode;
180 E : client_dll_ = kCallTraceClientDllProfiler;
181 E : } else {
182 : return Usage(cmd_line,
183 : base::StringPrintf("Unknown instrumentation mode: %s.",
184 i : mode.c_str()).c_str());
185 : }
186 :
187 E : LOG(INFO) << "Default agent for mode " << mode << " is \""
188 : << client_dll_ << "\".";
189 E : }
190 E : DCHECK_NE(kInstrumentInvalidMode, mode_);
191 :
192 : // Parse the custom agent if one is specified.
193 E : if (cmd_line->HasSwitch("agent")) {
194 E : if (mode_ == kInstrumentAsanMode) {
195 : // TODO(siggi): Make this work properly!
196 E : LOG(WARNING) << "Ignoring --agent in ASAN mode.";
197 E : } else {
198 E : client_dll_ = cmd_line->GetSwitchValueASCII("agent");
199 E : LOG(INFO) << "Got custom agent \"" << client_dll_ << "\".";
200 : }
201 : }
202 :
203 : // Parse the remaining command line arguments. Not all of these are valid in
204 : // all modes, but we don't care too much about ignored arguments.
205 E : input_pdb_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("input-pdb"));
206 E : output_pdb_path_ = AbsolutePath(cmd_line->GetSwitchValuePath("output-pdb"));
207 E : allow_overwrite_ = cmd_line->HasSwitch("overwrite");
208 E : no_augment_pdb_ = cmd_line->HasSwitch("no-augment-pdb");
209 E : no_strip_strings_ = cmd_line->HasSwitch("no-strip-strings");
210 E : debug_friendly_ = cmd_line->HasSwitch("debug-friendly");
211 E : thunk_imports_ = cmd_line->HasSwitch("instrument-imports");
212 E : instrument_unsafe_references_ = !cmd_line->HasSwitch("no-unsafe-refs");
213 E : module_entry_only_ = cmd_line->HasSwitch("module-entry-only");
214 :
215 :
216 : // Set per-mode overrides as necessary.
217 E : switch (mode_) {
218 : case kInstrumentCoverageMode: {
219 E : thunk_imports_ = false;
220 E : instrument_unsafe_references_ = false;
221 E : module_entry_only_ = true;
222 E : } break;
223 :
224 : case kInstrumentProfilerMode: {
225 E : instrument_unsafe_references_ = false;
226 E : module_entry_only_ = false;
227 : } break;
228 :
229 : default: break;
230 : }
231 :
232 E : return true;
233 E : }
234 :
235 E : int InstrumentApp::Run() {
236 E : DCHECK_NE(kInstrumentInvalidMode, mode_);
237 :
238 E : pe::PERelinker& relinker = GetRelinker();
239 E : relinker.set_input_path(input_dll_path_);
240 E : relinker.set_input_pdb_path(input_pdb_path_);
241 E : relinker.set_output_path(output_dll_path_);
242 E : relinker.set_output_pdb_path(output_pdb_path_);
243 E : relinker.set_allow_overwrite(allow_overwrite_);
244 E : relinker.set_augment_pdb(!no_augment_pdb_);
245 E : relinker.set_strip_strings(!no_strip_strings_);
246 :
247 : // Initialize the relinker. This does the decomposition, etc.
248 E : if (!relinker.Init()) {
249 E : LOG(ERROR) << "Failed to initialize relinker.";
250 E : return 1;
251 : }
252 :
253 : // A list of all possible transforms that we will need.
254 E : scoped_ptr<instrument::transforms::AsanTransform> asan_transform;
255 E : scoped_ptr<instrument::transforms::EntryThunkTransform> entry_thunk_tx;
256 : scoped_ptr<instrument::transforms::ThunkImportReferencesTransform>
257 E : import_thunk_tx;
258 : scoped_ptr<instrument::transforms::CoverageInstrumentationTransform>
259 E : coverage_tx;
260 : scoped_ptr<instrument::mutators::AddBasicBlockRangesStreamPdbMutator>
261 E : add_bb_addr_stream_mutator;
262 :
263 : // We are instrumenting in ASAN mode.
264 E : if (mode_ == kInstrumentAsanMode) {
265 i : asan_transform.reset(new instrument::transforms::AsanTransform);
266 i : relinker.AppendTransform(asan_transform.get());
267 i : } else {
268 : // We're either in call-trace mode, coverage or profiler mode. Each of these
269 : // use the entry_thunk_tx, so we handle them in the same manner.
270 :
271 : // If we're in coverage mode, we need to add coverage structures to
272 : // the image and we need to augment the PDB file with the basic block
273 : // addresses.
274 E : if (mode_ == kInstrumentCoverageMode) {
275 : coverage_tx.reset(
276 i : new instrument::transforms::CoverageInstrumentationTransform);
277 i : relinker.AppendTransform(coverage_tx.get());
278 :
279 : add_bb_addr_stream_mutator.reset(
280 : new instrument::mutators::AddBasicBlockRangesStreamPdbMutator(
281 i : coverage_tx->bb_ranges()));
282 i : relinker.AppendPdbMutator(add_bb_addr_stream_mutator.get());
283 : }
284 :
285 : // Set up the entry thunk instrumenting transform and add it to the
286 : // relinker.
287 E : entry_thunk_tx.reset(new instrument::transforms::EntryThunkTransform);
288 E : entry_thunk_tx->set_instrument_dll_name(client_dll_);
289 : entry_thunk_tx->set_instrument_unsafe_references(
290 E : instrument_unsafe_references_);
291 E : entry_thunk_tx->set_src_ranges_for_thunks(debug_friendly_);
292 E : entry_thunk_tx->set_only_instrument_module_entry(module_entry_only_);
293 E : relinker.AppendTransform(entry_thunk_tx.get());
294 :
295 : // If we are thunking imports then add the appropriate transform.
296 E : if (thunk_imports_) {
297 : import_thunk_tx.reset(
298 i : new instrument::transforms::ThunkImportReferencesTransform);
299 i : import_thunk_tx->ExcludeModule(client_dll_);
300 i : relinker.AppendTransform(import_thunk_tx.get());
301 : }
302 : }
303 :
304 : // We let the PERelinker use the implicit OriginalOrderer.
305 E : if (!relinker.Relink()) {
306 E : LOG(ERROR) << "Unable to relink input image.";
307 E : return 1;
308 : }
309 :
310 E : return 0;
311 E : }
312 :
313 : bool InstrumentApp::Usage(const CommandLine* cmd_line,
314 E : const base::StringPiece& message) const {
315 E : if (!message.empty()) {
316 E : ::fwrite(message.data(), 1, message.length(), err());
317 E : ::fprintf(err(), "\n\n");
318 : }
319 :
320 : ::fprintf(err(),
321 : kUsageFormatStr,
322 E : cmd_line->GetProgram().BaseName().value().c_str());
323 :
324 E : return false;
325 E : }
326 :
327 : } // namespace instrument
|