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 : // Implements the BasicBlockEntryHookTransform class.
16 :
17 : #include "syzygy/instrument/transforms/basic_block_entry_hook_transform.h"
18 :
19 : #include "base/logging.h"
20 : #include "base/string_util.h"
21 : #include "base/stringprintf.h"
22 : #include "syzygy/agent/basic_block_entry/basic_block_entry.h"
23 : #include "syzygy/block_graph/block_builder.h"
24 : #include "syzygy/block_graph/block_util.h"
25 : #include "syzygy/common/defs.h"
26 : #include "syzygy/common/indexed_frequency_data.h"
27 : #include "syzygy/instrument/transforms/entry_thunk_transform.h"
28 : #include "syzygy/pe/pe_utils.h"
29 : #include "syzygy/pe/transforms/pe_add_imports_transform.h"
30 :
31 : namespace instrument {
32 : namespace transforms {
33 :
34 : namespace {
35 :
36 : using agent::basic_block_entry::BasicBlockEntry;
37 : using block_graph::BasicBlock;
38 : using block_graph::BasicBlockAssembler;
39 : using block_graph::BasicBlockReference;
40 : using block_graph::BasicCodeBlock;
41 : using block_graph::BlockBuilder;
42 : using block_graph::BlockGraph;
43 : using block_graph::Displacement;
44 : using block_graph::Immediate;
45 : using block_graph::Operand;
46 : using block_graph::Successor;
47 : using block_graph::TransformPolicyInterface;
48 : using common::kBasicBlockEntryAgentId;
49 : using pe::transforms::PEAddImportsTransform;
50 :
51 : typedef BasicBlockEntryHookTransform::RelativeAddressRange RelativeAddressRange;
52 : typedef BasicBlockEntry::BasicBlockIndexedFrequencyData
53 : BasicBlockIndexedFrequencyData;
54 : typedef pe::transforms::ImportedModule ImportedModule;
55 :
56 : const char kDefaultModuleName[] = "basic_block_entry_client.dll";
57 : const char kBasicBlockEnter[] = "_increment_indexed_freq_data";
58 :
59 : // Compares two relative address ranges to see if they overlap. Assumes they
60 : // are already sorted. This is used to validate basic-block ranges.
61 : struct RelativeAddressRangesOverlapFunctor {
62 : bool operator()(const RelativeAddressRange& r1,
63 E : const RelativeAddressRange& r2) const {
64 E : DCHECK_LT(r1.start(), r2.start());
65 :
66 E : if (r1.end() > r2.start())
67 i : return true;
68 :
69 E : return false;
70 E : }
71 : };
72 :
73 : // Sets up the basic-block entry and the frequency data hooks import.
74 : bool SetupEntryHooks(const TransformPolicyInterface* policy,
75 : BlockGraph* block_graph,
76 : BlockGraph::Block* header_block,
77 : const std::string& module_name,
78 E : BlockGraph::Reference* basic_block_enter) {
79 E : DCHECK(block_graph != NULL);
80 E : DCHECK(header_block != NULL);
81 E : DCHECK(basic_block_enter != NULL);
82 :
83 : // Setup the import module.
84 E : ImportedModule module(module_name);
85 : size_t bb_index = module.AddSymbol(kBasicBlockEnter,
86 E : ImportedModule::kAlwaysImport);
87 :
88 : // Setup the add-imports transform.
89 E : PEAddImportsTransform add_imports;
90 E : add_imports.AddModule(&module);
91 :
92 : // Add the imports to the block-graph.
93 : if (!ApplyBlockGraphTransform(
94 E : &add_imports, policy, block_graph, header_block)) {
95 i : LOG(ERROR) << "Unable to add import entry hook functions.";
96 i : return false;
97 : }
98 :
99 : // Get a reference to the entry-hook function.
100 E : if (!module.GetSymbolReference(bb_index, basic_block_enter)) {
101 i : LOG(ERROR) << "Unable to get " << kBasicBlockEnter << ".";
102 i : return false;
103 : }
104 E : DCHECK(basic_block_enter->IsValid());
105 :
106 E : return true;
107 E : }
108 :
109 : void AddSuccessorBetween(Successor::Condition condition,
110 : BasicCodeBlock* from,
111 : BasicCodeBlock* to) {
112 : from->successors().push_back(
113 : Successor(condition,
114 : BasicBlockReference(BlockGraph::RELATIVE_REF,
115 : BlockGraph::Reference::kMaximumSize,
116 : to),
117 : 0));
118 : }
119 :
120 : } // namespace
121 :
122 : const char BasicBlockEntryHookTransform::kTransformName[] =
123 : "BasicBlockEntryHookTransform";
124 :
125 : BasicBlockEntryHookTransform::BasicBlockEntryHookTransform()
126 : : add_frequency_data_(kBasicBlockEntryAgentId,
127 : "Basic-Block Frequency Data",
128 : common::kBasicBlockFrequencyDataVersion,
129 : common::IndexedFrequencyData::BASIC_BLOCK_ENTRY,
130 : sizeof(BasicBlockIndexedFrequencyData)),
131 : thunk_section_(NULL),
132 : instrument_dll_name_(kDefaultModuleName),
133 : set_src_ranges_for_thunks_(false),
134 E : set_inline_fast_path_(false) {
135 E : }
136 :
137 : bool BasicBlockEntryHookTransform::PreBlockGraphIteration(
138 : const TransformPolicyInterface* policy,
139 : BlockGraph* block_graph,
140 E : BlockGraph::Block* header_block) {
141 E : DCHECK(policy != NULL);
142 E : DCHECK(block_graph != NULL);
143 E : DCHECK(header_block != NULL);
144 :
145 : // Setup basic block entry and the frequency data hooks.
146 : if (!SetupEntryHooks(policy,
147 : block_graph,
148 : header_block,
149 : instrument_dll_name_,
150 E : &bb_entry_hook_ref_)) {
151 i : return false;
152 : }
153 :
154 : // Add the static basic-block frequency data.
155 : if (!ApplyBlockGraphTransform(
156 E : &add_frequency_data_, policy, block_graph, header_block)) {
157 i : LOG(ERROR) << "Failed to insert basic-block frequency data.";
158 i : return false;
159 : }
160 :
161 : // Find or create the section we put our thunks in.
162 : thunk_section_ = block_graph->FindOrAddSection(common::kThunkSectionName,
163 E : pe::kCodeCharacteristics);
164 E : DCHECK(thunk_section_ != NULL);
165 :
166 E : return true;
167 E : }
168 :
169 : bool BasicBlockEntryHookTransform::OnBlock(
170 : const TransformPolicyInterface* policy,
171 : BlockGraph* block_graph,
172 E : BlockGraph::Block* block) {
173 E : DCHECK(policy != NULL);
174 E : DCHECK(block_graph != NULL);
175 E : DCHECK(block != NULL);
176 E : DCHECK(thunk_section_ != NULL);
177 :
178 : // Skip blocks created by this transform.
179 E : if (block->section() == thunk_section_->id())
180 i : return true;
181 :
182 E : if (block->type() != BlockGraph::CODE_BLOCK)
183 E : return true;
184 :
185 E : if (!policy->BlockIsSafeToBasicBlockDecompose(block)) {
186 E : if (!ThunkNonDecomposableCodeBlock(block_graph, block))
187 i : return false;
188 E : return true;
189 : }
190 :
191 E : if (!ApplyBasicBlockSubGraphTransform(this, policy, block_graph, block, NULL))
192 i : return false;
193 :
194 E : return true;
195 E : }
196 :
197 : bool BasicBlockEntryHookTransform::TransformBasicBlockSubGraph(
198 : const TransformPolicyInterface* policy,
199 : BlockGraph* block_graph ,
200 E : BasicBlockSubGraph* subgraph) {
201 : // TODO(rogerm): A lot of this is boilerplate that can be hoisted to an
202 : // IterativeBasicBlockSubgraphTransform (or some such). In particular,
203 : // iterating the subgraph, dispatch on code/data basic block, and the
204 : // bb_ranges_ are duplicated in the coverage transform.
205 E : DCHECK(policy != NULL);
206 E : DCHECK(block_graph != NULL);
207 E : DCHECK(subgraph != NULL);
208 E : DCHECK(bb_entry_hook_ref_.IsValid());
209 E : DCHECK(add_frequency_data_.frequency_data_block() != NULL);
210 :
211 : // Insert a call to the basic-block entry hook at the top of each code
212 : // basic-block.
213 : BasicBlockSubGraph::BBCollection::iterator it =
214 E : subgraph->basic_blocks().begin();
215 E : for (; it != subgraph->basic_blocks().end(); ++it) {
216 E : BasicCodeBlock* bb = BasicCodeBlock::Cast(*it);
217 E : if (bb == NULL || bb->is_padding())
218 E : continue;
219 :
220 : // Find the source range associated with this basic-block.
221 E : BlockGraph::Block::SourceRange source_range;
222 E : if (!GetBasicBlockSourceRange(*bb, &source_range)) {
223 i : LOG(ERROR) << "Unable to get source range for basic block '"
224 : << bb->name() << "'";
225 i : return false;
226 : }
227 :
228 : // We use the location/index in the bb_ranges vector of the current
229 : // basic-block range as the basic_block_id, and we pass a pointer to
230 : // the frequency data block as the module_data parameter. We then make
231 : // a memory indirect call to the bb_entry_hook.
232 E : Immediate basic_block_id(bb_ranges_.size(), core::kSize32Bit);
233 E : Immediate module_data(add_frequency_data_.frequency_data_block(), 0);
234 : Operand bb_entry_hook(Displacement(bb_entry_hook_ref_.referenced(),
235 E : bb_entry_hook_ref_.offset()));
236 :
237 : // Assemble entry hook instrumentation into the instruction stream.
238 E : BasicBlockAssembler bb_asm(bb->instructions().begin(), &bb->instructions());
239 :
240 E : bb_asm.push(basic_block_id);
241 E : bb_asm.push(module_data);
242 E : bb_asm.call(bb_entry_hook);
243 :
244 E : bb_ranges_.push_back(source_range);
245 E : }
246 :
247 E : return true;
248 E : }
249 :
250 : bool BasicBlockEntryHookTransform::PostBlockGraphIteration(
251 : const TransformPolicyInterface* policy,
252 : BlockGraph* block_graph,
253 E : BlockGraph::Block* header_block) {
254 E : DCHECK(policy != NULL);
255 E : DCHECK(block_graph != NULL);
256 E : DCHECK(header_block != NULL);
257 :
258 E : size_t num_basic_blocks = bb_ranges_.size();
259 E : if (num_basic_blocks == 0) {
260 i : LOG(WARNING) << "Encountered no basic code blocks during instrumentation.";
261 i : return true;
262 : }
263 :
264 : if (!add_frequency_data_.ConfigureFrequencyDataBuffer(num_basic_blocks,
265 : 1,
266 E : sizeof(uint32))) {
267 i : LOG(ERROR) << "Failed to configure frequency data buffer.";
268 i : return false;
269 : }
270 :
271 : // Initialized BasicBlock agent specific fields
272 E : block_graph::TypedBlock<BasicBlockIndexedFrequencyData> frequency_data;
273 E : CHECK(frequency_data.Init(0, add_frequency_data_.frequency_data_block()));
274 E : frequency_data->fs_slot = 0;
275 E : frequency_data->tls_index = TLS_OUT_OF_INDEXES;
276 :
277 : // Add the module entry thunks.
278 E : EntryThunkTransform add_thunks;
279 E : add_thunks.set_only_instrument_module_entry(true);
280 E : add_thunks.set_instrument_dll_name(instrument_dll_name_);
281 E : add_thunks.set_src_ranges_for_thunks(set_src_ranges_for_thunks_);
282 :
283 E : Immediate module_data(add_frequency_data_.frequency_data_block(), 0);
284 E : if (!add_thunks.SetEntryThunkParameter(module_data)) {
285 i : LOG(ERROR) << "Failed to configure the entry thunks with the module_data "
286 : << "parameter.";
287 i : return false;
288 : }
289 :
290 : if (!ApplyBlockGraphTransform(
291 E : &add_thunks, policy, block_graph, header_block)) {
292 i : LOG(ERROR) << "Unable to thunk module entry points.";
293 i : return false;
294 : }
295 E : thunk_section_ = add_thunks.thunk_section();
296 E : DCHECK(thunk_section_ != NULL);
297 :
298 : #ifndef NDEBUG
299 : // If we're in debug mode then sanity check the basic block ranges. When
300 : // sorted, they should not overlap.
301 E : RelativeAddressRangeVector bb_ranges_copy(bb_ranges_);
302 E : std::sort(bb_ranges_copy.begin(), bb_ranges_copy.end());
303 : DCHECK(std::adjacent_find(bb_ranges_copy.begin(),
304 : bb_ranges_copy.end(),
305 : RelativeAddressRangesOverlapFunctor()) ==
306 E : bb_ranges_copy.end());
307 : #endif
308 :
309 E : return true;
310 E : }
311 :
312 : bool BasicBlockEntryHookTransform::ThunkNonDecomposableCodeBlock(
313 E : BlockGraph* block_graph, BlockGraph::Block* code_block) {
314 E : DCHECK(block_graph != NULL);
315 E : DCHECK(code_block != NULL);
316 :
317 : // Typedef for the thunk block map. The key is the offset within the callee
318 : // block and the value is the thunk block that forwards to the callee at that
319 : // offset.
320 : typedef std::map<BlockGraph::Offset, BlockGraph::Block*> ThunkBlockMap;
321 :
322 : // We keep a cache of thunks we've already created (by target offset of the
323 : // entry-point into the block) so that we only create one thunk per entry
324 : // point.
325 E : ThunkBlockMap thunk_block_map;
326 :
327 : // Iterate through all the block's referrers, creating thunks as we go.
328 : // We copy the referrer set for simplicity, as it's potentially mutated
329 : // in the loop.
330 E : BlockGraph::Block::ReferrerSet referrers = code_block->referrers();
331 E : BlockGraph::Block::ReferrerSet::const_iterator referrer_it(referrers.begin());
332 E : for (; referrer_it != referrers.end(); ++referrer_it) {
333 E : const BlockGraph::Block::Referrer& referrer = *referrer_it;
334 : if (!EnsureReferrerIsThunked(
335 E : referrer, block_graph, code_block, &thunk_block_map)) {
336 i : return false;
337 : }
338 E : }
339 :
340 E : return true;
341 E : }
342 :
343 : bool BasicBlockEntryHookTransform::EnsureReferrerIsThunked(
344 : const BlockGraph::Block::Referrer& referrer,
345 : BlockGraph* block_graph,
346 : BlockGraph::Block* code_block,
347 E : ThunkBlockMap* thunk_block_map) {
348 E : DCHECK(block_graph != NULL);
349 E : DCHECK(code_block != NULL);
350 E : DCHECK(thunk_block_map != NULL);
351 :
352 : // Get the reference.
353 E : BlockGraph::Reference ref;
354 E : if (!referrer.first->GetReference(referrer.second, &ref)) {
355 i : LOG(ERROR) << "Unable to get reference from referrer.";
356 i : return false;
357 : }
358 E : DCHECK_EQ(code_block, ref.referenced());
359 :
360 : // Skip self-references, except long references to the start of the block.
361 : // Note: This may currently miss important cases. Notably if a block contains
362 : // more than one function, and the functions are mutually recursive, we'll
363 : // only record the original entry to the block, but will miss the internal
364 : // recursion. As-is, this does work for the common case where a block
365 : // contains one self-recursive function, however.
366 E : if (referrer.first == code_block) {
367 : // Skip short references.
368 E : if (ref.size() < sizeof(core::AbsoluteAddress))
369 E : return true;
370 :
371 : // Skip interior references. The block is not bb-decomposable so there is
372 : // nothing for us to do with them.
373 E : if (ref.offset() != 0)
374 E : return true;
375 : }
376 :
377 : // Get a thunk for the referenced offset from the thunk block map, creating
378 : // a new one if one does not already exist.
379 E : BlockGraph::Block* thunk_block = NULL;
380 : if (!FindOrCreateThunk(block_graph, thunk_block_map, code_block, ref.offset(),
381 E : &thunk_block)) {
382 i : LOG(ERROR) << "Unable to create thunk block.";
383 i : return false;
384 : }
385 E : DCHECK(thunk_block != NULL);
386 :
387 : // Update the referrer to point to the thunk.
388 : BlockGraph::Reference new_ref(ref.type(),
389 : ref.size(),
390 : thunk_block,
391 E : 0, 0);
392 E : referrer.first->SetReference(referrer.second, new_ref);
393 :
394 E : return true;
395 E : }
396 :
397 : bool BasicBlockEntryHookTransform::FindOrCreateThunk(
398 : BlockGraph* block_graph,
399 : ThunkBlockMap* thunk_block_map,
400 : BlockGraph::Block* code_block,
401 : BlockGraph::Offset offset,
402 E : BlockGraph::Block** thunk) {
403 E : DCHECK(block_graph != NULL);
404 E : DCHECK(thunk_block_map != NULL);
405 E : DCHECK(code_block != NULL);
406 E : DCHECK(thunk != NULL);
407 E : DCHECK_EQ(BlockGraph::CODE_BLOCK, code_block->type());
408 :
409 : // Do we already have a thunk defined for this offset? If so, return it.
410 E : ThunkBlockMap::const_iterator thunk_it = thunk_block_map->find(offset);
411 E : if (thunk_it != thunk_block_map->end()) {
412 E : *thunk = thunk_it->second;
413 E : return true;
414 : }
415 :
416 E : *thunk = NULL;
417 :
418 : // Determine the name for this thunk.
419 E : std::string name;
420 E : if (offset == 0) {
421 : name = base::StringPrintf("%s%s",
422 : code_block->name().c_str(),
423 E : common::kThunkSuffix);
424 E : } else {
425 : name = base::StringPrintf("%s%s+%d",
426 : code_block->name().c_str(),
427 : common::kThunkSuffix,
428 E : offset);
429 : }
430 :
431 : // Set up a basic block subgraph containing a single block description, with
432 : // that block description containing a single empty basic block, and get an
433 : // assembler writing into that basic block.
434 E : BasicBlockSubGraph subgraph;
435 E : BasicCodeBlock* bb = subgraph.AddBasicCodeBlock(name);
436 : BasicBlockSubGraph::BlockDescription* desc = subgraph.AddBlockDescription(
437 : name,
438 : NULL,
439 : BlockGraph::CODE_BLOCK,
440 : thunk_section_->id(),
441 : 1,
442 E : 0);
443 E : desc->basic_block_order.push_back(bb);
444 :
445 : // Find the source range associated with this block.
446 E : BlockGraph::Block::SourceRange source_range;
447 E : if (!code_block->source_ranges().empty())
448 E : source_range = code_block->source_ranges().range_pair(0).second;
449 :
450 : // Make sure we only push the source range if we have not already created
451 : // a source range mapping for this block (i.e., if the non-decomposable
452 : // block has multiple entry points, we want them to share an id). We can do
453 : // this because we handle one block at a time; so, all of a block's thunks
454 : // will be created as a group. This assertion is sanity checked in the
455 : // PostBlockGraphIteration function's check for overlapping source ranges.
456 E : if (bb_ranges_.empty() || source_range != bb_ranges_.back())
457 E : bb_ranges_.push_back(source_range);
458 :
459 : // We use the location/index in the bb_ranges vector of the current
460 : // basic-block range as the basic_block_id, and we pass a pointer to
461 : // the frequency data block as the module_data parameter. We then make
462 : // a memory indirect call to the bb_entry_hook.
463 E : Immediate basic_block_id(bb_ranges_.size()-1, core::kSize32Bit);
464 E : Immediate module_data(add_frequency_data_.frequency_data_block(), 0);
465 E : Immediate original_function(Displacement(code_block, offset));
466 : Operand bb_entry_hook(Displacement(bb_entry_hook_ref_.referenced(),
467 E : bb_entry_hook_ref_.offset()));
468 :
469 : // Assemble entry hook instrumentation into the thunk's instruction stream.
470 : // Note that we turn this into a simulated call, so that the return from
471 : // the bb entry hook continues from the thunked function.
472 E : BasicBlockAssembler bb_asm(bb->instructions().begin(), &bb->instructions());
473 E : bb_asm.push(basic_block_id);
474 E : bb_asm.push(module_data);
475 E : bb_asm.push(original_function);
476 E : bb_asm.jmp(bb_entry_hook);
477 :
478 : // Condense the whole mess into a block.
479 E : BlockBuilder block_builder(block_graph);
480 E : if (!block_builder.Merge(&subgraph)) {
481 i : LOG(ERROR) << "Failed to build thunk block.";
482 i : return false;
483 : }
484 :
485 : // Exactly one new block should have been created.
486 E : DCHECK_EQ(1u, block_builder.new_blocks().size());
487 E : *thunk = block_builder.new_blocks().front();
488 E : (*thunk_block_map)[offset] = *thunk;
489 :
490 E : return true;
491 E : }
492 :
493 : } // namespace transforms
494 : } // namespace instrument
|