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/transforms/entry_call_transform.h"
16 :
17 : #include "base/logging.h"
18 : #include "base/strings/stringprintf.h"
19 : #include "syzygy/block_graph/basic_block_assembler.h"
20 : #include "syzygy/block_graph/block_builder.h"
21 : #include "syzygy/block_graph/block_util.h"
22 : #include "syzygy/common/defs.h"
23 : #include "syzygy/pe/pe_utils.h"
24 : #include "syzygy/pe/transforms/pe_add_imports_transform.h"
25 :
26 : namespace instrument {
27 : namespace transforms {
28 :
29 : const char EntryCallBasicBlockTransform::kTransformName[] =
30 : "EntryCallBasicBlockTransform";
31 :
32 : const char EntryCallTransform::kTransformName[] =
33 : "EntryCallTransform";
34 :
35 : const char EntryCallTransform::kEntryHookName[] = "_indirect_penter";
36 : const char EntryCallTransform::kDllMainEntryHookName[] =
37 : "_indirect_penter_dllmain";
38 : const char EntryCallTransform::kExeMainEntryHookName[] =
39 : "_indirect_penter_exemain";
40 : const char EntryCallTransform::kDefaultInstrumentDll[] =
41 : "profile_client.dll";
42 :
43 : EntryCallBasicBlockTransform::EntryCallBasicBlockTransform(
44 : const BlockGraph::Reference& hook_reference, bool debug_friendly)
45 E : : hook_reference_(hook_reference), debug_friendly_(debug_friendly) {
46 E : }
47 :
48 : bool EntryCallBasicBlockTransform::TransformBasicBlockSubGraph(
49 : const TransformPolicyInterface* policy,
50 : BlockGraph* block_graph,
51 E : BasicBlockSubGraph* basic_block_subgraph) {
52 E : DCHECK_NE(static_cast<TransformPolicyInterface*>(NULL), policy);
53 E : DCHECK_NE(static_cast<BlockGraph*>(NULL), block_graph);
54 E : DCHECK_NE(static_cast<BasicBlockSubGraph*>(NULL), basic_block_subgraph);
55 :
56 : using block_graph::BasicCodeBlock;
57 : typedef block_graph::BasicBlockSubGraph::BBCollection BBCollection;
58 :
59 : // We expect to be looking into a newly-decomposed basic block graph, with
60 : // precisely one block description for the originating block.
61 E : DCHECK_EQ(1U, basic_block_subgraph->block_descriptions().size());
62 : BasicBlockSubGraph::BasicBlockOrdering& bb_order =
63 E : basic_block_subgraph->block_descriptions().front().basic_block_order;
64 :
65 : // An empty BB ordering is nonsensical.
66 E : DCHECK_NE(0U, bb_order.size());
67 :
68 : // Cast the first block to a code block - this should always succeed
69 : // for code coming from MSVC, but we do a runtime check for proper
70 : // belt-and-suspenders.
71 E : BasicCodeBlock* bb = BasicCodeBlock::Cast(bb_order.front());
72 E : if (bb == NULL) {
73 i : LOG(ERROR) << "No code at the head of function \""
74 : << basic_block_subgraph->original_block()->name()
75 : << "\"";
76 i : return false;
77 : }
78 :
79 E : DCHECK_NE(static_cast<BasicCodeBlock*>(NULL), bb);
80 E : DCHECK_EQ(0, bb->offset());
81 :
82 : // Create a new basic block for the entry hook.
83 : BasicCodeBlock* entry_hook =
84 E : basic_block_subgraph->AddBasicCodeBlock("EntryHook");
85 E : DCHECK_NE(static_cast<BasicCodeBlock*>(NULL), entry_hook);
86 :
87 : // Add a call instruction to the new block.
88 : using block_graph::BasicBlockAssembler;
89 : using block_graph::Operand;
90 : using block_graph::Displacement;
91 : block_graph::BasicBlockAssembler assm(entry_hook->instructions().begin(),
92 E : &entry_hook->instructions());
93 :
94 : // In debug friendly mode we assign the previously first instruction's
95 : // address to the inserted call.
96 E : if (debug_friendly_) {
97 E : if (bb->instructions().size() != 0) {
98 E : assm.set_source_range(bb->instructions().front().source_range());
99 E : } else {
100 i : LOG(WARNING) << "Function \""
101 : << basic_block_subgraph->original_block()->name()
102 : << "\" starts with an empty basic block. "
103 : << "Not inserting a source range for it.";
104 : }
105 : }
106 :
107 : assm.call(Operand(Displacement(hook_reference_.referenced(),
108 E : hook_reference_.offset())));
109 :
110 : // Put the new BB at the top of the function.
111 E : bb_order.push_front(entry_hook);
112 :
113 : // Nominate the original entry point BB as successor for the new block.
114 : using block_graph::Successor;
115 : using block_graph::BasicBlockReference;
116 : entry_hook->successors().push_back(
117 : Successor(Successor::kConditionTrue,
118 : BasicBlockReference(BlockGraph::PC_RELATIVE_REF, 4, bb),
119 E : 0));
120 :
121 : // Transfer the external referrers from the old head of function to the
122 : // entry hook.
123 E : bb->referrers().swap(entry_hook->referrers());
124 :
125 : // Now run through the code BBs in the function, and re-route any refs to the
126 : // former head of function to the entry hook. The point of this is to route
127 : // explicit self-recursion or self-references through the entry hook, while
128 : // leaving loops alone.
129 : // Loops will be implemented as either explicit control flow in successors,
130 : // or else may involve computed jumps through data "BBs", and by diverting
131 : // only instructions, we're sure to not divert loops through the entry hook.
132 : //
133 : // Note that this is not comprehensive, as it's in general impossible to
134 : // distinguish tail recursion elimination from a loop at the semantic
135 : // level of instructions.
136 : //
137 : // We choose to err on the side of performance and robustness, as
138 : // mis-instrumenting a loop will result in pushing the profiler's shadow
139 : // stack for every loop iteration, and then popping it as many times on exit.
140 : // This will lead to poor performance at best, but may also cause the
141 : // shadow stack to blow up in the extreme.
142 E : BasicBlockSubGraph::BasicBlockOrdering::iterator bb_iter = bb_order.begin();
143 :
144 : // Walk past the entry hook BB.
145 E : DCHECK_EQ(entry_hook, *bb_iter);
146 E : ++bb_iter;
147 :
148 : // Walk through all the BBs (in order).
149 E : for (; bb_iter != bb_order.end(); ++bb_iter) {
150 E : BasicCodeBlock* curr_block = BasicCodeBlock::Cast(*bb_iter);
151 E : if (curr_block != NULL) {
152 E : BasicCodeBlock::Instructions& instr = curr_block->instructions();
153 E : BasicCodeBlock::Instructions::iterator inst_it(instr.begin());
154 :
155 : // Walk through instructions for each code block.
156 E : for (; inst_it != instr.end(); ++inst_it) {
157 : block_graph::Instruction::BasicBlockReferenceMap&
158 E : refs(inst_it->references());
159 : block_graph::Instruction::BasicBlockReferenceMap::iterator
160 E : ref_it(refs.begin());
161 :
162 : // For each instruction, walk through refrences.
163 E : for (; ref_it != refs.end(); ++ref_it) {
164 E : BasicBlockReference& ref(ref_it->second);
165 E : if (ref.basic_block() == bb) {
166 : // And if the reference pointed to bb, redirect it to entry_hook.
167 : ref = BasicBlockReference(ref.reference_type(),
168 : ref.size(),
169 E : entry_hook);
170 : }
171 E : }
172 E : }
173 : }
174 E : }
175 :
176 E : return true;
177 E : }
178 :
179 : EntryCallTransform::EntryCallTransform(bool debug_friendly)
180 : : instrument_dll_name_(kDefaultInstrumentDll),
181 E : debug_friendly_(debug_friendly) {
182 E : }
183 :
184 : bool EntryCallTransform::PreBlockGraphIteration(
185 : const TransformPolicyInterface* policy,
186 : BlockGraph* block_graph,
187 E : BlockGraph::Block* header_block) {
188 E : DCHECK_NE(reinterpret_cast<TransformPolicyInterface*>(NULL), policy);
189 E : DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
190 E : DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), header_block);
191 E : DCHECK_EQ(BlockGraph::PE_IMAGE, block_graph->image_format());
192 :
193 E : if (!GetEntryPoints(header_block))
194 i : return false;
195 :
196 : using pe::transforms::ImportedModule;
197 E : ImportedModule import_module(instrument_dll_name_);
198 :
199 : // We import the minimal set of symbols necessary, depending on the types of
200 : // entry points we find in the module. We maintain a list of symbol indices/
201 : // reference pointers, which will be traversed after the import to populate
202 : // the references.
203 : typedef std::pair<size_t, BlockGraph::Reference*> ImportHook;
204 E : std::vector<ImportHook> import_hooks;
205 :
206 : // If there are any DllMain-like entry points (TLS initializers or DllMain
207 : // itself) then we need the DllMain entry hook.
208 E : if (dllmain_entrypoints_.size() > 0) {
209 : import_hooks.push_back(std::make_pair(
210 : import_module.AddSymbol(kDllMainEntryHookName,
211 : ImportedModule::kAlwaysImport),
212 E : &hook_dllmain_ref_));
213 : }
214 :
215 : // If this was an EXE then we need the EXE entry hook.
216 E : if (exe_entry_point_.first != NULL) {
217 : import_hooks.push_back(std::make_pair(
218 : import_module.AddSymbol(kExeMainEntryHookName,
219 : ImportedModule::kAlwaysImport),
220 i : &hook_exe_entry_ref_));
221 : }
222 :
223 : import_hooks.push_back(std::make_pair(
224 : import_module.AddSymbol(kEntryHookName,
225 : ImportedModule::kAlwaysImport),
226 E : &hook_ref_));
227 :
228 : // Nothing to do if we don't need any import hooks.
229 E : if (import_hooks.empty())
230 i : return true;
231 :
232 : // Run the transform.
233 E : pe::transforms::PEAddImportsTransform add_imports_transform;
234 E : add_imports_transform.AddModule(&import_module);
235 : if (!add_imports_transform.TransformBlockGraph(
236 E : policy, block_graph, header_block)) {
237 i : LOG(ERROR) << "Unable to add imports for instrumentation DLL.";
238 i : return false;
239 : }
240 :
241 : // Get references to each of the imported symbols.
242 E : for (size_t i = 0; i < import_hooks.size(); ++i) {
243 : if (!import_module.GetSymbolReference(import_hooks[i].first,
244 E : import_hooks[i].second)) {
245 i : LOG(ERROR) << "Unable to get reference to import.";
246 i : return false;
247 : }
248 E : }
249 :
250 E : return true;
251 E : }
252 :
253 : bool EntryCallTransform::OnBlock(const TransformPolicyInterface* policy,
254 : BlockGraph* block_graph,
255 E : BlockGraph::Block* block) {
256 E : DCHECK_NE(static_cast<TransformPolicyInterface*>(NULL), policy);
257 E : DCHECK_NE(static_cast<BlockGraph*>(NULL), block_graph);
258 E : DCHECK_NE(static_cast<BlockGraph::Block*>(NULL), block);
259 :
260 : // Skip blocks that aren't eligible for basic-block decomposition.
261 E : if (!policy->BlockIsSafeToBasicBlockDecompose(block))
262 E : return true;
263 :
264 : // Apply the basic block transform.
265 : // See whether this is one of the DLL entrypoints.
266 E : pe::EntryPoint entry(block, 0);
267 : pe::EntryPointSet::const_iterator entry_it(dllmain_entrypoints_.find(
268 E : entry));
269 E : bool is_dllmain_entry = entry_it != dllmain_entrypoints_.end();
270 :
271 : // Determine if this is an EXE entry point.
272 E : bool is_exe_entry = entry == exe_entry_point_;
273 :
274 : // It can't be both an EXE and a DLL entry.
275 E : DCHECK(!is_dllmain_entry || !is_exe_entry);
276 :
277 : // Determine which hook function to use.
278 E : BlockGraph::Reference* hook_ref = &hook_ref_;
279 E : if (is_dllmain_entry)
280 E : hook_ref = &hook_dllmain_ref_;
281 E : else if (is_exe_entry)
282 i : hook_ref = &hook_exe_entry_ref_;
283 :
284 E : EntryCallBasicBlockTransform entry_call_transform(*hook_ref, debug_friendly_);
285 : if (!ApplyBasicBlockSubGraphTransform(
286 E : &entry_call_transform, policy, block_graph, block, NULL)) {
287 i : return false;
288 : }
289 :
290 E : return true;
291 E : }
292 :
293 E : bool EntryCallTransform::GetEntryPoints(BlockGraph::Block* header_block) {
294 : // Get the TLS initializer entry-points. These have the same signature and
295 : // call patterns to DllMain.
296 E : if (!pe::GetTlsInitializers(header_block, &dllmain_entrypoints_)) {
297 i : LOG(ERROR) << "Failed to populate the TLS Initializer entry-points.";
298 i : return false;
299 : }
300 :
301 : // Get the DLL entry-point.
302 E : pe::EntryPoint dll_entry_point;
303 E : if (!pe::GetDllEntryPoint(header_block, &dll_entry_point)) {
304 i : LOG(ERROR) << "Failed to resolve the DLL entry-point.";
305 i : return false;
306 : }
307 :
308 : // If the image is an EXE or is a DLL that does not specify an entry-point
309 : // (the entry-point is optional for DLLs) then the dll_entry_point will have
310 : // a NULL block pointer. Otherwise, add it to the entry-point set.
311 E : if (dll_entry_point.first != NULL) {
312 E : dllmain_entrypoints_.insert(dll_entry_point);
313 E : } else {
314 : // Get the EXE entry point. We only need to bother looking if we didn't get
315 : // a DLL entry point, as we can't have both.
316 i : if (!pe::GetExeEntryPoint(header_block, &exe_entry_point_)) {
317 i : LOG(ERROR) << "Failed to resolve the EXE entry-point.";
318 i : return false;
319 : }
320 : }
321 :
322 E : return true;
323 E : }
324 :
325 : bool EntryCallTransform::PostBlockGraphIteration(
326 : const TransformPolicyInterface* policy,
327 : BlockGraph* block_graph,
328 E : BlockGraph::Block* header_block) {
329 : // Make sure the thunks section contains at least one block, as its existence
330 : // is what Chrome's glue code looks for to see whether it's instrumented.
331 : BlockGraph::Section* thunk_section =
332 E : block_graph->FindSection(common::kThunkSectionName);
333 E : if (thunk_section != NULL) {
334 : // It already exists - we're done!
335 i : return true;
336 : }
337 :
338 : // The section didn't already exist, create it.
339 : thunk_section = block_graph->FindOrAddSection(common::kThunkSectionName,
340 E : pe::kCodeCharacteristics);
341 E : DCHECK(thunk_section != NULL);
342 :
343 : // Create a one-byte marker block and assign it to the thunks segment.
344 : BlockGraph::Block* marker =
345 E : block_graph->AddBlock(BlockGraph::CODE_BLOCK, 1, "InstrumentationMarker");
346 E : DCHECK(marker != NULL);
347 :
348 E : marker->set_section(thunk_section->id());
349 :
350 : // Provide the marker function with valid code.
351 : static const uint8 kRet[] = { 0xC3 };
352 E : marker->SetData(kRet, sizeof(kRet));
353 :
354 E : return true;
355 E : }
356 :
357 : } // namespace transforms
358 : } // namespace instrument
|