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