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 : // Implements the BranchHookTransform class.
16 :
17 : #include "syzygy/instrument/transforms/branch_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::BasicBlockSubGraph;
41 : using block_graph::BasicCodeBlock;
42 : using block_graph::BlockBuilder;
43 : using block_graph::BlockGraph;
44 : using block_graph::Displacement;
45 : using block_graph::Immediate;
46 : using block_graph::Operand;
47 : using block_graph::Successor;
48 : using block_graph::TransformPolicyInterface;
49 : using common::kBasicBlockEntryAgentId;
50 : using pe::transforms::PEAddImportsTransform;
51 :
52 : typedef BasicBlockEntry::BasicBlockIndexedFrequencyData
53 : BasicBlockIndexedFrequencyData;
54 : typedef BasicBlockSubGraph::BlockDescriptionList
55 : BlockDescriptionList;
56 : typedef pe::transforms::ImportedModule ImportedModule;
57 :
58 : const char kDefaultModuleName[] = "basic_block_entry_client.dll";
59 : const char kBranchFunctionEnter[] = "_function_enter";
60 : const char kBranchEnter[] = "_branch_enter";
61 : const char kBranchEnterBuffered[] = "_branch_enter_buffered";
62 : const char kBranchExit[] = "_branch_exit";
63 : const size_t kNumBranchSlot = 4;
64 :
65 : // Sets up the entry and the exit hooks import.
66 : bool SetupEntryHooks(const TransformPolicyInterface* policy,
67 : BlockGraph* block_graph,
68 : BlockGraph::Block* header_block,
69 : const std::string& module_name,
70 : bool buffering,
71 : uint32 fs_slot,
72 : BlockGraph::Reference* function_enter,
73 : BlockGraph::Reference* branch_enter,
74 E : BlockGraph::Reference* branch_exit) {
75 E : DCHECK(policy != NULL);
76 E : DCHECK(block_graph != NULL);
77 E : DCHECK(header_block != NULL);
78 E : DCHECK(branch_enter != NULL);
79 E : DCHECK(branch_exit != NULL);
80 :
81 : // Determine which hooks to use.
82 E : std::string function_enter_name;
83 E : std::string branch_enter_name;
84 E : std::string branch_exit_name;
85 :
86 E : if (buffering) {
87 E : branch_enter_name = kBranchEnterBuffered;
88 E : branch_exit_name = kBranchExit;
89 E : } else {
90 E : branch_enter_name = kBranchEnter;
91 E : branch_exit_name = kBranchExit;
92 : }
93 :
94 E : if (fs_slot != 0) {
95 : function_enter_name =
96 E : base::StringPrintf("%s_s%d", kBranchFunctionEnter, fs_slot);
97 : branch_enter_name =
98 E : base::StringPrintf("%s_s%d", branch_enter_name.c_str(), fs_slot);
99 : branch_exit_name =
100 E : base::StringPrintf("%s_s%d", branch_exit_name.c_str(), fs_slot);
101 : }
102 :
103 : // Setup the import module.
104 E : ImportedModule module(module_name);
105 : size_t enter_index = module.AddSymbol(branch_enter_name,
106 E : ImportedModule::kAlwaysImport);
107 :
108 : size_t exit_index = module.AddSymbol(branch_exit_name,
109 E : ImportedModule::kAlwaysImport);
110 E : size_t function_enter_index = 0;
111 E : if (!function_enter_name.empty()) {
112 : function_enter_index = module.AddSymbol(function_enter_name,
113 E : ImportedModule::kAlwaysImport);
114 : }
115 :
116 : // Setup the add-imports transform.
117 E : PEAddImportsTransform add_imports;
118 E : add_imports.AddModule(&module);
119 :
120 : // Add the imports to the block-graph.
121 : if (!ApplyBlockGraphTransform(
122 E : &add_imports, policy, block_graph, header_block)) {
123 i : LOG(ERROR) << "Unable to add import entry hook functions.";
124 i : return false;
125 : }
126 :
127 : // Get a reference to the entry-hook function.
128 E : if (!module.GetSymbolReference(enter_index, branch_enter)) {
129 i : LOG(ERROR) << "Unable to get " << kBranchEnter << ".";
130 i : return false;
131 : }
132 E : DCHECK(branch_enter->IsValid());
133 :
134 : // Get a reference to the exit-hook function.
135 E : if (!module.GetSymbolReference(exit_index, branch_exit)) {
136 i : LOG(ERROR) << "Unable to get " << kBranchExit << ".";
137 i : return false;
138 : }
139 E : DCHECK(branch_exit->IsValid());
140 :
141 E : if (!function_enter_name.empty()) {
142 E : if (!module.GetSymbolReference(function_enter_index, function_enter)) {
143 i : LOG(ERROR) << "Unable to get " << function_enter_name << ".";
144 i : return false;
145 : }
146 E : DCHECK(function_enter->IsValid());
147 : }
148 :
149 E : return true;
150 E : }
151 :
152 : } // namespace
153 :
154 : const char BranchHookTransform::kTransformName[] = "BranchTransform";
155 :
156 : BranchHookTransform::BranchHookTransform()
157 : : add_frequency_data_(kBasicBlockEntryAgentId,
158 : "Basic-Block Branch Information Data",
159 : common::kBranchFrequencyDataVersion,
160 : common::IndexedFrequencyData::BRANCH,
161 : sizeof(BasicBlockIndexedFrequencyData)),
162 : thunk_section_(NULL),
163 : instrument_dll_name_(kDefaultModuleName),
164 : buffering_(false),
165 E : fs_slot_(0U) {
166 E : }
167 :
168 : bool BranchHookTransform::PreBlockGraphIteration(
169 : const TransformPolicyInterface* policy,
170 : BlockGraph* block_graph,
171 E : BlockGraph::Block* header_block) {
172 E : DCHECK(policy != NULL);
173 E : DCHECK(block_graph != NULL);
174 E : DCHECK(header_block != NULL);
175 :
176 : // Setup instrumentation functions hooks.
177 : if (!SetupEntryHooks(policy,
178 : block_graph,
179 : header_block,
180 : instrument_dll_name_,
181 : buffering_,
182 : fs_slot_,
183 : &function_enter_hook_ref_,
184 : &enter_hook_ref_,
185 E : &exit_hook_ref_)) {
186 i : return false;
187 : }
188 :
189 : // Add the static basic-block frequency data.
190 : if (!ApplyBlockGraphTransform(
191 E : &add_frequency_data_, policy, block_graph, header_block)) {
192 i : LOG(ERROR) << "Failed to insert basic-block frequency data.";
193 i : return false;
194 : }
195 :
196 E : return true;
197 E : }
198 :
199 : bool BranchHookTransform::OnBlock(const TransformPolicyInterface* policy,
200 : BlockGraph* block_graph,
201 E : BlockGraph::Block* block) {
202 E : DCHECK(block_graph != NULL);
203 E : DCHECK(block != NULL);
204 :
205 : // Ignore non-decomposable blocks.
206 E : if (!policy->BlockIsSafeToBasicBlockDecompose(block))
207 E : return true;
208 :
209 E : if (!ApplyBasicBlockSubGraphTransform(this, policy, block_graph, block, NULL))
210 i : return false;
211 :
212 E : return true;
213 E : }
214 :
215 : bool BranchHookTransform::TransformBasicBlockSubGraph(
216 : const TransformPolicyInterface* policy,
217 : BlockGraph* block_graph,
218 E : BasicBlockSubGraph* subgraph) {
219 E : DCHECK(policy != NULL);
220 E : DCHECK(block_graph != NULL);
221 E : DCHECK(subgraph != NULL);
222 E : DCHECK(enter_hook_ref_.IsValid());
223 E : DCHECK(exit_hook_ref_.IsValid());
224 E : DCHECK(add_frequency_data_.frequency_data_block() != NULL);
225 :
226 : // Determine whether we must pass the module_data pointer to the hooks.
227 E : bool need_module_data = true;
228 E : if (fs_slot_ != 0)
229 E : need_module_data = false;
230 :
231 E : BlockDescriptionList& descriptions = subgraph->block_descriptions();
232 E : BlockDescriptionList::iterator description = descriptions.begin();
233 E : for (; description != descriptions.end(); ++description) {
234 : BasicBlockSubGraph::BasicBlockOrdering& original_order =
235 E : (*description).basic_block_order;
236 :
237 : // Get the first basic block of this ordering.
238 E : DCHECK(!original_order.empty());
239 E : BasicCodeBlock* first_bb = BasicCodeBlock::Cast(*original_order.begin());
240 E : DCHECK(first_bb != NULL);
241 :
242 : // Insert a call to the basic-block entry hook at the beginning and the end
243 : // of each code basic-block.
244 : BasicBlockSubGraph::BasicBlockOrdering::const_iterator it =
245 E : original_order.begin();
246 E : for (; it != original_order.end(); ++it) {
247 E : BasicCodeBlock* bb = BasicCodeBlock::Cast(*it);
248 E : if (bb == NULL || bb->is_padding())
249 E : continue;
250 :
251 : // Find the source range associated with this basic-block.
252 E : BlockGraph::Block::SourceRange source_range;
253 E : if (!GetBasicBlockSourceRange(*bb, &source_range)) {
254 i : LOG(ERROR) << "Unable to get source range for basic block '"
255 : << bb->name() << "'";
256 i : return false;
257 : }
258 :
259 : // We use the index in the bb_ranges vector of the current basic-block
260 : // range as the basic_block_id, and we pass a pointer to the frequency
261 : // data block as the module_data parameter. We then make a memory indirect
262 : // call to the bb_entry_hook.
263 E : Immediate basic_block_id(bb_ranges_.size(), core::kSize32Bit);
264 E : Immediate module_data(add_frequency_data_.frequency_data_block(), 0);
265 :
266 : // Assemble entry hook instrumentation into the instruction stream.
267 E : BlockGraph::Reference* enter_hook_ref = &enter_hook_ref_;
268 : Operand enter_hook(Displacement(enter_hook_ref->referenced(),
269 E : enter_hook_ref->offset()));
270 : BasicBlockAssembler bb_asm_enter(bb->instructions().begin(),
271 E : &bb->instructions());
272 E : bb_asm_enter.push(basic_block_id);
273 E : if (need_module_data)
274 E : bb_asm_enter.push(module_data);
275 E : bb_asm_enter.call(enter_hook);
276 :
277 : // Find the last non jumping instruction in the basic block.
278 E : BasicBlock::Instructions::iterator last = bb->instructions().begin();
279 E : BasicBlock::Instructions::iterator last_instruction = last;
280 E : for (; last != bb->instructions().end(); ++last) {
281 E : if (!last->IsReturn() && !last->IsBranch()) {
282 E : last_instruction = last;
283 E : ++last_instruction;
284 : }
285 E : }
286 :
287 : if (last == bb->instructions().end() ||
288 E : !last->CallsNonReturningFunction()) {
289 : // Assemble exit hook instrumentation into the instruction stream.
290 : Operand exit_hook(Displacement(exit_hook_ref_.referenced(),
291 E : exit_hook_ref_.offset()));
292 : BasicBlockAssembler bb_asm_exit(last_instruction,
293 E : &bb->instructions());
294 E : bb_asm_exit.push(basic_block_id);
295 E : if (need_module_data)
296 E : bb_asm_exit.push(module_data);
297 E : bb_asm_exit.call(exit_hook);
298 E : }
299 :
300 : // Push the range for the current basic block.
301 E : bb_ranges_.push_back(source_range);
302 E : }
303 :
304 : // Insert a call to the function entry hook at the beginning of the
305 : // function.
306 E : if (function_enter_hook_ref_.IsValid()) {
307 : // Assemble function enter hook instrumentation into the instruction
308 : // stream.
309 E : Immediate module_data(add_frequency_data_.frequency_data_block(), 0);
310 : Operand func_hook(Displacement(function_enter_hook_ref_.referenced(),
311 E : function_enter_hook_ref_.offset()));
312 : BasicBlockAssembler func_asm_enter(first_bb->instructions().begin(),
313 E : &first_bb->instructions());
314 E : func_asm_enter.push(module_data);
315 E : func_asm_enter.call(func_hook);
316 E : }
317 E : }
318 :
319 E : return true;
320 E : }
321 :
322 : bool BranchHookTransform::PostBlockGraphIteration(
323 : const TransformPolicyInterface* policy,
324 : BlockGraph* block_graph,
325 E : BlockGraph::Block* header_block) {
326 E : DCHECK(policy != NULL);
327 E : DCHECK(block_graph != NULL);
328 E : DCHECK(header_block != NULL);
329 :
330 E : size_t num_basic_blocks = bb_ranges_.size();
331 E : if (num_basic_blocks == 0) {
332 i : LOG(WARNING) << "Encountered no basic code blocks during instrumentation.";
333 i : return true;
334 : }
335 :
336 : if (!add_frequency_data_.ConfigureFrequencyDataBuffer(num_basic_blocks,
337 : 3,
338 E : sizeof(uint32))) {
339 i : LOG(ERROR) << "Failed to configure frequency data buffer.";
340 i : return false;
341 : }
342 :
343 : // Initialized BasicBlock agent specific fields
344 E : block_graph::TypedBlock<BasicBlockIndexedFrequencyData> frequency_data;
345 E : CHECK(frequency_data.Init(0, add_frequency_data_.frequency_data_block()));
346 E : frequency_data->fs_slot = fs_slot_;
347 E : frequency_data->tls_index = TLS_OUT_OF_INDEXES;
348 :
349 : // Add the module entry thunks.
350 E : EntryThunkTransform add_thunks;
351 E : add_thunks.set_only_instrument_module_entry(true);
352 E : add_thunks.set_instrument_dll_name(instrument_dll_name_);
353 E : add_thunks.set_src_ranges_for_thunks(true);
354 :
355 E : Immediate module_data(add_frequency_data_.frequency_data_block(), 0);
356 E : if (!add_thunks.SetEntryThunkParameter(module_data)) {
357 i : LOG(ERROR) << "Failed to configure the entry thunks with the module_data "
358 : << "parameter.";
359 i : return false;
360 : }
361 :
362 : if (!ApplyBlockGraphTransform(
363 E : &add_thunks, policy, block_graph, header_block)) {
364 i : LOG(ERROR) << "Unable to thunk module entry points.";
365 i : return false;
366 : }
367 :
368 : // Find or create the section we put our thunks in.
369 : thunk_section_ = block_graph->FindOrAddSection(common::kThunkSectionName,
370 E : pe::kCodeCharacteristics);
371 E : DCHECK(thunk_section_ != NULL);
372 :
373 E : return true;
374 E : }
375 :
376 : } // namespace transforms
377 : } // namespace instrument
|