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 : // Branch hook instrumentation transform unit-tests.
16 :
17 : #include "syzygy/instrument/transforms/branch_hook_transform.h"
18 :
19 : #include "gtest/gtest.h"
20 : #include "syzygy/agent/basic_block_entry/basic_block_entry.h"
21 : #include "syzygy/block_graph/basic_block.h"
22 : #include "syzygy/block_graph/basic_block_decomposer.h"
23 : #include "syzygy/block_graph/basic_block_subgraph.h"
24 : #include "syzygy/block_graph/block_graph.h"
25 : #include "syzygy/block_graph/typed_block.h"
26 : #include "syzygy/common/indexed_frequency_data.h"
27 : #include "syzygy/core/unittest_util.h"
28 : #include "syzygy/instrument/transforms/unittest_util.h"
29 : #include "syzygy/pe/unittest_util.h"
30 :
31 : #include "mnemonics.h" // NOLINT
32 :
33 : namespace instrument {
34 : namespace transforms {
35 : namespace {
36 :
37 : using agent::basic_block_entry::BasicBlockEntry;
38 : using block_graph::BasicBlock;
39 : using block_graph::BasicBlockDecomposer;
40 : using block_graph::BasicBlockSubGraph;
41 : using block_graph::BasicCodeBlock;
42 : using block_graph::BlockGraph;
43 : using block_graph::Instruction;
44 : using common::IndexedFrequencyData;
45 : using common::ThreadLocalIndexedFrequencyData;
46 : using common::kBasicBlockEntryAgentId;
47 : using common::kBasicBlockFrequencyDataVersion;
48 :
49 : class TestBranchHookTransform : public BranchHookTransform {
50 : public:
51 : using BranchHookTransform::function_enter_hook_ref_;
52 : using BranchHookTransform::enter_hook_ref_;
53 : using BranchHookTransform::exit_hook_ref_;
54 : using BranchHookTransform::thunk_section_;
55 : using BranchHookTransform::buffering_;
56 : using BranchHookTransform::fs_slot_;
57 :
58 E : BlockGraph::Block* frequency_data_block() {
59 E : return add_frequency_data_.frequency_data_block();
60 E : }
61 :
62 E : BlockGraph::Block* frequency_data_buffer_block() {
63 E : return add_frequency_data_.frequency_data_buffer_block();
64 E : }
65 : };
66 :
67 : class BranchHookTransformTest : public testing::TestDllTransformTest {
68 : public:
69 : void CheckBasicBlockInstrumentation();
70 :
71 : protected:
72 : TestBranchHookTransform tx_;
73 : };
74 :
75 E : void BranchHookTransformTest::CheckBasicBlockInstrumentation() {
76 : // Determine whether the module_data should be passed as an argument.
77 E : bool need_module_data = true;
78 E : if (tx_.fs_slot_ != 0)
79 E : need_module_data = false;
80 :
81 : // Determine whether a function_enter hook should be added at the beginning of
82 : // each instrumented block.
83 E : bool need_function_enter = false;
84 E : if (tx_.function_enter_hook_ref_.IsValid())
85 E : need_function_enter = true;
86 :
87 : // Let's examine each eligible block to verify that its basic blocks have been
88 : // instrumented.
89 : BlockGraph::BlockMap::const_iterator block_iter =
90 E : block_graph_.blocks().begin();
91 E : for (; block_iter != block_graph_.blocks().end(); ++block_iter) {
92 E : const BlockGraph::Block& block = block_iter->second;
93 :
94 : // Skip non-code blocks.
95 E : if (block.type() != BlockGraph::CODE_BLOCK)
96 E : continue;
97 :
98 : // We'll skip thunks, they're a mixed bag of things.
99 E : if (block.section() == tx_.thunk_section_->id())
100 E : continue;
101 :
102 : // Skip non-decomposable blocks.
103 E : if (!policy_->BlockIsSafeToBasicBlockDecompose(&block))
104 E : continue;
105 :
106 : // Decompose the block to basic-blocks.
107 E : BasicBlockSubGraph subgraph;
108 E : BasicBlockDecomposer bb_decomposer(&block, &subgraph);
109 E : ASSERT_TRUE(bb_decomposer.Decompose());
110 :
111 : // Retrieve the first basic block.
112 E : DCHECK_EQ(1U, subgraph.block_descriptions().size());
113 : const BasicBlockSubGraph::BasicBlockOrdering& original_order =
114 E : subgraph.block_descriptions().front().basic_block_order;
115 E : BasicCodeBlock* first_bb = BasicCodeBlock::Cast(*original_order.begin());
116 E : DCHECK(first_bb != NULL);
117 :
118 : // Check if each non-padding basic code-block begins with the
119 : // instrumentation sequence.
120 : BasicBlockSubGraph::BBCollection::const_iterator bb_iter =
121 E : subgraph.basic_blocks().begin();
122 E : for (; bb_iter != subgraph.basic_blocks().end(); ++bb_iter) {
123 E : const BasicCodeBlock* bb = BasicCodeBlock::Cast(*bb_iter);
124 E : if (bb == NULL || bb->is_padding())
125 E : continue;
126 :
127 : // Check entry hook function call.
128 : BasicBlock::Instructions::const_iterator inst_iter =
129 E : bb->instructions().begin();
130 E : ASSERT_TRUE(inst_iter != bb->instructions().end());
131 :
132 E : if (need_function_enter && bb == first_bb) {
133 : // Instruction 1 should push the frequency data block pointer.
134 E : const Instruction& inst1 = *inst_iter;
135 E : EXPECT_EQ(I_PUSH, inst1.representation().opcode);
136 E : ASSERT_EQ(1U, inst1.references().size());
137 : EXPECT_EQ(tx_.frequency_data_block(),
138 E : inst1.references().begin()->second.block());
139 E : ASSERT_TRUE(++inst_iter != bb->instructions().end());
140 :
141 : // Instruction 2 should be a call to the function enter hook.
142 E : const Instruction& inst2 = *inst_iter;
143 E : EXPECT_EQ(I_CALL, inst2.representation().opcode);
144 E : ASSERT_EQ(1U, inst2.references().size());
145 : EXPECT_EQ(tx_.function_enter_hook_ref_.referenced(),
146 E : inst2.references().begin()->second.block());
147 E : ASSERT_TRUE(++inst_iter != bb->instructions().end());
148 : }
149 :
150 : // Instruction 1 should push the basic block id.
151 E : const Instruction& inst1 = *inst_iter;
152 E : EXPECT_EQ(I_PUSH, inst1.representation().opcode);
153 E : ASSERT_EQ(0U, inst1.references().size());
154 E : ASSERT_TRUE(++inst_iter != bb->instructions().end());
155 :
156 : // Instruction 2 should push the frequency data block pointer.
157 E : if (need_module_data) {
158 E : const Instruction& inst2 = *inst_iter;
159 E : EXPECT_EQ(I_PUSH, inst2.representation().opcode);
160 E : ASSERT_EQ(1U, inst2.references().size());
161 : EXPECT_EQ(tx_.frequency_data_block(),
162 E : inst2.references().begin()->second.block());
163 E : ASSERT_TRUE(++inst_iter != bb->instructions().end());
164 : }
165 :
166 : // Instruction 3 should be a call to the enter hook.
167 E : const Instruction& inst3 = *inst_iter;
168 E : EXPECT_EQ(I_CALL, inst3.representation().opcode);
169 E : ASSERT_EQ(1U, inst3.references().size());
170 : EXPECT_EQ(tx_.enter_hook_ref_.referenced(),
171 E : inst3.references().begin()->second.block());
172 E : ASSERT_TRUE(++inst_iter != bb->instructions().end());
173 :
174 : // Check exit hook function call.
175 : BasicBlock::Instructions::const_reverse_iterator rev_inst_iter =
176 E : bb->instructions().rbegin();
177 E : ASSERT_TRUE(rev_inst_iter != bb->instructions().rend());
178 :
179 : // Find last non branching instruction.
180 E : for (; rev_inst_iter != bb->instructions().rend(); ++rev_inst_iter) {
181 E : if (!rev_inst_iter->IsBranch() && !rev_inst_iter->IsReturn())
182 E : break;
183 E : }
184 E : ASSERT_TRUE(rev_inst_iter != bb->instructions().rend());
185 :
186 : // Skip non returning basic block.
187 E : if (rev_inst_iter->CallsNonReturningFunction())
188 E : continue;
189 :
190 : // Instruction 3 should be a call to the exit hook.
191 E : const Instruction& rev_inst3 = *rev_inst_iter;
192 E : EXPECT_EQ(I_CALL, rev_inst3.representation().opcode);
193 E : ASSERT_EQ(1U, rev_inst3.references().size());
194 : EXPECT_EQ(tx_.enter_hook_ref_.referenced(),
195 E : rev_inst3.references().begin()->second.block());
196 E : ASSERT_TRUE(++rev_inst_iter != bb->instructions().rend());
197 :
198 : // Instruction 2 should push the frequency data block pointer.
199 E : if (need_module_data) {
200 E : const Instruction& rev_inst2 = *rev_inst_iter;
201 E : EXPECT_EQ(I_PUSH, rev_inst2.representation().opcode);
202 E : ASSERT_EQ(1U, rev_inst2.references().size());
203 : EXPECT_EQ(tx_.frequency_data_block(),
204 E : rev_inst2.references().begin()->second.block());
205 E : ASSERT_TRUE(++rev_inst_iter != bb->instructions().rend());
206 : }
207 :
208 : // Instruction 1 should push the basic block id.
209 E : const Instruction& rev_inst1 = *rev_inst_iter;
210 E : EXPECT_EQ(I_PUSH, rev_inst1.representation().opcode);
211 E : ASSERT_TRUE(++rev_inst_iter != bb->instructions().rend());
212 E : }
213 E : }
214 E : }
215 :
216 : } // namespace
217 :
218 E : TEST_F(BranchHookTransformTest, ApplyAgentInstrumentation) {
219 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
220 :
221 : // Apply the transform.
222 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
223 E : &tx_, policy_, &block_graph_, header_block_));
224 E : ASSERT_TRUE(tx_.frequency_data_block() != NULL);
225 E : ASSERT_TRUE(tx_.enter_hook_ref_.IsValid());
226 E : ASSERT_TRUE(tx_.exit_hook_ref_.IsValid());
227 E : ASSERT_LT(0u, tx_.bb_ranges().size());
228 :
229 : // Validate the basic-block frequency data structure.
230 E : block_graph::ConstTypedBlock<IndexedFrequencyData> frequency_data;
231 E : ASSERT_TRUE(frequency_data.Init(0, tx_.frequency_data_block()));
232 E : EXPECT_EQ(kBasicBlockEntryAgentId, frequency_data->agent_id);
233 E : EXPECT_EQ(kBasicBlockFrequencyDataVersion, frequency_data->version);
234 E : EXPECT_EQ(IndexedFrequencyData::BRANCH, frequency_data->data_type);
235 E : EXPECT_EQ(tx_.bb_ranges().size(), frequency_data->num_entries);
236 E : EXPECT_EQ(3U, frequency_data->num_columns);
237 E : EXPECT_EQ(sizeof(uint32), frequency_data->frequency_size);
238 : EXPECT_TRUE(frequency_data.HasReferenceAt(
239 E : frequency_data.OffsetOf(frequency_data->frequency_data)));
240 : EXPECT_EQ(sizeof(ThreadLocalIndexedFrequencyData),
241 E : tx_.frequency_data_block()->size());
242 : EXPECT_EQ(sizeof(ThreadLocalIndexedFrequencyData),
243 E : tx_.frequency_data_block()->data_size());
244 :
245 : uint32 expected_size = frequency_data->num_entries *
246 E : frequency_data->num_columns * frequency_data->frequency_size;
247 E : EXPECT_EQ(expected_size, tx_.frequency_data_buffer_block()->size());
248 :
249 : // Validate that all basic blocks have been instrumented.
250 E : ASSERT_NO_FATAL_FAILURE(CheckBasicBlockInstrumentation());
251 E : }
252 :
253 E : TEST_F(BranchHookTransformTest, ApplyBufferedAgentInstrumentation) {
254 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
255 :
256 : // Activate buffering.
257 E : tx_.set_buffering(true);
258 E : ASSERT_TRUE(tx_.buffering_);
259 :
260 : // Apply the transform.
261 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
262 E : &tx_, policy_, &block_graph_, header_block_));
263 E : ASSERT_FALSE(tx_.function_enter_hook_ref_.IsValid());
264 E : ASSERT_TRUE(tx_.enter_hook_ref_.IsValid());
265 :
266 : // Validate that all basic blocks have been instrumented.
267 E : ASSERT_NO_FATAL_FAILURE(CheckBasicBlockInstrumentation());
268 E : }
269 :
270 E : TEST_F(BranchHookTransformTest, ApplySlotAgentInstrumentation) {
271 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
272 :
273 : // Activate fs-slot.
274 E : tx_.set_buffering(false);
275 E : tx_.set_fs_slot(1);
276 E : ASSERT_FALSE(tx_.buffering_);
277 E : ASSERT_EQ(tx_.fs_slot_, 1U);
278 :
279 : // Apply the transform.
280 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
281 E : &tx_, policy_, &block_graph_, header_block_));
282 E : ASSERT_TRUE(tx_.function_enter_hook_ref_.IsValid());
283 E : ASSERT_TRUE(tx_.enter_hook_ref_.IsValid());
284 :
285 : // Validate that all basic blocks have been instrumented.
286 E : ASSERT_NO_FATAL_FAILURE(CheckBasicBlockInstrumentation());
287 E : }
288 :
289 E : TEST_F(BranchHookTransformTest, ApplySlotBufferedAgentInstrumentation) {
290 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
291 :
292 : // Activate buffering and fs-slot.
293 E : tx_.set_buffering(true);
294 E : tx_.set_fs_slot(1);
295 E : ASSERT_TRUE(tx_.buffering_);
296 E : ASSERT_EQ(tx_.fs_slot_, 1U);
297 :
298 : // Apply the transform.
299 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
300 E : &tx_, policy_, &block_graph_, header_block_));
301 E : ASSERT_TRUE(tx_.function_enter_hook_ref_.IsValid());
302 E : ASSERT_TRUE(tx_.enter_hook_ref_.IsValid());
303 :
304 : // Validate that all basic blocks have been instrumented.
305 E : ASSERT_NO_FATAL_FAILURE(CheckBasicBlockInstrumentation());
306 E : }
307 :
308 : } // namespace transforms
309 : } // namespace instrument
|