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