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/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_NE(reinterpret_cast<TransformPolicyInterface*>(NULL), policy);
173 E : DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
174 E : DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), header_block);
175 E : DCHECK_EQ(BlockGraph::PE_IMAGE, block_graph->image_format());
176 :
177 : // Setup instrumentation functions hooks.
178 : if (!SetupEntryHooks(policy,
179 : block_graph,
180 : header_block,
181 : instrument_dll_name_,
182 : buffering_,
183 : fs_slot_,
184 : &function_enter_hook_ref_,
185 : &enter_hook_ref_,
186 E : &exit_hook_ref_)) {
187 i : return false;
188 : }
189 :
190 : // Add the static basic-block frequency data.
191 : if (!ApplyBlockGraphTransform(
192 E : &add_frequency_data_, policy, block_graph, header_block)) {
193 i : LOG(ERROR) << "Failed to insert basic-block frequency data.";
194 i : return false;
195 : }
196 :
197 E : return true;
198 E : }
199 :
200 : bool BranchHookTransform::OnBlock(const TransformPolicyInterface* policy,
201 : BlockGraph* block_graph,
202 E : BlockGraph::Block* block) {
203 E : DCHECK(block_graph != NULL);
204 E : DCHECK(block != NULL);
205 :
206 : // Ignore non-decomposable blocks.
207 E : if (!policy->BlockIsSafeToBasicBlockDecompose(block))
208 E : return true;
209 :
210 E : if (!ApplyBasicBlockSubGraphTransform(this, policy, block_graph, block, NULL))
211 i : return false;
212 :
213 E : return true;
214 E : }
215 :
216 : bool BranchHookTransform::TransformBasicBlockSubGraph(
217 : const TransformPolicyInterface* policy,
218 : BlockGraph* block_graph,
219 E : BasicBlockSubGraph* subgraph) {
220 E : DCHECK(policy != NULL);
221 E : DCHECK(block_graph != NULL);
222 E : DCHECK(subgraph != NULL);
223 E : DCHECK(enter_hook_ref_.IsValid());
224 E : DCHECK(exit_hook_ref_.IsValid());
225 E : DCHECK(add_frequency_data_.frequency_data_block() != NULL);
226 :
227 : // Determine whether we must pass the module_data pointer to the hooks.
228 E : bool need_module_data = true;
229 E : if (fs_slot_ != 0)
230 E : need_module_data = false;
231 :
232 E : BlockDescriptionList& descriptions = subgraph->block_descriptions();
233 E : BlockDescriptionList::iterator description = descriptions.begin();
234 E : for (; description != descriptions.end(); ++description) {
235 : BasicBlockSubGraph::BasicBlockOrdering& original_order =
236 E : (*description).basic_block_order;
237 :
238 : // Get the first basic block of this ordering.
239 E : DCHECK(!original_order.empty());
240 E : BasicCodeBlock* first_bb = BasicCodeBlock::Cast(*original_order.begin());
241 E : DCHECK(first_bb != NULL);
242 :
243 : // Insert a call to the basic-block entry hook at the beginning and the end
244 : // of each code basic-block.
245 : BasicBlockSubGraph::BasicBlockOrdering::const_iterator it =
246 E : original_order.begin();
247 E : for (; it != original_order.end(); ++it) {
248 E : BasicCodeBlock* bb = BasicCodeBlock::Cast(*it);
249 E : if (bb == NULL || bb->is_padding())
250 E : continue;
251 :
252 : // Find the source range associated with this basic-block.
253 E : BlockGraph::Block::SourceRange source_range;
254 E : if (!GetBasicBlockSourceRange(*bb, &source_range)) {
255 i : LOG(ERROR) << "Unable to get source range for basic block '"
256 : << bb->name() << "'";
257 i : return false;
258 : }
259 :
260 : // We use the index in the bb_ranges vector of the current basic-block
261 : // range as the basic_block_id, and we pass a pointer to the frequency
262 : // data block as the module_data parameter. We then make a memory indirect
263 : // call to the bb_entry_hook.
264 E : auto basic_block_id(Immediate(bb_ranges_.size(), assm::kSize32Bit));
265 : auto module_data(
266 E : Immediate(add_frequency_data_.frequency_data_block(), 0));
267 :
268 : // Assemble entry hook instrumentation into the instruction stream.
269 E : BlockGraph::Reference* enter_hook_ref = &enter_hook_ref_;
270 : auto enter_hook(Operand(Displacement(enter_hook_ref->referenced(),
271 E : enter_hook_ref->offset())));
272 : BasicBlockAssembler bb_asm_enter(bb->instructions().begin(),
273 E : &bb->instructions());
274 E : bb_asm_enter.push(basic_block_id);
275 E : if (need_module_data)
276 E : bb_asm_enter.push(module_data);
277 E : bb_asm_enter.call(enter_hook);
278 :
279 : // Find the last non jumping instruction in the basic block.
280 E : BasicBlock::Instructions::iterator last = bb->instructions().begin();
281 E : BasicBlock::Instructions::iterator last_instruction = last;
282 E : for (; last != bb->instructions().end(); ++last) {
283 E : if (!last->IsReturn() && !last->IsBranch()) {
284 E : last_instruction = last;
285 E : ++last_instruction;
286 : }
287 E : }
288 :
289 : if (last == bb->instructions().end() ||
290 E : !last->CallsNonReturningFunction()) {
291 : // Assemble exit hook instrumentation into the instruction stream.
292 : auto exit_hook(Operand(Displacement(exit_hook_ref_.referenced(),
293 E : exit_hook_ref_.offset())));
294 : BasicBlockAssembler bb_asm_exit(last_instruction,
295 E : &bb->instructions());
296 E : bb_asm_exit.push(basic_block_id);
297 E : if (need_module_data)
298 E : bb_asm_exit.push(module_data);
299 E : bb_asm_exit.call(exit_hook);
300 : }
301 :
302 : // Push the range for the current basic block.
303 E : bb_ranges_.push_back(source_range);
304 E : }
305 :
306 : // Insert a call to the function entry hook at the beginning of the
307 : // function.
308 E : if (function_enter_hook_ref_.IsValid()) {
309 : // Assemble function enter hook instrumentation into the instruction
310 : // stream.
311 : auto module_data(
312 E : Immediate(add_frequency_data_.frequency_data_block(), 0));
313 : auto func_hook(Operand(
314 : Displacement(function_enter_hook_ref_.referenced(),
315 E : function_enter_hook_ref_.offset())));
316 : BasicBlockAssembler func_asm_enter(first_bb->instructions().begin(),
317 E : &first_bb->instructions());
318 E : func_asm_enter.push(module_data);
319 E : func_asm_enter.call(func_hook);
320 : }
321 E : }
322 :
323 E : return true;
324 E : }
325 :
326 : bool BranchHookTransform::PostBlockGraphIteration(
327 : const TransformPolicyInterface* policy,
328 : BlockGraph* block_graph,
329 E : BlockGraph::Block* header_block) {
330 E : DCHECK(policy != NULL);
331 E : DCHECK(block_graph != NULL);
332 E : DCHECK(header_block != NULL);
333 :
334 E : size_t num_basic_blocks = bb_ranges_.size();
335 E : if (num_basic_blocks == 0) {
336 i : LOG(WARNING) << "Encountered no basic code blocks during instrumentation.";
337 i : return true;
338 : }
339 :
340 : if (!add_frequency_data_.ConfigureFrequencyDataBuffer(num_basic_blocks,
341 : 3,
342 E : sizeof(uint32))) {
343 i : LOG(ERROR) << "Failed to configure frequency data buffer.";
344 i : return false;
345 : }
346 :
347 : // Initialized BasicBlock agent specific fields
348 E : block_graph::TypedBlock<BasicBlockIndexedFrequencyData> frequency_data;
349 E : CHECK(frequency_data.Init(0, add_frequency_data_.frequency_data_block()));
350 E : frequency_data->fs_slot = fs_slot_;
351 E : frequency_data->tls_index = TLS_OUT_OF_INDEXES;
352 :
353 : // Add the module entry thunks.
354 E : EntryThunkTransform add_thunks;
355 E : add_thunks.set_only_instrument_module_entry(true);
356 E : add_thunks.set_instrument_dll_name(instrument_dll_name_);
357 E : add_thunks.set_src_ranges_for_thunks(true);
358 :
359 E : auto module_data(Immediate(add_frequency_data_.frequency_data_block(), 0));
360 E : if (!add_thunks.SetEntryThunkParameter(module_data)) {
361 i : LOG(ERROR) << "Failed to configure the entry thunks with the module_data "
362 : << "parameter.";
363 i : return false;
364 : }
365 :
366 : if (!ApplyBlockGraphTransform(
367 E : &add_thunks, policy, block_graph, header_block)) {
368 i : LOG(ERROR) << "Unable to thunk module entry points.";
369 i : return false;
370 : }
371 :
372 : // Find or create the section we put our thunks in.
373 : thunk_section_ = block_graph->FindOrAddSection(common::kThunkSectionName,
374 E : pe::kCodeCharacteristics);
375 E : DCHECK(thunk_section_ != NULL);
376 :
377 E : return true;
378 E : }
379 :
380 : } // namespace transforms
381 : } // namespace instrument
|