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/instrument/transforms/entry_thunk_transform.h"
16 :
17 : #include "base/logging.h"
18 : #include "base/stringprintf.h"
19 : #include "syzygy/block_graph/block_builder.h"
20 : #include "syzygy/block_graph/block_util.h"
21 : #include "syzygy/common/defs.h"
22 : #include "syzygy/pe/pe_utils.h"
23 : #include "syzygy/pe/transforms/pe_add_imports_transform.h"
24 :
25 : namespace instrument {
26 : namespace transforms {
27 :
28 : using block_graph::BasicBlock;
29 : using block_graph::BasicCodeBlock;
30 : using block_graph::BasicBlockAssembler;
31 : using block_graph::BasicBlockReference;
32 : using block_graph::BasicBlockSubGraph;
33 : using block_graph::BlockBuilder;
34 : using block_graph::BlockGraph;
35 : using block_graph::Displacement;
36 : using block_graph::Operand;
37 : using block_graph::TransformPolicyInterface;
38 : using pe::transforms::PEAddImportsTransform;
39 :
40 : typedef pe::transforms::ImportedModule ImportedModule;
41 :
42 : const char EntryThunkTransform::kTransformName[] =
43 : "EntryThunkTransform";
44 :
45 : const char EntryThunkTransform::kEntryHookName[] = "_indirect_penter";
46 : const char EntryThunkTransform::kDllMainEntryHookName[] =
47 : "_indirect_penter_dllmain";
48 : const char EntryThunkTransform::kExeMainEntryHookName[] =
49 : "_indirect_penter_exemain";
50 : const char EntryThunkTransform::kDefaultInstrumentDll[] =
51 : "call_trace_client.dll";
52 :
53 : EntryThunkTransform::EntryThunkTransform()
54 : : thunk_section_(NULL),
55 : instrument_unsafe_references_(true),
56 : src_ranges_for_thunks_(false),
57 : only_instrument_module_entry_(false),
58 E : instrument_dll_name_(kDefaultInstrumentDll) {
59 E : }
60 :
61 E : bool EntryThunkTransform::SetEntryThunkParameter(const Immediate& immediate) {
62 : if (immediate.size() != core::kSizeNone &&
63 E : immediate.size() != core::kSize32Bit) {
64 E : return false;
65 : }
66 E : entry_thunk_parameter_ = immediate;
67 E : return true;
68 E : }
69 :
70 : bool EntryThunkTransform::SetFunctionThunkParameter(
71 E : const Immediate& immediate) {
72 : if (immediate.size() != core::kSizeNone &&
73 E : immediate.size() != core::kSize32Bit) {
74 E : return false;
75 : }
76 E : function_thunk_parameter_ = immediate;
77 E : return true;
78 E : }
79 :
80 E : bool EntryThunkTransform::EntryThunkIsParameterized() const {
81 E : return entry_thunk_parameter_.size() != core::kSizeNone;
82 E : }
83 :
84 E : bool EntryThunkTransform::FunctionThunkIsParameterized() const {
85 E : return function_thunk_parameter_.size() != core::kSizeNone;
86 E : }
87 :
88 : bool EntryThunkTransform::PreBlockGraphIteration(
89 : const TransformPolicyInterface* policy,
90 : BlockGraph* block_graph,
91 E : BlockGraph::Block* header_block) {
92 E : DCHECK(policy != NULL);
93 E : DCHECK(block_graph != NULL);
94 E : DCHECK(header_block != NULL);
95 E : DCHECK(thunk_section_ == NULL);
96 :
97 E : if (!GetEntryPoints(header_block))
98 i : return false;
99 :
100 E : ImportedModule import_module(instrument_dll_name_);
101 :
102 : // We import the minimal set of symbols necessary, depending on the types of
103 : // entry points we find in the module. We maintain a list of symbol indices/
104 : // reference pointers, which will be traversed after the import to populate
105 : // the references.
106 : typedef std::pair<size_t, BlockGraph::Reference*> ImportHook;
107 E : std::vector<ImportHook> import_hooks;
108 :
109 : // If there are any DllMain-like entry points (TLS initializers or DllMain
110 : // itself) then we need the DllMain entry hook.
111 E : if (dllmain_entrypoints_.size() > 0) {
112 : import_hooks.push_back(std::make_pair(
113 : import_module.AddSymbol(kDllMainEntryHookName,
114 : ImportedModule::kAlwaysImport),
115 E : &hook_dllmain_ref_));
116 : }
117 :
118 : // If this was an EXE then we need the EXE entry hook.
119 E : if (exe_entry_point_.first != NULL) {
120 : import_hooks.push_back(std::make_pair(
121 : import_module.AddSymbol(kExeMainEntryHookName,
122 : ImportedModule::kAlwaysImport),
123 E : &hook_exe_entry_ref_));
124 : }
125 :
126 : // If we're not only instrumenting module entry then we need the function
127 : // entry hook.
128 E : if (!only_instrument_module_entry_) {
129 : import_hooks.push_back(std::make_pair(
130 : import_module.AddSymbol(kEntryHookName,
131 : ImportedModule::kAlwaysImport),
132 E : &hook_ref_));
133 : }
134 :
135 : // Nothing to do if we don't need any import hooks.
136 E : if (import_hooks.empty())
137 E : return true;
138 :
139 : // Run the transform.
140 E : PEAddImportsTransform add_imports_transform;
141 E : add_imports_transform.AddModule(&import_module);
142 : if (!add_imports_transform.TransformBlockGraph(
143 E : policy, block_graph, header_block)) {
144 i : LOG(ERROR) << "Unable to add imports for instrumentation DLL.";
145 i : return false;
146 : }
147 :
148 : // Get references to each of the imported symbols.
149 E : for (size_t i = 0; i < import_hooks.size(); ++i) {
150 : if (!import_module.GetSymbolReference(import_hooks[i].first,
151 E : import_hooks[i].second)) {
152 i : LOG(ERROR) << "Unable to get reference to import.";
153 i : return false;
154 : }
155 E : }
156 :
157 : // Find or create the section we put our thunks in.
158 : thunk_section_ = block_graph->FindOrAddSection(common::kThunkSectionName,
159 E : pe::kCodeCharacteristics);
160 E : DCHECK(thunk_section_ != NULL);
161 :
162 E : return true;
163 E : }
164 :
165 : bool EntryThunkTransform::OnBlock(const TransformPolicyInterface* policy,
166 : BlockGraph* block_graph,
167 E : BlockGraph::Block* block) {
168 E : DCHECK(block != NULL);
169 :
170 E : if (block->type() != BlockGraph::CODE_BLOCK)
171 E : return true;
172 :
173 E : return InstrumentCodeBlock(block_graph, block);
174 E : }
175 :
176 : bool EntryThunkTransform::InstrumentCodeBlock(
177 E : BlockGraph* block_graph, BlockGraph::Block* block) {
178 E : DCHECK(block_graph != NULL);
179 E : DCHECK(block != NULL);
180 :
181 : // Typedef for the thunk block map. The key is the offset within the callee
182 : // block and the value is the thunk block that forwards to the callee at that
183 : // offset.
184 E : ThunkBlockMap thunk_block_map;
185 :
186 : // Iterate through all the block's referrers, creating thunks as we go.
187 : // We copy the referrer set for simplicity, as it's potentially mutated
188 : // in the loop.
189 E : BlockGraph::Block::ReferrerSet referrers = block->referrers();
190 E : BlockGraph::Block::ReferrerSet::const_iterator referrer_it(referrers.begin());
191 E : for (; referrer_it != referrers.end(); ++referrer_it) {
192 E : const BlockGraph::Block::Referrer& referrer = *referrer_it;
193 : if (!InstrumentCodeBlockReferrer(
194 E : referrer, block_graph, block, &thunk_block_map)) {
195 i : return false;
196 : }
197 E : }
198 :
199 E : return true;
200 E : }
201 :
202 : bool EntryThunkTransform::InstrumentCodeBlockReferrer(
203 : const BlockGraph::Block::Referrer& referrer,
204 : BlockGraph* block_graph,
205 : BlockGraph::Block* block,
206 E : ThunkBlockMap* thunk_block_map) {
207 E : DCHECK(block_graph != NULL);
208 E : DCHECK(block != NULL);
209 E : DCHECK(thunk_block_map != NULL);
210 :
211 : // Get the reference.
212 E : BlockGraph::Reference ref;
213 E : if (!referrer.first->GetReference(referrer.second, &ref)) {
214 i : LOG(ERROR) << "Unable to get reference from referrer.";
215 i : return false;
216 : }
217 :
218 : // Skip self-references, except long references to the start of the block.
219 : // TODO(siggi): This needs refining, as it may currently miss important
220 : // cases. Notably if a block contains more than one function, and the
221 : // functions are mutually recursive, we'll only record the original
222 : // entry to the block, but will miss the internal recursion.
223 : // As-is, this does work for the common case where a block contains
224 : // one self-recursive function, however.
225 E : if (referrer.first == block) {
226 : // Skip short references.
227 E : if (ref.size() < sizeof(core::AbsoluteAddress))
228 E : return true;
229 :
230 : // Skip interior references. The rationale for this is because these
231 : // references will tend to be switch tables, and we don't need the
232 : // overhead of instrumenting and recording all switch statement executions
233 : // for now.
234 E : if (ref.offset() != 0)
235 E : return true;
236 : }
237 :
238 : // See whether this is one of the DLL entrypoints.
239 E : pe::EntryPoint entry(ref.referenced(), ref.offset());
240 : pe::EntryPointSet::const_iterator entry_it(dllmain_entrypoints_.find(
241 E : entry));
242 E : bool is_dllmain_entry = entry_it != dllmain_entrypoints_.end();
243 :
244 : // Determine if this is an EXE entry point.
245 E : bool is_exe_entry = entry == exe_entry_point_;
246 :
247 : // It can't be both an EXE and a DLL entry.
248 E : DCHECK(!is_dllmain_entry || !is_exe_entry);
249 :
250 : // If we're only instrumenting entry points and this isn't one, then skip it.
251 E : if (only_instrument_module_entry_ && !is_dllmain_entry && !is_exe_entry)
252 E : return true;
253 :
254 : if (!instrument_unsafe_references_ &&
255 E : block_graph::IsUnsafeReference(referrer.first, ref)) {
256 E : LOG(INFO) << "Skipping reference between unsafe block pair '"
257 : << referrer.first->name() << "' and '"
258 : << block->name() << "'";
259 E : return true;
260 : }
261 :
262 : // Determine which hook function to use.
263 E : BlockGraph::Reference* hook_ref = &hook_ref_;
264 E : if (is_dllmain_entry)
265 E : hook_ref = &hook_dllmain_ref_;
266 E : else if (is_exe_entry)
267 E : hook_ref = &hook_exe_entry_ref_;
268 E : DCHECK(hook_ref->referenced() != NULL);
269 :
270 : // Determine which parameter to use, if any.
271 E : const Immediate* param = NULL;
272 E : if ((is_dllmain_entry || is_exe_entry) && EntryThunkIsParameterized()) {
273 E : param = &entry_thunk_parameter_;
274 E : } else if (FunctionThunkIsParameterized()) {
275 E : param = &function_thunk_parameter_;
276 : }
277 :
278 : // Look for the reference in the thunk block map, and only create a new one
279 : // if it does not already exist.
280 E : BlockGraph::Block* thunk_block = NULL;
281 E : ThunkBlockMap::const_iterator thunk_it = thunk_block_map->find(ref.offset());
282 E : if (thunk_it == thunk_block_map->end()) {
283 E : thunk_block = CreateOneThunk(block_graph, ref, *hook_ref, param);
284 E : if (thunk_block == NULL) {
285 i : LOG(ERROR) << "Unable to create thunk block.";
286 i : return false;
287 : }
288 E : (*thunk_block_map)[ref.offset()] = thunk_block;
289 E : } else {
290 E : thunk_block = thunk_it->second;
291 : }
292 E : DCHECK(thunk_block != NULL);
293 :
294 : // Update the referrer to point to the thunk.
295 : BlockGraph::Reference new_ref(ref.type(),
296 : ref.size(),
297 : thunk_block,
298 E : 0, 0);
299 E : referrer.first->SetReference(referrer.second, new_ref);
300 :
301 E : return true;
302 E : }
303 :
304 : BlockGraph::Block* EntryThunkTransform::CreateOneThunk(
305 : BlockGraph* block_graph,
306 : const BlockGraph::Reference& destination,
307 : const BlockGraph::Reference& hook,
308 E : const Immediate* parameter) {
309 E : std::string name;
310 E : if (destination.offset() == 0) {
311 : name = base::StringPrintf("%s%s",
312 : destination.referenced()->name().c_str(),
313 E : common::kThunkSuffix);
314 E : } else {
315 : name = base::StringPrintf("%s%s+%d",
316 : destination.referenced()->name().c_str(),
317 : common::kThunkSuffix,
318 E : destination.offset());
319 : }
320 :
321 : // Set up a basic block subgraph containing a single block description, with
322 : // that block description containing a single empty basic block, and get an
323 : // assembler writing into that basic block.
324 : // TODO(chrisha): Make this reusable somehow. Creating a code block via an
325 : // assembler is likely to be pretty common.
326 E : BasicBlockSubGraph bbsg;
327 : BasicBlockSubGraph::BlockDescription* block_desc = bbsg.AddBlockDescription(
328 : name,
329 : NULL,
330 : BlockGraph::CODE_BLOCK,
331 : thunk_section_->id(),
332 : 1,
333 E : 0);
334 E : BasicCodeBlock* bb = bbsg.AddBasicCodeBlock(name);
335 E : block_desc->basic_block_order.push_back(bb);
336 : BasicBlockAssembler assm(bb->instructions().begin(),
337 E : &bb->instructions());
338 :
339 : // Set up our thunk:
340 : // 1. push parameter
341 : // 2. push func_addr
342 : // 3. jmp hook_addr
343 E : if (parameter != NULL)
344 E : assm.push(*parameter);
345 E : assm.push(Immediate(destination.referenced(), destination.offset()));
346 E : assm.jmp(Operand(Displacement(hook.referenced(), hook.offset())));
347 :
348 : // Condense the whole mess into a block.
349 E : BlockBuilder block_builder(block_graph);
350 E : if (!block_builder.Merge(&bbsg)) {
351 i : LOG(ERROR) << "Failed to build thunk block.";
352 i : return NULL;
353 : }
354 :
355 : // Exactly one new block should have been created.
356 E : DCHECK_EQ(1u, block_builder.new_blocks().size());
357 E : BlockGraph::Block* thunk = block_builder.new_blocks().front();
358 :
359 E : if (src_ranges_for_thunks_) {
360 : // Give the thunk a source range synonymous with the destination.
361 : // That way the debugger will resolve calls and jumps to the thunk to the
362 : // destination function's name, which makes the assembly much easier to
363 : // read. The downside to this is that the symbols are now no longer unique,
364 : // and searching for a function by name may turn up either the function or
365 : // the thunk.
366 : const BlockGraph::Block::SourceRanges& source_ranges =
367 E : destination.referenced()->source_ranges();
368 : const BlockGraph::Block::SourceRanges::RangePair* source =
369 E : source_ranges.FindRangePair(destination.offset(), thunk->size());
370 E : if (source != NULL) {
371 : // Calculate the offset into the range.
372 E : size_t offs = destination.offset() - source->first.start();
373 E : BlockGraph::Block::DataRange data(0, thunk->size());
374 : BlockGraph::Block::SourceRange src(source->second.start() + offs,
375 E : thunk->size());
376 E : bool pushed = thunk->source_ranges().Push(data, src);
377 E : DCHECK(pushed);
378 : }
379 : }
380 :
381 E : return thunk;
382 E : }
383 :
384 E : bool EntryThunkTransform::GetEntryPoints(BlockGraph::Block* header_block) {
385 : // Get the TLS initializer entry-points. These have the same signature and
386 : // call patterns to DllMain.
387 E : if (!pe::GetTlsInitializers(header_block, &dllmain_entrypoints_)) {
388 i : LOG(ERROR) << "Failed to populate the TLS Initializer entry-points.";
389 i : return false;
390 : }
391 :
392 : // Get the DLL entry-point.
393 E : pe::EntryPoint dll_entry_point;
394 E : if (!pe::GetDllEntryPoint(header_block, &dll_entry_point)) {
395 i : LOG(ERROR) << "Failed to resolve the DLL entry-point.";
396 i : return false;
397 : }
398 :
399 : // If the image is an EXE or is a DLL that does not specify an entry-point
400 : // (the entry-point is optional for DLLs) then the dll_entry_point will have
401 : // a NULL block pointer. Otherwise, add it to the entry-point set.
402 E : if (dll_entry_point.first != NULL) {
403 E : dllmain_entrypoints_.insert(dll_entry_point);
404 E : } else {
405 : // Get the EXE entry point. We only need to bother looking if we didn't get
406 : // a DLL entry point, as we can't have both.
407 E : if (!pe::GetExeEntryPoint(header_block, &exe_entry_point_)) {
408 i : LOG(ERROR) << "Failed to resolve the EXE entry-point.";
409 i : return false;
410 : }
411 : }
412 :
413 E : return true;
414 E : }
415 :
416 : } // namespace transforms
417 : } // namespace instrument
|