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/basic_block_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::thunk_section_;
51 :
52 E : BlockGraph::Block* frequency_data_block() {
53 E : return add_frequency_data_.frequency_data_block();
54 E : }
55 :
56 E : BlockGraph::Block* frequency_data_buffer_block() {
57 E : return add_frequency_data_.frequency_data_buffer_block();
58 E : }
59 : };
60 :
61 : typedef testing::TestDllTransformTest BasicBlockEntryHookTransformTest;
62 :
63 : } // namespace
64 :
65 E : TEST_F(BasicBlockEntryHookTransformTest, Apply) {
66 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
67 :
68 : // Apply the transform.
69 E : TestBasicBlockEntryHookTransform tx;
70 E : tx.set_src_ranges_for_thunks(true);
71 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(&tx, &block_graph_,
72 E : dos_header_block_));
73 E : ASSERT_TRUE(tx.frequency_data_block() != NULL);
74 E : ASSERT_TRUE(tx.thunk_section_ != NULL);
75 E : ASSERT_TRUE(tx.bb_entry_hook_ref_.IsValid());
76 E : ASSERT_LT(0u, tx.bb_ranges().size());
77 :
78 : // Validate the basic-block frequency data structure.
79 E : block_graph::ConstTypedBlock<IndexedFrequencyData> frequency_data;
80 E : ASSERT_TRUE(frequency_data.Init(0, tx.frequency_data_block()));
81 E : EXPECT_EQ(kBasicBlockEntryAgentId, frequency_data->agent_id);
82 E : EXPECT_EQ(kBasicBlockFrequencyDataVersion, frequency_data->version);
83 E : EXPECT_EQ(tx.bb_ranges().size(), frequency_data->num_entries);
84 E : EXPECT_EQ(sizeof(uint32), frequency_data->frequency_size);
85 : EXPECT_TRUE(frequency_data.HasReferenceAt(
86 E : frequency_data.OffsetOf(frequency_data->frequency_data)));
87 E : EXPECT_EQ(sizeof(IndexedFrequencyData), tx.frequency_data_block()->size());
88 : EXPECT_EQ(sizeof(IndexedFrequencyData),
89 E : tx.frequency_data_block()->data_size());
90 : EXPECT_EQ(frequency_data->num_entries * frequency_data->frequency_size,
91 E : tx.frequency_data_buffer_block()->size());
92 :
93 : // Let's examine each eligible block to verify that its BB's have been
94 : // instrumented.
95 E : size_t num_decomposed_blocks = 0;
96 E : size_t total_basic_blocks = 0;
97 : BlockGraph::BlockMap::const_iterator block_iter =
98 E : block_graph_.blocks().begin();
99 E : for (; block_iter != block_graph_.blocks().end(); ++block_iter) {
100 E : const BlockGraph::Block& block = block_iter->second;
101 :
102 : // Skip non-code blocks.
103 E : if (block.type() != BlockGraph::CODE_BLOCK)
104 E : continue;
105 :
106 : // We'll skip thunks, they're a mixed bag of things.
107 E : if (block.section() == tx.thunk_section_->id())
108 E : continue;
109 :
110 : // Blocks which are not bb-decomposable should be thunked. While there may
111 : // be some internal referrers, the only external referrers should be thunks.
112 E : if (!pe::CodeBlockIsBasicBlockDecomposable(&block)) {
113 E : size_t num_external_thunks = 0;
114 : BlockGraph::Block::ReferrerSet::const_iterator ref_iter =
115 E : block.referrers().begin();
116 E : for (; ref_iter != block.referrers().end(); ++ref_iter) {
117 E : if (ref_iter->first != &block) {
118 E : ASSERT_EQ(tx.thunk_section_->id(), ref_iter->first->section());
119 E : ++num_external_thunks;
120 : }
121 E : }
122 :
123 : // Each of the thunks for a non-decomposable block will reuse the same
124 : // id to source range map entry, so we increment total_basic_blocks once
125 : // if num_external_thunks is non-zero. Note that we cannot assert that
126 : // num_external_thunks > 0 because the block could be statically dead
127 : // (in a debug build, for example).
128 E : if (num_external_thunks != 0U)
129 E : ++total_basic_blocks;
130 E : continue;
131 : }
132 :
133 : // Note that we have attempted to validate a block.
134 E : ++num_decomposed_blocks;
135 :
136 : // Decompose the block to basic-blocks.
137 E : BasicBlockSubGraph subgraph;
138 E : BasicBlockDecomposer bb_decomposer(&block, &subgraph);
139 E : ASSERT_TRUE(bb_decomposer.Decompose());
140 :
141 : // Check if each non-padding basic code-block begins with the
142 : // instrumentation sequence.
143 E : size_t num_basic_blocks = 0;
144 : BasicBlockSubGraph::BBCollection::const_iterator bb_iter =
145 E : subgraph.basic_blocks().begin();
146 E : for (; bb_iter != subgraph.basic_blocks().end(); ++bb_iter) {
147 E : const BasicCodeBlock* bb = BasicCodeBlock::Cast(*bb_iter);
148 E : if (bb == NULL || bb->is_padding())
149 E : continue;
150 E : ++num_basic_blocks;
151 E : ASSERT_LE(3U, bb->instructions().size());
152 : BasicBlock::Instructions::const_iterator inst_iter =
153 E : bb->instructions().begin();
154 :
155 : // Instruction 1 should push the basic block id.
156 E : const Instruction& inst1 = *inst_iter;
157 E : EXPECT_EQ(I_PUSH, inst1.representation().opcode);
158 :
159 : // Instruction 2 should push the frequency data block pointer.
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 :
166 : // Instruction 3 should be a call to the bb entry 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.bb_entry_hook_ref_.referenced(),
171 E : inst3.references().begin()->second.block());
172 E : }
173 E : EXPECT_NE(0U, num_basic_blocks);
174 E : total_basic_blocks += num_basic_blocks;
175 E : }
176 :
177 E : EXPECT_NE(0U, num_decomposed_blocks);
178 E : EXPECT_EQ(total_basic_blocks, tx.bb_ranges().size());
179 E : }
180 :
181 : } // namespace transforms
182 : } // namespace instrument
|