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 : // Unittests for iteration primitives.
16 :
17 : #include "syzygy/instrument/transforms/entry_call_transform.h"
18 :
19 : #include <vector>
20 :
21 : #include "gmock/gmock.h"
22 : #include "gtest/gtest.h"
23 : #include "syzygy/block_graph/basic_block_assembler.h"
24 : #include "syzygy/block_graph/basic_block_test_util.h"
25 : #include "syzygy/block_graph/transform.h"
26 : #include "syzygy/block_graph/unittest_util.h"
27 : #include "syzygy/common/defs.h"
28 : #include "syzygy/instrument/transforms/unittest_util.h"
29 : #include "syzygy/pe/pe_utils.h"
30 :
31 : namespace instrument {
32 : namespace transforms {
33 :
34 : namespace {
35 :
36 : class TestingEntryCallBasicBlockTransform
37 : : public EntryCallBasicBlockTransform {
38 : public:
39 : TestingEntryCallBasicBlockTransform(
40 : const BlockGraph::Reference& hook_reference, bool debug_friendly) :
41 E : EntryCallBasicBlockTransform(hook_reference, debug_friendly) {
42 E : }
43 :
44 : // Expose for testing.
45 : using EntryCallBasicBlockTransform::TransformBasicBlockSubGraph;
46 : };
47 :
48 :
49 : class EntryCallTransformTest : public testing::TestDllTransformTest {
50 : };
51 :
52 : class EntryCallBasicBlockTransformTest : public testing::BasicBlockTest {
53 : public:
54 E : virtual void SetUp() override {
55 : // Create a dummy IAT block for our reference.
56 E : dummy_iat_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK, 1024, "IAT");
57 : import_ref_ =
58 : BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
59 : sizeof(core::AbsoluteAddress),
60 : dummy_iat_,
61 : 103 * sizeof(core::AbsoluteAddress),
62 E : 0);
63 E : }
64 :
65 E : BlockGraph::Block* TransformAssemblyFunc(bool debug_friendly) {
66 E : EntryCallBasicBlockTransform tx(import_ref_, debug_friendly);
67 :
68 : // Apply the transform.
69 E : block_graph::BlockVector created_blocks;
70 E : EXPECT_TRUE(ApplyBasicBlockSubGraphTransform(
71 : &tx, &policy_, &block_graph_, assembly_func_, &created_blocks));
72 :
73 E : if (created_blocks.size() != 1)
74 i : return NULL;
75 :
76 E : return created_blocks[0];
77 E : }
78 :
79 : protected:
80 : BlockGraph::Block* dummy_iat_;
81 : BlockGraph::Reference import_ref_;
82 : };
83 :
84 : } // namespace
85 :
86 E : TEST_F(EntryCallBasicBlockTransformTest, AccessorsAndMutators) {
87 E : EntryCallBasicBlockTransform tx(import_ref_, false);
88 :
89 E : EXPECT_STREQ("EntryCallBasicBlockTransform", tx.name());
90 E : }
91 :
92 E : TEST_F(EntryCallBasicBlockTransformTest, ApplyTransform) {
93 E : ASSERT_NO_FATAL_FAILURE(InitBlockGraph());
94 E : ASSERT_NO_FATAL_FAILURE(InitBasicBlockSubGraph());
95 :
96 : // Transform and return a function.
97 E : BlockGraph::Block* created_block = TransformAssemblyFunc(false);
98 E : ASSERT_NE(static_cast<BlockGraph::Block*>(NULL), created_block);
99 :
100 E : ASSERT_LE(2U, created_block->size());
101 : // We expect an indirect call instruction at the start of the block.
102 : // That's opcode FF, followed by a 0x15 modrm byte.
103 E : ASSERT_EQ(0xFF, created_block->data()[0]);
104 E : ASSERT_EQ(0x15, created_block->data()[1]);
105 :
106 : // Retrieve and check the inserted reference.
107 E : BlockGraph::Reference ref;
108 E : ASSERT_TRUE(created_block->GetReference(2, &ref));
109 E : ASSERT_EQ(import_ref_.referenced(), ref.referenced());
110 E : ASSERT_EQ(import_ref_.offset(), ref.offset());
111 :
112 : // Check that there's no source range for the first byte of the new block.
113 E : EXPECT_TRUE(created_block->source_ranges().FindRangePair(0, 1) == NULL);
114 :
115 : // We expect that the data_ block refers to the head of the new block.
116 E : ASSERT_TRUE(data_->GetReference(0, &ref));
117 E : EXPECT_EQ(created_block, ref.referenced());
118 E : EXPECT_EQ(0, ref.offset());
119 E : }
120 :
121 E : TEST_F(EntryCallBasicBlockTransformTest, ApplyTransformDebugFriendly) {
122 E : ASSERT_NO_FATAL_FAILURE(InitBlockGraph());
123 E : ASSERT_NO_FATAL_FAILURE(InitBasicBlockSubGraph());
124 :
125 : // Transform and return a function.
126 E : BlockGraph::Block* created_block = TransformAssemblyFunc(true);
127 E : ASSERT_NE(static_cast<BlockGraph::Block*>(NULL), created_block);
128 :
129 : // Check that there's a source range for the first byte of the new block.
130 E : EXPECT_TRUE(created_block->source_ranges().FindRangePair(0, 1) != NULL);
131 E : }
132 :
133 E : TEST_F(EntryCallBasicBlockTransformTest, CorrectlyInstrumentsSelfRecursion) {
134 : using block_graph::BasicBlockAssembler;
135 : using block_graph::BasicBlockReference;
136 : using block_graph::BasicCodeBlock;
137 : using block_graph::Instruction;
138 : using block_graph::Immediate;
139 : using block_graph::Successor;
140 :
141 E : BasicCodeBlock* code_block = subgraph_.AddBasicCodeBlock("Foo()");
142 E : ASSERT_NE(static_cast<BasicCodeBlock*>(NULL), code_block);
143 E : code_block->set_offset(0);
144 :
145 : // Create the minimal self-recursive function, that also loops to itself.
146 : BasicBlockAssembler assm(code_block->instructions().begin(),
147 E : &code_block->instructions());
148 E : assm.call(Immediate(code_block));
149 : code_block->successors().push_back(
150 : Successor(Successor::kConditionTrue,
151 : BasicBlockReference(BlockGraph::PC_RELATIVE_REF,
152 : 4,
153 : code_block),
154 E : 5));
155 :
156 : BasicBlockSubGraph::BlockDescription* desc =
157 : subgraph_.AddBlockDescription("Foo()", "foo.obj",
158 E : BlockGraph::CODE_BLOCK, 1, 1, 0);
159 E : ASSERT_NE(static_cast<BasicBlockSubGraph::BlockDescription*>(NULL), desc);
160 E : desc->basic_block_order.push_back(code_block);
161 :
162 E : TestingEntryCallBasicBlockTransform tx(import_ref_, false);
163 :
164 : // Apply the transform.
165 : ASSERT_TRUE(
166 E : tx.TransformBasicBlockSubGraph(&policy_, &block_graph_, &subgraph_));
167 :
168 : // Get the entry hook block.
169 E : BasicBlock* entry_hook = desc->basic_block_order.front();
170 E : ASSERT_NE(code_block, entry_hook);
171 :
172 : // Now make sure the self-referential call instruction has been redirected,
173 : // while the self-referential successor has not.
174 E : ASSERT_EQ(1U, code_block->instructions().size());
175 E : const Instruction& call_inst = code_block->instructions().front();
176 E : ASSERT_EQ(1U, call_inst.references().size());
177 E : ASSERT_TRUE(call_inst.references().begin() != call_inst.references().end());
178 E : const BasicBlockReference& call_ref = call_inst.references().begin()->second;
179 E : EXPECT_EQ(entry_hook, call_ref.basic_block());
180 :
181 E : ASSERT_EQ(1U, code_block->successors().size());
182 E : Successor& succ = code_block->successors().front();
183 E : EXPECT_EQ(code_block, succ.reference().basic_block());
184 E : }
185 :
186 E : TEST_F(EntryCallTransformTest, AccessorsAndMutators) {
187 E : EntryCallTransform tx(false);
188 :
189 E : EXPECT_STREQ("EntryCallTransform", tx.name());
190 E : EXPECT_STREQ("profile_client.dll", tx.instrument_dll_name());
191 E : EXPECT_EQ(false, tx.debug_friendly());
192 :
193 E : tx.set_instrument_dll_name("HulaBonga.dll");
194 E : EXPECT_STREQ("HulaBonga.dll", tx.instrument_dll_name());
195 E : }
196 :
197 E : TEST_F(EntryCallTransformTest, TransformCreatesThunkSection) {
198 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
199 :
200 : using block_graph::BlockGraph;
201 : ASSERT_EQ(static_cast<BlockGraph::Section*>(NULL),
202 E : block_graph_.FindSection(common::kThunkSectionName));
203 :
204 E : EntryCallTransform transform(false);
205 :
206 : // Run the transform.
207 : ASSERT_TRUE(ApplyBlockGraphTransform(
208 E : &transform, policy_, &block_graph_, header_block_));
209 :
210 : // Check that the thunks section now exists.
211 : ASSERT_NE(static_cast<BlockGraph::Section*>(NULL),
212 E : block_graph_.FindSection(common::kThunkSectionName));
213 E : }
214 :
215 : } // namespace transforms
216 : } // namespace instrument
|