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