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/strings/string_util.h"
21 : #include "base/strings/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_NE(reinterpret_cast<TransformPolicyInterface*>(NULL), policy);
142 E : DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
143 E : DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), header_block);
144 E : DCHECK_EQ(BlockGraph::PE_IMAGE, block_graph->image_format());
145 :
146 : // Setup basic block entry and the frequency data hooks.
147 : if (!SetupEntryHooks(policy,
148 : block_graph,
149 : header_block,
150 : instrument_dll_name_,
151 E : &bb_entry_hook_ref_)) {
152 i : return false;
153 : }
154 :
155 : // Add the static basic-block frequency data.
156 : if (!ApplyBlockGraphTransform(
157 E : &add_frequency_data_, policy, block_graph, header_block)) {
158 i : LOG(ERROR) << "Failed to insert basic-block frequency data.";
159 i : return false;
160 : }
161 :
162 : // Find or create the section we put our thunks in.
163 : thunk_section_ = block_graph->FindOrAddSection(common::kThunkSectionName,
164 E : pe::kCodeCharacteristics);
165 E : DCHECK(thunk_section_ != NULL);
166 :
167 E : return true;
168 E : }
169 :
170 : bool BasicBlockEntryHookTransform::OnBlock(
171 : const TransformPolicyInterface* policy,
172 : BlockGraph* block_graph,
173 E : BlockGraph::Block* block) {
174 E : DCHECK(policy != NULL);
175 E : DCHECK(block_graph != NULL);
176 E : DCHECK(block != NULL);
177 E : DCHECK(thunk_section_ != NULL);
178 :
179 : // Skip blocks created by this transform.
180 E : if (block->section() == thunk_section_->id())
181 i : return true;
182 :
183 E : if (block->type() != BlockGraph::CODE_BLOCK)
184 E : return true;
185 :
186 E : if (!policy->BlockIsSafeToBasicBlockDecompose(block)) {
187 E : if (!ThunkNonDecomposableCodeBlock(block_graph, block))
188 i : return false;
189 E : return true;
190 : }
191 :
192 E : if (!ApplyBasicBlockSubGraphTransform(this, policy, block_graph, block, NULL))
193 i : return false;
194 :
195 E : return true;
196 E : }
197 :
198 : bool BasicBlockEntryHookTransform::TransformBasicBlockSubGraph(
199 : const TransformPolicyInterface* policy,
200 : BlockGraph* block_graph ,
201 E : BasicBlockSubGraph* subgraph) {
202 : // TODO(rogerm): A lot of this is boilerplate that can be hoisted to an
203 : // IterativeBasicBlockSubgraphTransform (or some such). In particular,
204 : // iterating the subgraph, dispatch on code/data basic block, and the
205 : // bb_ranges_ are duplicated in the coverage transform.
206 E : DCHECK(policy != NULL);
207 E : DCHECK(block_graph != NULL);
208 E : DCHECK(subgraph != NULL);
209 E : DCHECK(bb_entry_hook_ref_.IsValid());
210 E : DCHECK(add_frequency_data_.frequency_data_block() != NULL);
211 :
212 : // Insert a call to the basic-block entry hook at the top of each code
213 : // basic-block.
214 : BasicBlockSubGraph::BBCollection::iterator it =
215 E : subgraph->basic_blocks().begin();
216 E : for (; it != subgraph->basic_blocks().end(); ++it) {
217 E : BasicCodeBlock* bb = BasicCodeBlock::Cast(*it);
218 E : if (bb == NULL || bb->is_padding())
219 E : continue;
220 :
221 : // Find the source range associated with this basic-block.
222 E : BlockGraph::Block::SourceRange source_range;
223 E : if (!GetBasicBlockSourceRange(*bb, &source_range)) {
224 i : LOG(ERROR) << "Unable to get source range for basic block '"
225 : << bb->name() << "'";
226 i : return false;
227 : }
228 :
229 : // We use the location/index in the bb_ranges vector of the current
230 : // basic-block range as the basic_block_id, and we pass a pointer to
231 : // the frequency data block as the module_data parameter. We then make
232 : // a memory indirect call to the bb_entry_hook.
233 E : auto basic_block_id(Immediate(bb_ranges_.size(), assm::kSize32Bit));
234 E : auto module_data(Immediate(add_frequency_data_.frequency_data_block(), 0));
235 : auto bb_entry_hook(Operand(Displacement(bb_entry_hook_ref_.referenced(),
236 E : bb_entry_hook_ref_.offset())));
237 :
238 : // Assemble entry hook instrumentation into the instruction stream.
239 E : BasicBlockAssembler bb_asm(bb->instructions().begin(), &bb->instructions());
240 :
241 E : bb_asm.push(basic_block_id);
242 E : bb_asm.push(module_data);
243 E : bb_asm.call(bb_entry_hook);
244 :
245 E : bb_ranges_.push_back(source_range);
246 E : }
247 :
248 E : return true;
249 E : }
250 :
251 : bool BasicBlockEntryHookTransform::PostBlockGraphIteration(
252 : const TransformPolicyInterface* policy,
253 : BlockGraph* block_graph,
254 E : BlockGraph::Block* header_block) {
255 E : DCHECK(policy != NULL);
256 E : DCHECK(block_graph != NULL);
257 E : DCHECK(header_block != NULL);
258 :
259 E : size_t num_basic_blocks = bb_ranges_.size();
260 E : if (num_basic_blocks == 0) {
261 i : LOG(WARNING) << "Encountered no basic code blocks during instrumentation.";
262 i : return true;
263 : }
264 :
265 : if (!add_frequency_data_.ConfigureFrequencyDataBuffer(num_basic_blocks,
266 : 1,
267 E : sizeof(uint32))) {
268 i : LOG(ERROR) << "Failed to configure frequency data buffer.";
269 i : return false;
270 : }
271 :
272 : // Initialized BasicBlock agent specific fields
273 E : block_graph::TypedBlock<BasicBlockIndexedFrequencyData> frequency_data;
274 E : CHECK(frequency_data.Init(0, add_frequency_data_.frequency_data_block()));
275 E : frequency_data->fs_slot = 0;
276 E : frequency_data->tls_index = TLS_OUT_OF_INDEXES;
277 :
278 : // Add the module entry thunks.
279 E : EntryThunkTransform add_thunks;
280 E : add_thunks.set_only_instrument_module_entry(true);
281 E : add_thunks.set_instrument_dll_name(instrument_dll_name_);
282 E : add_thunks.set_src_ranges_for_thunks(set_src_ranges_for_thunks_);
283 :
284 E : auto module_data(Immediate(add_frequency_data_.frequency_data_block(), 0));
285 E : if (!add_thunks.SetEntryThunkParameter(module_data)) {
286 i : LOG(ERROR) << "Failed to configure the entry thunks with the module_data "
287 : << "parameter.";
288 i : return false;
289 : }
290 :
291 : if (!ApplyBlockGraphTransform(
292 E : &add_thunks, policy, block_graph, header_block)) {
293 i : LOG(ERROR) << "Unable to thunk module entry points.";
294 i : return false;
295 : }
296 E : thunk_section_ = add_thunks.thunk_section();
297 E : DCHECK(thunk_section_ != NULL);
298 :
299 : #ifndef NDEBUG
300 : // If we're in debug mode then sanity check the basic block ranges. When
301 : // sorted, they should not overlap.
302 E : RelativeAddressRangeVector bb_ranges_copy(bb_ranges_);
303 E : std::sort(bb_ranges_copy.begin(), bb_ranges_copy.end());
304 : DCHECK(std::adjacent_find(bb_ranges_copy.begin(),
305 : bb_ranges_copy.end(),
306 : RelativeAddressRangesOverlapFunctor()) ==
307 E : bb_ranges_copy.end());
308 : #endif
309 :
310 E : return true;
311 E : }
312 :
313 : bool BasicBlockEntryHookTransform::ThunkNonDecomposableCodeBlock(
314 E : BlockGraph* block_graph, BlockGraph::Block* code_block) {
315 E : DCHECK(block_graph != NULL);
316 E : DCHECK(code_block != NULL);
317 :
318 : // Typedef for the thunk block map. The key is the offset within the callee
319 : // block and the value is the thunk block that forwards to the callee at that
320 : // offset.
321 : typedef std::map<BlockGraph::Offset, BlockGraph::Block*> ThunkBlockMap;
322 :
323 : // We keep a cache of thunks we've already created (by target offset of the
324 : // entry-point into the block) so that we only create one thunk per entry
325 : // point.
326 E : ThunkBlockMap thunk_block_map;
327 :
328 : // Iterate through all the block's referrers, creating thunks as we go.
329 : // We copy the referrer set for simplicity, as it's potentially mutated
330 : // in the loop.
331 E : BlockGraph::Block::ReferrerSet referrers = code_block->referrers();
332 E : BlockGraph::Block::ReferrerSet::const_iterator referrer_it(referrers.begin());
333 E : for (; referrer_it != referrers.end(); ++referrer_it) {
334 E : const BlockGraph::Block::Referrer& referrer = *referrer_it;
335 : if (!EnsureReferrerIsThunked(
336 E : referrer, block_graph, code_block, &thunk_block_map)) {
337 i : return false;
338 : }
339 E : }
340 :
341 E : return true;
342 E : }
343 :
344 : bool BasicBlockEntryHookTransform::EnsureReferrerIsThunked(
345 : const BlockGraph::Block::Referrer& referrer,
346 : BlockGraph* block_graph,
347 : BlockGraph::Block* code_block,
348 E : ThunkBlockMap* thunk_block_map) {
349 E : DCHECK(block_graph != NULL);
350 E : DCHECK(code_block != NULL);
351 E : DCHECK(thunk_block_map != NULL);
352 :
353 : // Get the reference.
354 E : BlockGraph::Reference ref;
355 E : if (!referrer.first->GetReference(referrer.second, &ref)) {
356 i : LOG(ERROR) << "Unable to get reference from referrer.";
357 i : return false;
358 : }
359 E : DCHECK_EQ(code_block, ref.referenced());
360 :
361 : // Skip self-references, except long references to the start of the block.
362 : // Note: This may currently miss important cases. Notably if a block contains
363 : // more than one function, and the functions are mutually recursive, we'll
364 : // only record the original entry to the block, but will miss the internal
365 : // recursion. As-is, this does work for the common case where a block
366 : // contains one self-recursive function, however.
367 E : if (referrer.first == code_block) {
368 : // Skip short references.
369 E : if (ref.size() < sizeof(core::AbsoluteAddress))
370 i : return true;
371 :
372 : // Skip interior references. The block is not bb-decomposable so there is
373 : // nothing for us to do with them.
374 E : if (ref.offset() != 0)
375 E : return true;
376 : }
377 :
378 : // Get a thunk for the referenced offset from the thunk block map, creating
379 : // a new one if one does not already exist.
380 E : BlockGraph::Block* thunk_block = NULL;
381 : if (!FindOrCreateThunk(block_graph, thunk_block_map, code_block, ref.offset(),
382 E : &thunk_block)) {
383 i : LOG(ERROR) << "Unable to create thunk block.";
384 i : return false;
385 : }
386 E : DCHECK(thunk_block != NULL);
387 :
388 : // Update the referrer to point to the thunk.
389 : BlockGraph::Reference new_ref(ref.type(),
390 : ref.size(),
391 : thunk_block,
392 E : 0, 0);
393 E : referrer.first->SetReference(referrer.second, new_ref);
394 :
395 E : return true;
396 E : }
397 :
398 : bool BasicBlockEntryHookTransform::FindOrCreateThunk(
399 : BlockGraph* block_graph,
400 : ThunkBlockMap* thunk_block_map,
401 : BlockGraph::Block* code_block,
402 : BlockGraph::Offset offset,
403 E : BlockGraph::Block** thunk) {
404 E : DCHECK(block_graph != NULL);
405 E : DCHECK(thunk_block_map != NULL);
406 E : DCHECK(code_block != NULL);
407 E : DCHECK(thunk != NULL);
408 E : DCHECK_EQ(BlockGraph::CODE_BLOCK, code_block->type());
409 :
410 : // Do we already have a thunk defined for this offset? If so, return it.
411 E : ThunkBlockMap::const_iterator thunk_it = thunk_block_map->find(offset);
412 E : if (thunk_it != thunk_block_map->end()) {
413 E : *thunk = thunk_it->second;
414 E : return true;
415 : }
416 :
417 E : *thunk = NULL;
418 :
419 : // Determine the name for this thunk.
420 E : std::string name;
421 E : if (offset == 0) {
422 : name = base::StringPrintf("%s%s",
423 : code_block->name().c_str(),
424 E : common::kThunkSuffix);
425 E : } else {
426 : name = base::StringPrintf("%s%s+%d",
427 : code_block->name().c_str(),
428 : common::kThunkSuffix,
429 E : offset);
430 : }
431 :
432 : // Set up a basic block subgraph containing a single block description, with
433 : // that block description containing a single empty basic block, and get an
434 : // assembler writing into that basic block.
435 E : BasicBlockSubGraph subgraph;
436 E : BasicCodeBlock* bb = subgraph.AddBasicCodeBlock(name);
437 : BasicBlockSubGraph::BlockDescription* desc = subgraph.AddBlockDescription(
438 : name,
439 : NULL,
440 : BlockGraph::CODE_BLOCK,
441 : thunk_section_->id(),
442 : 1,
443 E : 0);
444 E : desc->basic_block_order.push_back(bb);
445 :
446 : // Find the source range associated with this block.
447 E : BlockGraph::Block::SourceRange source_range;
448 E : if (!code_block->source_ranges().empty())
449 E : source_range = code_block->source_ranges().range_pair(0).second;
450 :
451 : // Make sure we only push the source range if we have not already created
452 : // a source range mapping for this block (i.e., if the non-decomposable
453 : // block has multiple entry points, we want them to share an id). We can do
454 : // this because we handle one block at a time; so, all of a block's thunks
455 : // will be created as a group. This assertion is sanity checked in the
456 : // PostBlockGraphIteration function's check for overlapping source ranges.
457 E : if (bb_ranges_.empty() || source_range != bb_ranges_.back())
458 E : bb_ranges_.push_back(source_range);
459 :
460 : // We use the location/index in the bb_ranges vector of the current
461 : // basic-block range as the basic_block_id, and we pass a pointer to
462 : // the frequency data block as the module_data parameter. We then make
463 : // a memory indirect call to the bb_entry_hook.
464 E : auto basic_block_id(Immediate(bb_ranges_.size()-1, assm::kSize32Bit));
465 E : auto module_data(Immediate(add_frequency_data_.frequency_data_block(), 0));
466 E : auto original_function(Immediate(code_block, offset));
467 : auto bb_entry_hook(Operand(Displacement(bb_entry_hook_ref_.referenced(),
468 E : bb_entry_hook_ref_.offset())));
469 :
470 : // Assemble entry hook instrumentation into the thunk's instruction stream.
471 : // Note that we turn this into a simulated call, so that the return from
472 : // the bb entry hook continues from the thunked function.
473 E : BasicBlockAssembler bb_asm(bb->instructions().begin(), &bb->instructions());
474 E : bb_asm.push(basic_block_id);
475 E : bb_asm.push(module_data);
476 E : bb_asm.push(original_function);
477 E : bb_asm.jmp(bb_entry_hook);
478 :
479 : // Condense the whole mess into a block.
480 E : BlockBuilder block_builder(block_graph);
481 E : if (!block_builder.Merge(&subgraph)) {
482 i : LOG(ERROR) << "Failed to build thunk block.";
483 i : return false;
484 : }
485 :
486 : // Exactly one new block should have been created.
487 E : DCHECK_EQ(1u, block_builder.new_blocks().size());
488 E : *thunk = block_builder.new_blocks().front();
489 E : (*thunk_block_map)[offset] = *thunk;
490 :
491 E : return true;
492 E : }
493 :
494 : } // namespace transforms
495 : } // namespace instrument
|