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 : // Unittests for BlockGraph transform wrapper.
16 :
17 : #include "syzygy/block_graph/transform.h"
18 :
19 : #include "gmock/gmock.h"
20 : #include "gtest/gtest.h"
21 : #include "syzygy/block_graph/unittest_util.h"
22 :
23 : namespace block_graph {
24 : namespace {
25 :
26 : using testing::_;
27 : using testing::DummyTransformPolicy;
28 : using testing::Invoke;
29 : using testing::Return;
30 :
31 : // A constant data structure/buffer from which to initialize a data block.
32 : // The structure contains a core reference (function pointer) and an integer
33 : // data element.
34 : struct MyData {
35 : void (*code)(int);
36 : int data;
37 : };
38 : const BlockGraph::Offset kOffsetOfReferenceToCode = offsetof(MyData, code);
39 : const BlockGraph::Offset kOffsetOfData = offsetof(MyData, data);
40 : const MyData kDataBytes = { reinterpret_cast<void(*)(int)>(0xCAFEBABE),
41 : 0xDEADBEEF };
42 :
43 : // A byte buffer from which to initialize a code block. The original C source
44 : // code for this function is:
45 : //
46 : // static int y = 1;
47 : // void add(int x) {
48 : // y += x;
49 : // }
50 : //
51 : // Note the reference to
52 : // y starts 5 bytes from the end.
53 : const uint8_t kCodeBytes[] = {
54 : 0x8B,
55 : 0x44,
56 : 0x24,
57 : 0x04, // mov eax,dword ptr [esp+4]
58 : 0x01,
59 : 0x05,
60 : 0x00,
61 : 0x00,
62 : 0x00,
63 : 0x00, // add dword ptr [_y],eax
64 : 0xC3 // ret
65 : };
66 : const BlockGraph::Offset kOffsetOfCode = 0;
67 : const BlockGraph::Offset kOffsetOfReferenceToData = sizeof(kCodeBytes) - 5;
68 :
69 : class ApplyBlockGraphTransformTest : public testing::Test {
70 : public:
71 E : virtual void SetUp() {
72 E : header_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK, 10, "Header");
73 E : }
74 :
75 : protected:
76 : DummyTransformPolicy policy_;
77 : BlockGraph block_graph_;
78 : BlockGraph::Block* header_block_;
79 : };
80 :
81 : class LenientMockBlockGraphTransform : public BlockGraphTransformInterface {
82 : public:
83 E : virtual ~LenientMockBlockGraphTransform() { }
84 :
85 E : virtual const char* name() const { return "MockBlockGraphTransform"; }
86 :
87 E : MOCK_METHOD3(TransformBlockGraph,
88 : bool(const TransformPolicyInterface*,
89 : BlockGraph*,
90 E : BlockGraph::Block*));
91 :
92 : bool DeleteHeader(const TransformPolicyInterface* policy,
93 : BlockGraph* block_graph,
94 E : BlockGraph::Block* header_block) {
95 E : CHECK(block_graph->RemoveBlock(header_block));
96 E : return true;
97 E : }
98 : };
99 : typedef testing::StrictMock<LenientMockBlockGraphTransform>
100 : MockBlockGraphTransform;
101 :
102 : class ApplyBasicBlockSubGraphTransformTest : public testing::Test {
103 : public:
104 E : ApplyBasicBlockSubGraphTransformTest()
105 E : : data_block_(NULL), code_block_(NULL) {
106 E : }
107 :
108 E : virtual void SetUp() {
109 : // Create some blocks to test with.
110 E : data_block_ = block_graph_.AddBlock(
111 : BlockGraph::DATA_BLOCK, sizeof(kDataBytes), "Data");
112 E : ASSERT_TRUE(data_block_ != NULL);
113 E : code_block_ = block_graph_.AddBlock(
114 : BlockGraph::CODE_BLOCK, sizeof(kCodeBytes), "Code");
115 E : ASSERT_TRUE(code_block_ != NULL);
116 :
117 : // Set up the data block.
118 E : data_block_->SetData(reinterpret_cast<const uint8_t*>(&kDataBytes),
119 : sizeof(kDataBytes));
120 :
121 : // Set up the code block.
122 E : ASSERT_TRUE(code_block_->SetLabel(
123 : kOffsetOfCode,
124 : BlockGraph::Label("Code", BlockGraph::CODE_LABEL)));
125 E : code_block_->SetData(kCodeBytes, sizeof(kCodeBytes));
126 :
127 : // Set up the references
128 E : ASSERT_TRUE(
129 : data_block_->SetReference(kOffsetOfReferenceToCode,
130 : MakeReference(code_block_, kOffsetOfCode)));
131 E : ASSERT_TRUE(
132 : code_block_->SetReference(kOffsetOfReferenceToData,
133 : MakeReference(data_block_, kOffsetOfData)));
134 E : }
135 :
136 : static BlockGraph::Reference MakeReference(BlockGraph::Block* target,
137 E : BlockGraph::Offset offset) {
138 E : EXPECT_TRUE(target != NULL);
139 E : return BlockGraph::Reference(BlockGraph::RELATIVE_REF,
140 : BlockGraph::Reference::kMaximumSize,
141 : target, offset, offset);
142 E : }
143 :
144 : protected:
145 : DummyTransformPolicy policy_;
146 : BlockGraph block_graph_;
147 : BlockGraph::Block* data_block_;
148 : BlockGraph::Block* code_block_;
149 : };
150 :
151 : class MockBasicBlockSubGraphTransform :
152 : public BasicBlockSubGraphTransformInterface {
153 : public:
154 E : virtual ~MockBasicBlockSubGraphTransform() { }
155 :
156 i : virtual const char* name() const { return "MockBasicBlockSubGraphTransform"; }
157 :
158 E : MOCK_METHOD3(TransformBasicBlockSubGraph,
159 : bool(const TransformPolicyInterface*,
160 : BlockGraph*,
161 E : BasicBlockSubGraph*));
162 : };
163 :
164 :
165 : class ApplyImageLayoutTransformTest : public testing::Test {
166 : public:
167 E : virtual void SetUp() {
168 E : BlockGraph block_graph_;
169 E : block_graph_.AddBlock(BlockGraph::DATA_BLOCK, 10, "Block1");
170 E : block_graph_.AddBlock(BlockGraph::DATA_BLOCK, 20, "Block2");
171 E : image_layout_ = new pe::ImageLayout(&block_graph_);
172 E : ordered_block_graph_ = new OrderedBlockGraph(&block_graph_);
173 E : }
174 :
175 : protected:
176 : DummyTransformPolicy policy_;
177 : pe::ImageLayout* image_layout_;
178 : OrderedBlockGraph* ordered_block_graph_;
179 : };
180 :
181 : class LenientMockImageLayoutTransform : public ImageLayoutTransformInterface {
182 : public:
183 E : virtual ~LenientMockImageLayoutTransform() { }
184 :
185 E : virtual const char* name() const { return "MockImageLayoutTransform"; }
186 :
187 E : MOCK_METHOD3(TransformImageLayout,
188 : bool(const TransformPolicyInterface*,
189 : const pe::ImageLayout*,
190 E : const OrderedBlockGraph*));
191 : };
192 : typedef testing::StrictMock<LenientMockImageLayoutTransform>
193 : MockImageLayoutTransform;
194 :
195 : } // namespace
196 :
197 E : TEST_F(ApplyBlockGraphTransformTest, NormalTransformSucceeds) {
198 E : MockBlockGraphTransform transform;
199 : EXPECT_CALL(transform, TransformBlockGraph(_, _, _)).Times(1).
200 E : WillOnce(Return(true));
201 E : EXPECT_TRUE(ApplyBlockGraphTransform(&transform,
202 : &policy_,
203 : &block_graph_,
204 E : header_block_));
205 E : }
206 :
207 E : TEST_F(ApplyBlockGraphTransformTest, DeletingHeaderFails) {
208 E : MockBlockGraphTransform transform;
209 E : EXPECT_CALL(transform, TransformBlockGraph(_, _, _)).Times(1).WillOnce(
210 : Invoke(&transform, &MockBlockGraphTransform::DeleteHeader));
211 E : EXPECT_FALSE(ApplyBlockGraphTransform(&transform,
212 : &policy_,
213 : &block_graph_,
214 E : header_block_));
215 E : }
216 :
217 E : TEST_F(ApplyBlockGraphTransformTest, VectorTransformSucceeds) {
218 E : MockBlockGraphTransform tx1, tx2, tx3;
219 E : std::vector<BlockGraphTransformInterface*> txs;
220 E : txs.push_back(&tx1);
221 E : txs.push_back(&tx2);
222 E : txs.push_back(&tx3);
223 :
224 : EXPECT_CALL(tx1, TransformBlockGraph(&policy_, &block_graph_, header_block_))
225 E : .WillOnce(Return(true));
226 : EXPECT_CALL(tx2, TransformBlockGraph(&policy_, &block_graph_, header_block_))
227 E : .WillOnce(Return(true));
228 : EXPECT_CALL(tx3, TransformBlockGraph(&policy_, &block_graph_, header_block_))
229 E : .WillOnce(Return(true));
230 :
231 E : EXPECT_TRUE(ApplyBlockGraphTransforms(
232 E : txs, &policy_, &block_graph_, header_block_));
233 E : }
234 :
235 E : TEST_F(ApplyBlockGraphTransformTest, VectorTransformFails) {
236 E : MockBlockGraphTransform tx1, tx2, tx3;
237 E : std::vector<BlockGraphTransformInterface*> txs;
238 E : txs.push_back(&tx1);
239 E : txs.push_back(&tx2);
240 E : txs.push_back(&tx3);
241 :
242 : EXPECT_CALL(tx1, TransformBlockGraph(&policy_, &block_graph_, header_block_))
243 E : .WillOnce(Return(true));
244 : EXPECT_CALL(tx2, TransformBlockGraph(&policy_, &block_graph_, header_block_))
245 E : .WillOnce(Return(false));
246 :
247 E : EXPECT_FALSE(ApplyBlockGraphTransforms(
248 E : txs, &policy_, &block_graph_, header_block_));
249 E : }
250 :
251 E : TEST_F(ApplyBasicBlockSubGraphTransformTest, TransformFails) {
252 : // Remember the block ids of the original blocks.
253 E : BlockGraph::BlockId data_block_id = data_block_->id();
254 E : BlockGraph::BlockId code_block_id = code_block_->id();
255 :
256 : // Apply an empty transform that reports failure.
257 E : MockBasicBlockSubGraphTransform transform;
258 : EXPECT_CALL(transform, TransformBasicBlockSubGraph(_, _, _)).Times(1).
259 E : WillOnce(Return(false));
260 E : EXPECT_FALSE(ApplyBasicBlockSubGraphTransform(&transform,
261 : &policy_,
262 : &block_graph_,
263 : code_block_,
264 E : NULL));
265 :
266 : // The original block graph should be unchanged.
267 E : EXPECT_EQ(2U, block_graph_.blocks().size());
268 E : EXPECT_EQ(data_block_, block_graph_.GetBlockById(data_block_id));
269 E : EXPECT_EQ(code_block_, block_graph_.GetBlockById(code_block_id));
270 E : }
271 :
272 E : TEST_F(ApplyBasicBlockSubGraphTransformTest, EmptyTransformSucceeds) {
273 : // Remember the block ids of the original blocks.
274 E : BlockGraph::BlockId data_block_id = data_block_->id();
275 E : BlockGraph::BlockId code_block_id = code_block_->id();
276 :
277 : // Apply an empty transform that reports success.
278 E : MockBasicBlockSubGraphTransform transform;
279 E : BlockVector new_blocks;
280 : EXPECT_CALL(transform, TransformBasicBlockSubGraph(_, _, _)).Times(1).
281 E : WillOnce(Return(true));
282 E : EXPECT_TRUE(ApplyBasicBlockSubGraphTransform(&transform,
283 : &policy_,
284 : &block_graph_,
285 : code_block_,
286 E : &new_blocks));
287 :
288 : // The code block should have been replaced with an equivalent one. We'll
289 : // have the same number of blocks, but the code block should no longer
290 : // be in the graph.
291 E : EXPECT_EQ(2U, block_graph_.blocks().size());
292 E : EXPECT_EQ(data_block_, block_graph_.GetBlockById(data_block_id));
293 E : EXPECT_EQ(NULL, block_graph_.GetBlockById(code_block_id));
294 :
295 : // Clean up our dangling pointer.
296 E : code_block_ = NULL;
297 :
298 : // Find the new block.
299 E : ASSERT_EQ(1U, new_blocks.size());
300 E : const BlockGraph::Block* new_block = new_blocks[0];
301 :
302 : // Validate the references.
303 E : EXPECT_EQ(1U, new_block->references().size());
304 E : BlockGraph::Reference ref;
305 E : EXPECT_TRUE(new_block->GetReference(kOffsetOfReferenceToData, &ref));
306 E : EXPECT_EQ(kOffsetOfData, ref.offset());
307 E : EXPECT_EQ(data_block_, ref.referenced());
308 :
309 : // Validate the referrers.
310 E : EXPECT_EQ(1U, new_block->referrers().size());
311 E : EXPECT_EQ(data_block_, new_block->referrers().begin()->first);
312 E : EXPECT_EQ(kOffsetOfReferenceToCode, new_block->referrers().begin()->second);
313 E : EXPECT_TRUE(new_block->referrers().begin()->first->GetReference(
314 E : kOffsetOfReferenceToCode, &ref));
315 E : EXPECT_EQ(new_block, ref.referenced());
316 E : }
317 :
318 E : TEST_F(ApplyBasicBlockSubGraphTransformTest, VectorTransformSucceeds) {
319 : // Validate applying a vector of transforms.
320 E : MockBasicBlockSubGraphTransform transform1;
321 E : MockBasicBlockSubGraphTransform transform2;
322 E : BlockVector new_blocks;
323 : EXPECT_CALL(transform1, TransformBasicBlockSubGraph(_, _, _)).Times(1).
324 E : WillOnce(Return(true));
325 : EXPECT_CALL(transform2, TransformBasicBlockSubGraph(_, _, _)).Times(1).
326 E : WillOnce(Return(true));
327 :
328 E : std::vector<BasicBlockSubGraphTransformInterface*> transforms;
329 E : transforms.push_back(&transform1);
330 E : transforms.push_back(&transform2);
331 E : EXPECT_TRUE(ApplyBasicBlockSubGraphTransforms(transforms,
332 : &policy_,
333 : &block_graph_,
334 : code_block_,
335 E : &new_blocks));
336 E : }
337 :
338 E : TEST_F(ApplyImageLayoutTransformTest, NormalTransformSucceeds) {
339 E : MockImageLayoutTransform transform;
340 : EXPECT_CALL(transform, TransformImageLayout(_, _, _)).Times(1).
341 E : WillOnce(Return(true));
342 E : EXPECT_TRUE(ApplyImageLayoutTransform(&transform,
343 : &policy_,
344 : image_layout_,
345 E : ordered_block_graph_));
346 E : }
347 :
348 E : TEST_F(ApplyImageLayoutTransformTest, VectorTransformSucceeds) {
349 E : MockImageLayoutTransform tx1, tx2, tx3;
350 E : std::vector<ImageLayoutTransformInterface*> txs;
351 E : txs.push_back(&tx1);
352 E : txs.push_back(&tx2);
353 E : txs.push_back(&tx3);
354 :
355 : EXPECT_CALL(tx1, TransformImageLayout(&policy_,
356 : image_layout_,
357 : ordered_block_graph_))
358 E : .WillOnce(Return(true));
359 : EXPECT_CALL(tx2, TransformImageLayout(&policy_,
360 : image_layout_,
361 : ordered_block_graph_))
362 E : .WillOnce(Return(true));
363 : EXPECT_CALL(tx3, TransformImageLayout(&policy_,
364 : image_layout_,
365 : ordered_block_graph_))
366 E : .WillOnce(Return(true));
367 :
368 E : EXPECT_TRUE(ApplyImageLayoutTransforms(
369 E : txs, &policy_, image_layout_, ordered_block_graph_));
370 E : }
371 : } // namespace block_graph
|