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