1 : // Copyright 2012 Google Inc.
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 : // Coverage instrumentation transform unittests.
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/transform.h"
25 : #include "syzygy/block_graph/typed_block.h"
26 : #include "syzygy/common/basic_block_frequency_data.h"
27 : #include "syzygy/core/unittest_util.h"
28 : #include "syzygy/pe/block_util.h"
29 : #include "syzygy/pe/decomposer.h"
30 : #include "syzygy/pe/unittest_util.h"
31 :
32 : #include "mnemonics.h" // NOLINT
33 :
34 : namespace instrument {
35 : namespace transforms {
36 : namespace {
37 :
38 : using block_graph::BasicBlock;
39 : using block_graph::BasicBlockDecomposer;
40 : using block_graph::BasicBlockSubGraph;
41 : using block_graph::BlockGraph;
42 : using block_graph::Instruction;
43 :
44 : class BasicBlockEntryHookTransformTest : public testing::PELibUnitTest {
45 : public:
46 E : BasicBlockEntryHookTransformTest() : dos_header_block_(NULL) { }
47 :
48 E : void DecomposeTestDll() {
49 : // Open the PE file.
50 E : ASSERT_TRUE(pe_file_.Init(::testing::GetOutputRelativePath(kDllName)));
51 :
52 : // Initialize the block-graph.
53 E : pe::ImageLayout layout(&block_graph_);
54 E : pe::Decomposer decomposer(pe_file_);
55 E : ASSERT_TRUE(decomposer.Decompose(&layout));
56 :
57 : // Get the DOS header block.
58 : dos_header_block_ = layout.blocks.GetBlockByAddress(
59 E : core::RelativeAddress(0));
60 E : ASSERT_TRUE(dos_header_block_ != NULL);
61 E : }
62 :
63 : pe::PEFile pe_file_;
64 : BlockGraph block_graph_;
65 : BlockGraph::Block* dos_header_block_;
66 : };
67 :
68 : } // namespace
69 :
70 E : TEST_F(BasicBlockEntryHookTransformTest, DefaultConstructor) {
71 E : BasicBlockEntryHookTransform tx;
72 :
73 E : std::string module_name(BasicBlockEntryHookTransform::kDefaultModuleName);
74 E : std::string function_name(BasicBlockEntryHookTransform::kDefaultFunctionName);
75 E : EXPECT_EQ(module_name, tx.module_name());
76 E : EXPECT_EQ(function_name, tx.function_name());
77 E : EXPECT_EQ(0U, tx.bb_addresses().size());
78 E : EXPECT_FALSE(tx.bb_entry_hook_ref().IsValid());
79 E : }
80 :
81 E : TEST_F(BasicBlockEntryHookTransformTest, NamedConstructor) {
82 E : std::string module_name("foo.dll");
83 E : std::string function_name("bar");
84 :
85 E : BasicBlockEntryHookTransform tx(module_name, function_name);
86 E : EXPECT_EQ(module_name, tx.module_name());
87 E : EXPECT_EQ(function_name, tx.function_name());
88 E : EXPECT_EQ(0U, tx.bb_addresses().size());
89 E : EXPECT_FALSE(tx.bb_entry_hook_ref().IsValid());
90 E : }
91 :
92 E : TEST_F(BasicBlockEntryHookTransformTest, SetNames) {
93 E : std::string module_name("foo.dll");
94 E : std::string function_name("bar");
95 :
96 E : BasicBlockEntryHookTransform tx;
97 E : tx.set_module_name(module_name);
98 E : tx.set_function_name(function_name);
99 E : EXPECT_EQ(module_name, tx.module_name());
100 E : EXPECT_EQ(function_name, tx.function_name());
101 E : }
102 :
103 E : TEST_F(BasicBlockEntryHookTransformTest, ApplyBasicBlockEntryHookTransform) {
104 E : DecomposeTestDll();
105 E : BasicBlockEntryHookTransform tx;
106 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(&tx, &block_graph_,
107 E : dos_header_block_));
108 :
109 : // Let's examine each eligible block to verify that its BB's have been
110 : // instrumented.
111 E : size_t num_decomposed_blocks = 0;
112 E : size_t total_basic_blocks = 0;
113 : BlockGraph::BlockMap::const_iterator block_iter =
114 E : block_graph_.blocks().begin();
115 E : for (; block_iter != block_graph_.blocks().end(); ++block_iter) {
116 E : const BlockGraph::Block& block = block_iter->second;
117 :
118 : // Skip ineligible blocks.
119 E : if (block.type() != BlockGraph::CODE_BLOCK)
120 E : continue;
121 E : if (!pe::CodeBlockIsBasicBlockDecomposable(&block))
122 E : continue;
123 :
124 : // Note that we have attempted to validate a block.
125 E : ++num_decomposed_blocks;
126 :
127 : // Decompose the block to basic-blocks.
128 E : BasicBlockSubGraph subgraph;
129 E : BasicBlockDecomposer bb_decomposer(&block, &subgraph);
130 E : ASSERT_TRUE(bb_decomposer.Decompose());
131 :
132 : // Check if each basic block begins with the instrumentation sequence.
133 E : size_t num_basic_blocks = 0;
134 : BasicBlockSubGraph::BBCollection::const_iterator bb_iter =
135 E : subgraph.basic_blocks().begin();
136 E : for (; bb_iter != subgraph.basic_blocks().end(); ++bb_iter) {
137 E : const BasicBlock& bb = bb_iter->second;
138 E : if (bb.type() != BasicBlock::BASIC_CODE_BLOCK)
139 E : continue;
140 E : ++num_basic_blocks;
141 E : ASSERT_LE(2U, bb.instructions().size());
142 E : const Instruction& inst1 = *(bb.instructions().begin());
143 E : const Instruction& inst2 = *(++bb.instructions().begin());
144 E : EXPECT_EQ(I_PUSH, inst1.representation().opcode);
145 E : EXPECT_EQ(I_CALL, inst2.representation().opcode);
146 E : EXPECT_EQ(1U, inst2.references().size());
147 : EXPECT_EQ(tx.bb_entry_hook_ref().referenced(),
148 E : inst2.references().begin()->second.block());
149 E : }
150 E : EXPECT_NE(0U, num_basic_blocks);
151 E : total_basic_blocks += num_basic_blocks;
152 E : }
153 :
154 E : EXPECT_NE(0U, num_decomposed_blocks);
155 E : EXPECT_EQ(total_basic_blocks, tx.bb_addresses().size());
156 E : }
157 :
158 : } // namespace transforms
159 : } // namespace instrument
|