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/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 TestBasicBlockEntryHookTransform : public BasicBlockEntryHookTransform {
52 : public:
53 : using BasicBlockEntryHookTransform::bb_entry_hook_ref_;
54 : using BasicBlockEntryHookTransform::thunk_section_;
55 :
56 E : BlockGraph::Block* frequency_data_block() {
57 E : return add_frequency_data_.frequency_data_block();
58 E : }
59 :
60 E : BlockGraph::Block* frequency_data_buffer_block() {
61 E : return add_frequency_data_.frequency_data_buffer_block();
62 E : }
63 : };
64 :
65 : class BasicBlockEntryHookTransformTest : public testing::TestDllTransformTest {
66 : public:
67 : enum InstrumentationKind {
68 : kAgentInstrumentation,
69 : kFastPathInstrumentation
70 : };
71 :
72 : void CheckBasicBlockInstrumentation(InstrumentationKind kind);
73 :
74 : protected:
75 : TestBasicBlockEntryHookTransform tx_;
76 : };
77 :
78 : void BasicBlockEntryHookTransformTest::CheckBasicBlockInstrumentation(
79 E : InstrumentationKind kind) {
80 : // Let's examine each eligible block to verify that its basic blocks have been
81 : // instrumented.
82 E : size_t num_decomposed_blocks = 0;
83 E : size_t total_basic_blocks = 0;
84 : BlockGraph::BlockMap::const_iterator block_iter =
85 E : block_graph_.blocks().begin();
86 E : for (; block_iter != block_graph_.blocks().end(); ++block_iter) {
87 E : const BlockGraph::Block& block = block_iter->second;
88 :
89 : // Skip non-code blocks.
90 E : if (block.type() != BlockGraph::CODE_BLOCK)
91 E : continue;
92 :
93 : // We'll skip thunks, they're a mixed bag of things.
94 E : if (block.section() == tx_.thunk_section_->id())
95 E : continue;
96 :
97 : // Blocks which are not bb-decomposable should be thunked. While there may
98 : // be some internal referrers, the only external referrers should be thunks.
99 E : if (!policy_->BlockIsSafeToBasicBlockDecompose(&block)) {
100 E : size_t num_external_thunks = 0;
101 : BlockGraph::Block::ReferrerSet::const_iterator ref_iter =
102 E : block.referrers().begin();
103 E : for (; ref_iter != block.referrers().end(); ++ref_iter) {
104 E : if (ref_iter->first != &block) {
105 E : ASSERT_EQ(tx_.thunk_section_->id(), ref_iter->first->section());
106 E : ++num_external_thunks;
107 : }
108 E : }
109 :
110 : // Each of the thunks for a non-decomposable block will reuse the same
111 : // id to source range map entry, so we increment total_basic_blocks once
112 : // if num_external_thunks is non-zero. Note that we cannot assert that
113 : // num_external_thunks > 0 because the block could be statically dead
114 : // (in a debug build, for example).
115 E : if (num_external_thunks != 0U)
116 E : ++total_basic_blocks;
117 E : continue;
118 : }
119 :
120 : // Note that we have attempted to validate a block.
121 E : ++num_decomposed_blocks;
122 :
123 : // Decompose the block to basic-blocks.
124 E : BasicBlockSubGraph subgraph;
125 E : BasicBlockDecomposer bb_decomposer(&block, &subgraph);
126 E : ASSERT_TRUE(bb_decomposer.Decompose());
127 :
128 : // Check if each non-padding basic code-block begins with the
129 : // instrumentation sequence.
130 E : size_t num_basic_blocks = 0;
131 : BasicBlockSubGraph::BBCollection::const_iterator bb_iter =
132 E : subgraph.basic_blocks().begin();
133 E : for (; bb_iter != subgraph.basic_blocks().end(); ++bb_iter) {
134 E : const BasicCodeBlock* bb = BasicCodeBlock::Cast(*bb_iter);
135 E : if (bb == NULL || bb->is_padding())
136 E : continue;
137 E : ++num_basic_blocks;
138 :
139 E : if (kind == kAgentInstrumentation) {
140 E : ASSERT_LE(3U, bb->instructions().size());
141 : BasicBlock::Instructions::const_iterator inst_iter =
142 E : bb->instructions().begin();
143 :
144 : // Instruction 1 should push the basic block id.
145 E : const Instruction& inst1 = *inst_iter;
146 E : EXPECT_EQ(I_PUSH, inst1.representation().opcode);
147 :
148 : // Instruction 2 should push the frequency data block pointer.
149 E : const Instruction& inst2 = *(++inst_iter);
150 E : EXPECT_EQ(I_PUSH, inst2.representation().opcode);
151 E : ASSERT_EQ(1U, inst2.references().size());
152 : EXPECT_EQ(tx_.frequency_data_block(),
153 E : inst2.references().begin()->second.block());
154 :
155 : // Instruction 3 should be a call to the bb entry hook.
156 E : const Instruction& inst3 = *(++inst_iter);
157 E : EXPECT_EQ(I_CALL, inst3.representation().opcode);
158 E : ASSERT_EQ(1U, inst3.references().size());
159 : EXPECT_EQ(tx_.bb_entry_hook_ref_.referenced(),
160 E : inst3.references().begin()->second.block());
161 E : } else {
162 i : DCHECK(kind == kFastPathInstrumentation);
163 i : ASSERT_LE(2U, bb->instructions().size());
164 : BasicBlock::Instructions::const_iterator inst_iter =
165 i : bb->instructions().begin();
166 :
167 : // Instruction 1 should push the basic block id.
168 i : const Instruction& inst1 = *inst_iter;
169 i : EXPECT_EQ(I_PUSH, inst1.representation().opcode);
170 :
171 : // Instruction 2 should be a call to the fast bb entry hook.
172 i : const Instruction& inst2 = *(++inst_iter);
173 i : EXPECT_EQ(I_CALL, inst2.representation().opcode);
174 i : ASSERT_EQ(1U, inst2.references().size());
175 : }
176 E : }
177 E : EXPECT_NE(0U, num_basic_blocks);
178 E : total_basic_blocks += num_basic_blocks;
179 E : }
180 :
181 E : EXPECT_NE(0U, num_decomposed_blocks);
182 E : EXPECT_EQ(total_basic_blocks, tx_.bb_ranges().size());
183 E : }
184 :
185 : } // namespace
186 :
187 E : TEST_F(BasicBlockEntryHookTransformTest, SetInlinePathFlag) {
188 E : EXPECT_FALSE(tx_.inline_fast_path());
189 E : tx_.set_inline_fast_path(true);
190 E : EXPECT_TRUE(tx_.inline_fast_path());
191 E : tx_.set_inline_fast_path(false);
192 E : EXPECT_FALSE(tx_.inline_fast_path());
193 E : }
194 :
195 E : TEST_F(BasicBlockEntryHookTransformTest, ApplyAgentInstrumentation) {
196 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
197 :
198 : // Apply the transform.
199 E : tx_.set_src_ranges_for_thunks(true);
200 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
201 E : &tx_, policy_, &block_graph_, header_block_));
202 E : ASSERT_TRUE(tx_.frequency_data_block() != NULL);
203 E : ASSERT_TRUE(tx_.thunk_section_ != NULL);
204 E : ASSERT_TRUE(tx_.bb_entry_hook_ref_.IsValid());
205 E : ASSERT_LT(0u, tx_.bb_ranges().size());
206 :
207 : // Validate the basic-block frequency data structure.
208 E : block_graph::ConstTypedBlock<IndexedFrequencyData> frequency_data;
209 E : ASSERT_TRUE(frequency_data.Init(0, tx_.frequency_data_block()));
210 E : EXPECT_EQ(kBasicBlockEntryAgentId, frequency_data->agent_id);
211 E : EXPECT_EQ(kBasicBlockFrequencyDataVersion, frequency_data->version);
212 E : EXPECT_EQ(IndexedFrequencyData::BASIC_BLOCK_ENTRY, frequency_data->data_type);
213 E : EXPECT_EQ(tx_.bb_ranges().size(), frequency_data->num_entries);
214 E : EXPECT_EQ(sizeof(uint32), frequency_data->frequency_size);
215 : EXPECT_TRUE(frequency_data.HasReferenceAt(
216 E : frequency_data.OffsetOf(frequency_data->frequency_data)));
217 : EXPECT_EQ(sizeof(BasicBlockIndexedFrequencyData),
218 E : tx_.frequency_data_block()->size());
219 : EXPECT_EQ(sizeof(BasicBlockIndexedFrequencyData),
220 E : tx_.frequency_data_block()->data_size());
221 : EXPECT_EQ(frequency_data->num_entries * frequency_data->frequency_size,
222 E : tx_.frequency_data_buffer_block()->size());
223 :
224 : // Validate that all basic block have been instrumented.
225 E : CheckBasicBlockInstrumentation(kAgentInstrumentation);
226 E : }
227 :
228 : } // namespace transforms
229 : } // namespace instrument
|