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[] = "jump_table_count_client.dll";
53 : const char kJumpTableCaseCounter[] = "_jump_table_case_counter";
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 : instrument_dll_name_(kDefaultModuleName),
106 E : jump_table_case_count_(0) {
107 E : }
108 :
109 : bool JumpTableCaseCountTransform::PreBlockGraphIteration(
110 : BlockGraph* block_graph,
111 E : BlockGraph::Block* header_block) {
112 E : DCHECK(block_graph != NULL);
113 E : DCHECK(header_block != NULL);
114 :
115 : // Setup the jump table counter entry hook.
116 : if (!SetupCounterHook(block_graph,
117 : header_block,
118 : instrument_dll_name_,
119 E : &jump_table_case_counter_hook_ref_)) {
120 i : return false;
121 : }
122 :
123 : // Add the static jump table count frequency data.
124 : if (!ApplyBlockGraphTransform(&add_frequency_data_,
125 : block_graph,
126 E : header_block)) {
127 i : LOG(ERROR) << "Failed to insert jump table count frequency data.";
128 i : return false;
129 : }
130 :
131 : // Find or create the section we put our thunks in.
132 : thunk_section_ = block_graph->FindOrAddSection(common::kThunkSectionName,
133 E : pe::kCodeCharacteristics);
134 E : DCHECK(thunk_section_ != NULL);
135 :
136 E : return true;
137 E : }
138 :
139 : bool JumpTableCaseCountTransform::OnBlock(BlockGraph* block_graph,
140 E : BlockGraph::Block* block) {
141 E : DCHECK(block_graph != NULL);
142 E : DCHECK(block != NULL);
143 :
144 E : if (block->type() != BlockGraph::CODE_BLOCK)
145 E : return true;
146 :
147 : // Iterate over the labels of the block to find the jump tables.
148 : for (BlockGraph::Block::LabelMap::const_iterator iter_label(
149 E : block->labels().begin());
150 E : iter_label != block->labels().end();
151 E : ++iter_label) {
152 E : if (!iter_label->second.has_attributes(BlockGraph::JUMP_TABLE_LABEL))
153 E : continue;
154 :
155 E : size_t table_size = 0;
156 E : if (!block_graph::GetJumpTableSize(block, iter_label, &table_size))
157 i : return false;
158 :
159 : jump_table_infos_.push_back(
160 E : std::make_pair(block->addr() + iter_label->first, table_size));
161 :
162 : BlockGraph::Block::ReferenceMap::const_iterator iter_ref =
163 E : block->references().find(iter_label->first);
164 :
165 : // Iterate over the references and thunk them.
166 E : for (size_t i = 0; i < table_size; ++i) {
167 E : DCHECK(iter_ref != block->references().end());
168 :
169 : BlockGraph::Block* thunk_block = CreateOneThunk(block_graph,
170 E : iter_ref->second);
171 E : if (thunk_block == NULL) {
172 i : jump_table_infos_.pop_back();
173 i : return false;
174 : }
175 :
176 : BlockGraph::Reference thunk_ref(BlockGraph::ABSOLUTE_REF,
177 : sizeof(iter_ref->second.size()),
178 E : thunk_block, 0, 0);
179 E : block->SetReference(iter_ref->first, thunk_ref);
180 E : ++iter_ref;
181 E : }
182 E : }
183 :
184 E : return true;
185 E : }
186 :
187 : bool JumpTableCaseCountTransform::PostBlockGraphIteration(
188 : BlockGraph* block_graph,
189 E : BlockGraph::Block* header_block) {
190 E : DCHECK(block_graph != NULL);
191 E : DCHECK(header_block != NULL);
192 :
193 E : if (jump_table_case_count_ == 0) {
194 i : LOG(INFO) << "Encountered no jump tables during instrumentation.";
195 i : return true;
196 : }
197 :
198 : if (!add_frequency_data_.ConfigureFrequencyDataBuffer(jump_table_case_count_,
199 E : sizeof(uint32))) {
200 i : LOG(ERROR) << "Failed to configure frequency data buffer.";
201 i : return false;
202 : }
203 :
204 : // Add the module entry thunks.
205 E : EntryThunkTransform add_thunks;
206 E : add_thunks.set_only_instrument_module_entry(true);
207 E : add_thunks.set_instrument_dll_name(instrument_dll_name_);
208 :
209 E : Immediate module_data(add_frequency_data_.frequency_data_block(), 0);
210 E : if (!add_thunks.SetEntryThunkParameter(module_data)) {
211 i : LOG(ERROR) << "Failed to configure the entry thunks with the module_data "
212 : << "parameter.";
213 i : return false;
214 : }
215 :
216 E : if (!ApplyBlockGraphTransform(&add_thunks, block_graph, header_block)) {
217 i : LOG(ERROR) << "Unable to thunk module entry points.";
218 i : return false;
219 : }
220 :
221 E : return true;
222 E : }
223 :
224 : BlockGraph::Block* JumpTableCaseCountTransform::CreateOneThunk(
225 : BlockGraph* block_graph,
226 E : const BlockGraph::Reference& destination) {
227 : // Construct the name for the new thunk.
228 E : std::string thunk_name(destination.referenced()->name() + kThunkSuffix);
229 :
230 : Operand jump_table_case_counter_hook(
231 : Displacement(jump_table_case_counter_hook_ref_.referenced(),
232 E : jump_table_case_counter_hook_ref_.offset()));
233 :
234 : // Construct the thunk basic block.
235 E : BasicBlockSubGraph bbsg;
236 : BasicBlockSubGraph::BlockDescription* block_desc = bbsg.AddBlockDescription(
237 E : thunk_name, BlockGraph::CODE_BLOCK, thunk_section_->id(), 1, 0);
238 E : BasicCodeBlock* bb = bbsg.AddBasicCodeBlock(thunk_name);
239 E : block_desc->basic_block_order.push_back(bb);
240 :
241 E : BasicBlockAssembler assm(bb->instructions().begin(), &bb->instructions());
242 E : DCHECK_LT(jump_table_case_count_, std::numeric_limits<size_t>::max());
243 E : assm.push(Immediate(jump_table_case_count_++, core::kSize32Bit));
244 E : assm.call(jump_table_case_counter_hook);
245 E : assm.jmp(Immediate(destination.referenced(), destination.offset()));
246 :
247 : // Condense into a block.
248 E : BlockBuilder block_builder(block_graph);
249 E : if (!block_builder.Merge(&bbsg)) {
250 i : LOG(ERROR) << "Failed to build thunk block.";
251 i : return NULL;
252 : }
253 :
254 : // Exactly one new block should have been created.
255 E : DCHECK_EQ(1u, block_builder.new_blocks().size());
256 E : return block_builder.new_blocks().front();
257 E : }
258 :
259 : } // namespace transforms
260 : } // namespace instrument
|