1 : // Copyright 2012 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 : // Basic-block entry hook instrumentation transform unit-tests.
16 :
17 : #include "syzygy/instrument/transforms/basic_block_entry_hook_transform.h"
18 :
19 : #include "gtest/gtest.h"
20 : #include "syzygy/block_graph/basic_block.h"
21 : #include "syzygy/block_graph/basic_block_decomposer.h"
22 : #include "syzygy/block_graph/basic_block_subgraph.h"
23 : #include "syzygy/block_graph/block_graph.h"
24 : #include "syzygy/block_graph/typed_block.h"
25 : #include "syzygy/common/indexed_frequency_data.h"
26 : #include "syzygy/core/unittest_util.h"
27 : #include "syzygy/instrument/transforms/unittest_util.h"
28 : #include "syzygy/pe/block_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 block_graph::BasicBlock;
38 : using block_graph::BasicCodeBlock;
39 : using block_graph::BasicBlockDecomposer;
40 : using block_graph::BasicBlockSubGraph;
41 : using block_graph::BlockGraph;
42 : using block_graph::Instruction;
43 : using common::IndexedFrequencyData;
44 : using common::kBasicBlockEntryAgentId;
45 : using common::kBasicBlockFrequencyDataVersion;
46 :
47 : class TestBasicBlockEntryHookTransform : public BasicBlockEntryHookTransform {
48 : public:
49 : using BasicBlockEntryHookTransform::bb_entry_hook_ref_;
50 : using BasicBlockEntryHookTransform::fast_bb_entry_block_;
51 : using BasicBlockEntryHookTransform::thunk_section_;
52 :
53 E : BlockGraph::Block* frequency_data_block() {
54 E : return add_frequency_data_.frequency_data_block();
55 E : }
56 :
57 E : BlockGraph::Block* frequency_data_buffer_block() {
58 E : return add_frequency_data_.frequency_data_buffer_block();
59 E : }
60 : };
61 :
62 : class BasicBlockEntryHookTransformTest : public testing::TestDllTransformTest {
63 : public:
64 : enum InstrumentationKind {
65 : kAgentInstrumentation,
66 : kFastPathInstrumentation
67 : };
68 :
69 : void CheckBasicBlockInstrumentation(InstrumentationKind kind);
70 :
71 : protected:
72 : TestBasicBlockEntryHookTransform tx_;
73 : };
74 :
75 : } // namespace
76 :
77 E : TEST_F(BasicBlockEntryHookTransformTest, SetInlinePathFlag) {
78 E : EXPECT_FALSE(tx_.inline_fast_path());
79 E : tx_.set_inline_fast_path(true);
80 E : EXPECT_TRUE(tx_.inline_fast_path());
81 E : tx_.set_inline_fast_path(false);
82 E : EXPECT_FALSE(tx_.inline_fast_path());
83 E : }
84 :
85 : void BasicBlockEntryHookTransformTest::CheckBasicBlockInstrumentation(
86 E : InstrumentationKind kind) {
87 : // Let's examine each eligible block to verify that its BB's have been
88 : // instrumented.
89 E : size_t num_decomposed_blocks = 0;
90 E : size_t total_basic_blocks = 0;
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 : // Blocks which are not bb-decomposable should be thunked. While there may
105 : // be some internal referrers, the only external referrers should be thunks.
106 E : if (!pe::CodeBlockIsBasicBlockDecomposable(&block)) {
107 E : size_t num_external_thunks = 0;
108 : BlockGraph::Block::ReferrerSet::const_iterator ref_iter =
109 E : block.referrers().begin();
110 E : for (; ref_iter != block.referrers().end(); ++ref_iter) {
111 E : if (ref_iter->first != &block) {
112 E : ASSERT_EQ(tx_.thunk_section_->id(), ref_iter->first->section());
113 E : ++num_external_thunks;
114 : }
115 E : }
116 :
117 : // Each of the thunks for a non-decomposable block will reuse the same
118 : // id to source range map entry, so we increment total_basic_blocks once
119 : // if num_external_thunks is non-zero. Note that we cannot assert that
120 : // num_external_thunks > 0 because the block could be statically dead
121 : // (in a debug build, for example).
122 E : if (num_external_thunks != 0U)
123 E : ++total_basic_blocks;
124 E : continue;
125 : }
126 :
127 : // Note that we have attempted to validate a block.
128 E : ++num_decomposed_blocks;
129 :
130 : // Decompose the block to basic-blocks.
131 E : BasicBlockSubGraph subgraph;
132 E : BasicBlockDecomposer bb_decomposer(&block, &subgraph);
133 E : ASSERT_TRUE(bb_decomposer.Decompose());
134 :
135 : // Check if each non-padding basic code-block begins with the
136 : // instrumentation sequence.
137 E : size_t num_basic_blocks = 0;
138 : BasicBlockSubGraph::BBCollection::const_iterator bb_iter =
139 E : subgraph.basic_blocks().begin();
140 E : for (; bb_iter != subgraph.basic_blocks().end(); ++bb_iter) {
141 E : const BasicCodeBlock* bb = BasicCodeBlock::Cast(*bb_iter);
142 E : if (bb == NULL || bb->is_padding())
143 E : continue;
144 E : ++num_basic_blocks;
145 :
146 E : if (kind == kAgentInstrumentation) {
147 E : ASSERT_LE(3U, bb->instructions().size());
148 : BasicBlock::Instructions::const_iterator inst_iter =
149 E : bb->instructions().begin();
150 :
151 : // Instruction 1 should push the basic block id.
152 E : const Instruction& inst1 = *inst_iter;
153 E : EXPECT_EQ(I_PUSH, inst1.representation().opcode);
154 :
155 : // Instruction 2 should push the frequency data block pointer.
156 E : const Instruction& inst2 = *(++inst_iter);
157 E : EXPECT_EQ(I_PUSH, inst2.representation().opcode);
158 E : ASSERT_EQ(1U, inst2.references().size());
159 : EXPECT_EQ(tx_.frequency_data_block(),
160 E : inst2.references().begin()->second.block());
161 :
162 : // Instruction 3 should be a call to the bb entry hook.
163 E : const Instruction& inst3 = *(++inst_iter);
164 E : EXPECT_EQ(I_CALL, inst3.representation().opcode);
165 E : ASSERT_EQ(1U, inst3.references().size());
166 : EXPECT_EQ(tx_.bb_entry_hook_ref_.referenced(),
167 E : inst3.references().begin()->second.block());
168 E : } else {
169 E : DCHECK(kind == kFastPathInstrumentation);
170 E : ASSERT_LE(2U, bb->instructions().size());
171 : BasicBlock::Instructions::const_iterator inst_iter =
172 E : bb->instructions().begin();
173 :
174 : // Instruction 1 should push the basic block id.
175 E : const Instruction& inst1 = *inst_iter;
176 E : EXPECT_EQ(I_PUSH, inst1.representation().opcode);
177 :
178 : // Instruction 2 should be a call to the fast bb entry hook.
179 E : const Instruction& inst2 = *(++inst_iter);
180 E : EXPECT_EQ(I_CALL, inst2.representation().opcode);
181 E : ASSERT_EQ(1U, inst2.references().size());
182 : EXPECT_EQ(tx_.fast_bb_entry_block_,
183 E : inst2.references().begin()->second.block());
184 : }
185 E : }
186 E : EXPECT_NE(0U, num_basic_blocks);
187 E : total_basic_blocks += num_basic_blocks;
188 E : }
189 :
190 E : EXPECT_NE(0U, num_decomposed_blocks);
191 E : EXPECT_EQ(total_basic_blocks, tx_.bb_ranges().size());
192 E : }
193 :
194 E : TEST_F(BasicBlockEntryHookTransformTest, ApplyAgentInstrumentation) {
195 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
196 :
197 : // Apply the transform.
198 E : tx_.set_src_ranges_for_thunks(true);
199 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(&tx_, &block_graph_,
200 E : dos_header_block_));
201 E : ASSERT_TRUE(tx_.frequency_data_block() != NULL);
202 E : ASSERT_TRUE(tx_.thunk_section_ != NULL);
203 E : ASSERT_TRUE(tx_.bb_entry_hook_ref_.IsValid());
204 E : ASSERT_LT(0u, tx_.bb_ranges().size());
205 :
206 : // Validate the basic-block frequency data structure.
207 E : block_graph::ConstTypedBlock<IndexedFrequencyData> frequency_data;
208 E : ASSERT_TRUE(frequency_data.Init(0, tx_.frequency_data_block()));
209 E : EXPECT_EQ(kBasicBlockEntryAgentId, frequency_data->agent_id);
210 E : EXPECT_EQ(kBasicBlockFrequencyDataVersion, frequency_data->version);
211 E : EXPECT_EQ(tx_.bb_ranges().size(), frequency_data->num_entries);
212 E : EXPECT_EQ(sizeof(uint32), frequency_data->frequency_size);
213 : EXPECT_TRUE(frequency_data.HasReferenceAt(
214 E : frequency_data.OffsetOf(frequency_data->frequency_data)));
215 E : EXPECT_EQ(sizeof(IndexedFrequencyData), tx_.frequency_data_block()->size());
216 : EXPECT_EQ(sizeof(IndexedFrequencyData),
217 E : tx_.frequency_data_block()->data_size());
218 : EXPECT_EQ(frequency_data->num_entries * frequency_data->frequency_size,
219 E : tx_.frequency_data_buffer_block()->size());
220 :
221 : // Validate that all basic block have been instrumented.
222 E : CheckBasicBlockInstrumentation(kAgentInstrumentation);
223 E : }
224 :
225 E : TEST_F(BasicBlockEntryHookTransformTest, ApplyFastPathInstrumentation) {
226 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
227 :
228 : // Apply the transform.
229 E : tx_.set_inline_fast_path(true);
230 :
231 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(&tx_, &block_graph_,
232 E : dos_header_block_));
233 E : ASSERT_TRUE(tx_.fast_bb_entry_block_ != NULL);
234 :
235 : // Validate that all basic block have been instrumented.
236 E : CheckBasicBlockInstrumentation(kFastPathInstrumentation);
237 E : }
238 :
239 : } // namespace transforms
240 : } // namespace instrument
|