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 JumpTableCaseCountTransform class.
16 :
17 : #include "syzygy/instrument/transforms/jump_table_count_transform.h"
18 :
19 : #include <limits>
20 :
21 : #include "base/logging.h"
22 : #include "base/string_util.h"
23 : #include "base/stringprintf.h"
24 : #include "syzygy/block_graph/block_builder.h"
25 : #include "syzygy/block_graph/block_util.h"
26 : #include "syzygy/common/defs.h"
27 : #include "syzygy/common/indexed_frequency_data.h"
28 : #include "syzygy/instrument/transforms/entry_thunk_transform.h"
29 : #include "syzygy/pe/pe_utils.h"
30 : #include "syzygy/pe/transforms/pe_add_imports_transform.h"
31 :
32 : namespace instrument {
33 : namespace transforms {
34 :
35 : namespace {
36 :
37 : using block_graph::BasicBlock;
38 : using block_graph::BasicBlockSubGraph;
39 : using block_graph::BasicCodeBlock;
40 : using block_graph::BasicDataBlock;
41 : using block_graph::BasicBlockAssembler;
42 : using block_graph::BasicBlockReference;
43 : using block_graph::BlockBuilder;
44 : using block_graph::BlockGraph;
45 : using block_graph::Displacement;
46 : using block_graph::Immediate;
47 : using block_graph::Instruction;
48 : using block_graph::Operand;
49 : using block_graph::TransformPolicyInterface;
50 : using pe::transforms::PEAddImportsTransform;
51 : using pe::transforms::ImportedModule;
52 :
53 : const char kDefaultModuleName[] = "basic_block_entry_client.dll";
54 : const char kJumpTableCaseCounter[] = "_increment_indexed_freq_data";
55 : const char kThunkSuffix[] = "_jump_table_thunk";
56 :
57 : // Sets up the jump table counter hook import.
58 : // @param policy The policy object restricting how the transform is applied.
59 : // @param block_graph The block-graph to populate.
60 : // @param header_block The header block from block_graph.
61 : // @param module_name The name of the module implementing the hooks.
62 : // @param jump_table_case_counter will refer to the imported hook function.
63 : // @returns true on success, false otherwise.
64 : bool SetupCounterHook(const TransformPolicyInterface* policy,
65 : BlockGraph* block_graph,
66 : BlockGraph::Block* header_block,
67 : const std::string& module_name,
68 E : BlockGraph::Reference* jump_table_case_counter) {
69 E : DCHECK(block_graph != NULL);
70 E : DCHECK(header_block != NULL);
71 E : DCHECK(jump_table_case_counter != NULL);
72 :
73 : // Setup the import module.
74 E : ImportedModule module(module_name);
75 : size_t index_case_counter = module.AddSymbol(
76 : kJumpTableCaseCounter,
77 E : ImportedModule::kAlwaysImport);
78 :
79 : // Setup the add-imports transform.
80 E : PEAddImportsTransform add_imports;
81 E : add_imports.AddModule(&module);
82 :
83 : // Add the imports to the block-graph.
84 : if (!ApplyBlockGraphTransform(
85 E : &add_imports, policy, block_graph, header_block)) {
86 i : LOG(ERROR) << "Unable to add import entry for jump table hook functions.";
87 i : return false;
88 : }
89 :
90 : // Get a reference to the hook function.
91 E : if (!module.GetSymbolReference(index_case_counter, jump_table_case_counter)) {
92 i : LOG(ERROR) << "Unable to get jump table hooks.";
93 i : return false;
94 : }
95 E : DCHECK(jump_table_case_counter->IsValid());
96 :
97 E : return true;
98 E : }
99 :
100 : } // namespace
101 :
102 : const char JumpTableCaseCountTransform::kTransformName[] =
103 : "JumpTableCountTransform";
104 :
105 : JumpTableCaseCountTransform::JumpTableCaseCountTransform()
106 : : add_frequency_data_(common::kJumpTableCountAgentId,
107 : "Jump Table Frequency Data",
108 : common::kJumpTableFrequencyDataVersion,
109 : common::IndexedFrequencyData::JUMP_TABLE,
110 : sizeof(common::IndexedFrequencyData)),
111 : instrument_dll_name_(kDefaultModuleName),
112 E : jump_table_case_count_(0) {
113 E : }
114 :
115 : bool JumpTableCaseCountTransform::PreBlockGraphIteration(
116 : const TransformPolicyInterface* policy,
117 : BlockGraph* block_graph,
118 E : BlockGraph::Block* header_block) {
119 E : DCHECK(policy != NULL);
120 E : DCHECK(block_graph != NULL);
121 E : DCHECK(header_block != NULL);
122 :
123 : // Setup the jump table counter entry hook.
124 : if (!SetupCounterHook(policy,
125 : block_graph,
126 : header_block,
127 : instrument_dll_name_,
128 E : &jump_table_case_counter_hook_ref_)) {
129 i : return false;
130 : }
131 :
132 : // Add the static jump table count frequency data.
133 : if (!ApplyBlockGraphTransform(&add_frequency_data_,
134 : policy,
135 : block_graph,
136 E : header_block)) {
137 i : LOG(ERROR) << "Failed to insert jump table count frequency data.";
138 i : return false;
139 : }
140 :
141 : // Find or create the section we put our thunks in.
142 : thunk_section_ = block_graph->FindOrAddSection(common::kThunkSectionName,
143 E : pe::kCodeCharacteristics);
144 E : DCHECK(thunk_section_ != NULL);
145 :
146 E : return true;
147 E : }
148 :
149 : bool JumpTableCaseCountTransform::OnBlock(
150 : const TransformPolicyInterface* policy,
151 : BlockGraph* block_graph,
152 E : BlockGraph::Block* block) {
153 E : DCHECK(policy != NULL);
154 E : DCHECK(block_graph != NULL);
155 E : DCHECK(block != NULL);
156 :
157 E : if (block->type() != BlockGraph::CODE_BLOCK)
158 E : return true;
159 :
160 : // Iterate over the labels of the block to find the jump tables.
161 : for (BlockGraph::Block::LabelMap::const_iterator iter_label(
162 E : block->labels().begin());
163 E : iter_label != block->labels().end();
164 E : ++iter_label) {
165 E : if (!iter_label->second.has_attributes(BlockGraph::JUMP_TABLE_LABEL))
166 E : continue;
167 :
168 E : size_t table_size = 0;
169 E : if (!block_graph::GetJumpTableSize(block, iter_label, &table_size))
170 i : return false;
171 :
172 : jump_table_infos_.push_back(
173 E : std::make_pair(block->addr() + iter_label->first, table_size));
174 :
175 : BlockGraph::Block::ReferenceMap::const_iterator iter_ref =
176 E : block->references().find(iter_label->first);
177 :
178 : // Iterate over the references and thunk them.
179 E : for (size_t i = 0; i < table_size; ++i) {
180 E : DCHECK(iter_ref != block->references().end());
181 :
182 : BlockGraph::Block* thunk_block = CreateOneThunk(block_graph,
183 E : iter_ref->second);
184 E : if (thunk_block == NULL) {
185 i : jump_table_infos_.pop_back();
186 i : return false;
187 : }
188 :
189 : BlockGraph::Reference thunk_ref(BlockGraph::ABSOLUTE_REF,
190 : sizeof(iter_ref->second.size()),
191 E : thunk_block, 0, 0);
192 E : block->SetReference(iter_ref->first, thunk_ref);
193 E : ++iter_ref;
194 E : }
195 E : }
196 :
197 E : return true;
198 E : }
199 :
200 : bool JumpTableCaseCountTransform::PostBlockGraphIteration(
201 : const TransformPolicyInterface* policy,
202 : BlockGraph* block_graph,
203 E : BlockGraph::Block* header_block) {
204 E : DCHECK(policy != NULL);
205 E : DCHECK(block_graph != NULL);
206 E : DCHECK(header_block != NULL);
207 :
208 E : if (jump_table_case_count_ == 0) {
209 i : LOG(INFO) << "Encountered no jump tables during instrumentation.";
210 i : return true;
211 : }
212 :
213 : if (!add_frequency_data_.ConfigureFrequencyDataBuffer(jump_table_case_count_,
214 : 1,
215 E : sizeof(uint32))) {
216 i : LOG(ERROR) << "Failed to configure frequency data buffer.";
217 i : return false;
218 : }
219 :
220 : // Add the module entry thunks.
221 E : EntryThunkTransform add_thunks;
222 E : add_thunks.set_only_instrument_module_entry(true);
223 E : add_thunks.set_instrument_dll_name(instrument_dll_name_);
224 :
225 E : Immediate module_data(add_frequency_data_.frequency_data_block(), 0);
226 E : if (!add_thunks.SetEntryThunkParameter(module_data)) {
227 i : LOG(ERROR) << "Failed to configure the entry thunks with the module_data "
228 : << "parameter.";
229 i : return false;
230 : }
231 :
232 : if (!ApplyBlockGraphTransform(
233 E : &add_thunks, policy, block_graph, header_block)) {
234 i : LOG(ERROR) << "Unable to thunk module entry points.";
235 i : return false;
236 : }
237 :
238 E : return true;
239 E : }
240 :
241 : BlockGraph::Block* JumpTableCaseCountTransform::CreateOneThunk(
242 : BlockGraph* block_graph,
243 E : const BlockGraph::Reference& destination) {
244 : // Construct the name for the new thunk.
245 E : std::string thunk_name(destination.referenced()->name() + kThunkSuffix);
246 :
247 : Operand jump_table_case_counter_hook(
248 : Displacement(jump_table_case_counter_hook_ref_.referenced(),
249 E : jump_table_case_counter_hook_ref_.offset()));
250 :
251 : // Construct the thunk basic block.
252 E : BasicBlockSubGraph bbsg;
253 : BasicBlockSubGraph::BlockDescription* block_desc = bbsg.AddBlockDescription(
254 : thunk_name,
255 : NULL,
256 : BlockGraph::CODE_BLOCK,
257 : thunk_section_->id(),
258 : 1,
259 E : 0);
260 E : BasicCodeBlock* bb = bbsg.AddBasicCodeBlock(thunk_name);
261 E : block_desc->basic_block_order.push_back(bb);
262 :
263 E : BasicBlockAssembler assm(bb->instructions().begin(), &bb->instructions());
264 E : DCHECK_LT(jump_table_case_count_, std::numeric_limits<size_t>::max());
265 E : assm.push(Immediate(jump_table_case_count_++, core::kSize32Bit));
266 E : assm.call(jump_table_case_counter_hook);
267 E : assm.jmp(Immediate(destination.referenced(), destination.offset()));
268 :
269 : // Condense into a block.
270 E : BlockBuilder block_builder(block_graph);
271 E : if (!block_builder.Merge(&bbsg)) {
272 i : LOG(ERROR) << "Failed to build thunk block.";
273 i : return NULL;
274 : }
275 :
276 : // Exactly one new block should have been created.
277 E : DCHECK_EQ(1u, block_builder.new_blocks().size());
278 E : return block_builder.new_blocks().front();
279 E : }
280 :
281 : } // namespace transforms
282 : } // namespace instrument
|