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/strings/string_util.h"
23 : #include "base/strings/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 E : size_t index_case_counter = module.AddSymbol(
76 : kJumpTableCaseCounter,
77 : 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 E : if (!ApplyBlockGraphTransform(
85 : &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 E : : add_frequency_data_(common::kJumpTableCountAgentId,
107 : "Jump Table Frequency Data",
108 : common::kJumpTableFrequencyDataVersion,
109 : common::IndexedFrequencyData::JUMP_TABLE,
110 : sizeof(common::IndexedFrequencyData)),
111 E : 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_NE(reinterpret_cast<TransformPolicyInterface*>(NULL), policy);
120 E : DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
121 E : DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), header_block);
122 E : DCHECK_EQ(BlockGraph::PE_IMAGE, block_graph->image_format());
123 :
124 : // Setup the jump table counter entry hook.
125 E : if (!SetupCounterHook(policy,
126 : block_graph,
127 : header_block,
128 : instrument_dll_name_,
129 : &jump_table_case_counter_hook_ref_)) {
130 i : return false;
131 : }
132 :
133 : // Add the static jump table count frequency data.
134 E : if (!ApplyBlockGraphTransform(&add_frequency_data_,
135 : policy,
136 : block_graph,
137 : header_block)) {
138 i : LOG(ERROR) << "Failed to insert jump table count frequency data.";
139 i : return false;
140 : }
141 :
142 : // Find or create the section we put our thunks in.
143 E : thunk_section_ = block_graph->FindOrAddSection(common::kThunkSectionName,
144 : pe::kCodeCharacteristics);
145 E : DCHECK(thunk_section_ != NULL);
146 :
147 E : return true;
148 E : }
149 :
150 : bool JumpTableCaseCountTransform::OnBlock(
151 : const TransformPolicyInterface* policy,
152 : BlockGraph* block_graph,
153 E : BlockGraph::Block* block) {
154 E : DCHECK(policy != NULL);
155 E : DCHECK(block_graph != NULL);
156 E : DCHECK(block != NULL);
157 :
158 E : if (block->type() != BlockGraph::CODE_BLOCK)
159 E : return true;
160 :
161 : // Iterate over the labels of the block to find the jump tables.
162 : for (BlockGraph::Block::LabelMap::const_iterator iter_label(
163 E : block->labels().begin());
164 E : iter_label != block->labels().end();
165 E : ++iter_label) {
166 E : if (!iter_label->second.has_attributes(BlockGraph::JUMP_TABLE_LABEL))
167 E : continue;
168 :
169 E : size_t table_size = 0;
170 E : if (!block_graph::GetJumpTableSize(block, iter_label, &table_size))
171 i : return false;
172 :
173 E : jump_table_infos_.push_back(
174 : std::make_pair(block->addr() + iter_label->first, table_size));
175 :
176 : BlockGraph::Block::ReferenceMap::const_iterator iter_ref =
177 E : block->references().find(iter_label->first);
178 :
179 : // Iterate over the references and thunk them.
180 E : for (size_t i = 0; i < table_size; ++i) {
181 E : DCHECK(iter_ref != block->references().end());
182 :
183 E : BlockGraph::Block* thunk_block = CreateOneThunk(block_graph,
184 : iter_ref->second);
185 E : if (thunk_block == NULL) {
186 i : jump_table_infos_.pop_back();
187 i : return false;
188 : }
189 :
190 E : BlockGraph::Reference thunk_ref(BlockGraph::ABSOLUTE_REF,
191 : sizeof(iter_ref->second.size()),
192 : thunk_block, 0, 0);
193 E : block->SetReference(iter_ref->first, thunk_ref);
194 E : ++iter_ref;
195 E : }
196 E : }
197 :
198 E : return true;
199 E : }
200 :
201 : bool JumpTableCaseCountTransform::PostBlockGraphIteration(
202 : const TransformPolicyInterface* policy,
203 : BlockGraph* block_graph,
204 E : BlockGraph::Block* header_block) {
205 E : DCHECK(policy != NULL);
206 E : DCHECK(block_graph != NULL);
207 E : DCHECK(header_block != NULL);
208 :
209 E : if (jump_table_case_count_ == 0) {
210 i : LOG(INFO) << "Encountered no jump tables during instrumentation.";
211 i : return true;
212 : }
213 :
214 E : if (!add_frequency_data_.ConfigureFrequencyDataBuffer(jump_table_case_count_,
215 : 1, sizeof(uint32_t))) {
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 : auto module_data(Immediate(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 E : if (!ApplyBlockGraphTransform(
233 : &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 : auto jump_table_case_counter_hook(
248 E : Operand(Displacement(jump_table_case_counter_hook_ref_.referenced(),
249 : jump_table_case_counter_hook_ref_.offset())));
250 :
251 : // Construct the thunk basic block.
252 E : BasicBlockSubGraph bbsg;
253 E : BasicBlockSubGraph::BlockDescription* block_desc = bbsg.AddBlockDescription(
254 : thunk_name,
255 : NULL,
256 : BlockGraph::CODE_BLOCK,
257 : thunk_section_->id(),
258 : 1,
259 : 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_++, assm::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
|