1 : // Copyright 2015 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 : #include "syzygy/pe/hot_patching_writer.h"
16 :
17 : #include <memory>
18 :
19 : #include "gtest/gtest.h"
20 : #include "syzygy/block_graph/basic_block_assembler.h"
21 : #include "syzygy/block_graph/basic_block_subgraph.h"
22 : #include "syzygy/block_graph/block_builder.h"
23 : #include "syzygy/pe/hot_patching_decomposer.h"
24 : #include "syzygy/pe/hot_patching_unittest_util.h"
25 :
26 : namespace pe {
27 :
28 : namespace {
29 :
30 : using block_graph::BlockGraph;
31 : using block_graph::BlockBuilder;
32 : using block_graph::BasicCodeBlock;
33 : using block_graph::BasicBlockSubGraph;
34 : using block_graph::BasicBlockAssembler;
35 : using block_graph::Displacement;
36 : using block_graph::Immediate;
37 : using block_graph::Operand;
38 :
39 : const size_t kTestMemorySize = 1024U * 1024U;
40 :
41 : // TODO(cseri): This is based on EntryThunkTransform::CreateOneThunk, where it
42 : // has a comment that is should be made reusable. This class should be
43 : // renamed and moved into a common location.
44 : class TestBlockCreator {
45 : public:
46 : // Set up a basic block subgraph containing a single block description, with
47 : // that block description containing a single empty basic block, and get an
48 : // assembler writing into that basic block.
49 E : TestBlockCreator() {
50 : BasicBlockSubGraph::BlockDescription* block_desc =
51 : bbsg_.AddBlockDescription("foo",
52 : NULL,
53 : BlockGraph::CODE_BLOCK,
54 : 1,
55 : 1,
56 E : 0);
57 E : BasicCodeBlock* bb = bbsg_.AddBasicCodeBlock("foo");
58 E : block_desc->basic_block_order.push_back(bb);
59 : assm_.reset(new BasicBlockAssembler(bb->instructions().begin(),
60 E : &bb->instructions()));
61 E : }
62 :
63 E : BasicBlockAssembler* assm() {
64 E : return assm_.get();
65 E : }
66 :
67 : // Builds a block from the instructions in the assembler.
68 : // @param block_graph A block graph in which the new block should be inserted.
69 : // @param new_block will contain the newly created block.
70 E : void ToBlock(BlockGraph* block_graph, BlockGraph::Block** new_block) {
71 E : BlockBuilder block_builder(block_graph);
72 E : if (!block_builder.Merge(&bbsg_)) {
73 i : LOG(ERROR) << "Failed to build test block.";
74 i : *new_block = nullptr;
75 : }
76 :
77 : // Exactly one new block should have been created.
78 E : ASSERT_EQ(1U, block_builder.new_blocks().size());
79 E : *new_block = block_builder.new_blocks()[0];
80 E : }
81 :
82 : private:
83 : BasicBlockSubGraph bbsg_;
84 : std::unique_ptr<BasicBlockAssembler> assm_;
85 : };
86 :
87 : // Creates a simple block with a return instruction.
88 : // @param return_value The return value of the function in the generated block.
89 : // @param block_graph A block graph in which the new block should be inserted.
90 : // @param new_block will contain the newly created block.
91 : void CreateSimpleTestBlock(int return_value,
92 : BlockGraph* block_graph,
93 E : BlockGraph::Block** new_block) {
94 :
95 E : TestBlockCreator block_creator;
96 :
97 : // The goal is to test with a function that returns return_value.
98 : // Set up our function:
99 : // 1. MOV EAX, [imm32: return_value]
100 : // 2. RET
101 :
102 E : block_creator.assm()->mov(assm::eax, Immediate(return_value));
103 E : block_creator.assm()->ret();
104 :
105 E : ASSERT_NO_FATAL_FAILURE(block_creator.ToBlock(block_graph, new_block));
106 E : ASSERT_NE(nullptr, new_block);
107 E : }
108 :
109 : // Creates a block that, when executed, calls another block using a PC-relative
110 : // reference.
111 : // @param block_to_call The block to be called.
112 : // @param block_graph A block graph in which the new block should be inserted.
113 : // @param new_block will contain the newly created block.
114 : void CreateTestBlockWithPCRelativeReference(BlockGraph::Block* block_to_call,
115 : BlockGraph* block_graph,
116 E : BlockGraph::Block** new_block) {
117 :
118 E : TestBlockCreator block_creator;
119 :
120 : // The goal is to test with a function that calls |block_to_call| both via
121 : // PC-relative reference.
122 : //
123 : // The assembly code for the block:
124 : // 1. MOV EAX, 0
125 : // 2. CALL block_to_call // PC-relative reference
126 : // 3. ADD EAX, 1
127 : // 4. RET
128 :
129 : // Reset EAX to 1.
130 E : block_creator.assm()->mov(assm::eax, Immediate(1));
131 : // Use a call instruction to get a PC-relative reference.
132 E : block_creator.assm()->call(Immediate(block_to_call, 0));
133 E : block_creator.assm()->add(assm::eax, Immediate(1));
134 E : block_creator.assm()->ret();
135 :
136 E : block_creator.ToBlock(block_graph, new_block);
137 E : ASSERT_NE(nullptr, new_block);
138 E : }
139 :
140 : // Creates a block that, when executed, returns the address of another block
141 : // using an absolute reference.
142 : // @param referenced_block The block that's address should be returned.
143 : // @param block_graph A block graph in which the new block should be inserted.
144 : // @param new_block will contain the newly created block.
145 : void CreateTestBlockWithAbsoluteReference(BlockGraph::Block* referenced_block,
146 : BlockGraph* block_graph,
147 E : BlockGraph::Block** new_block) {
148 :
149 E : TestBlockCreator block_creator;
150 :
151 : // This test function returns the address of the block in |block_to_call|.
152 : //
153 : // The assembly code for the block:
154 : // 1. MOV EAX, block_to_call // absolute reference
155 : // 2. RET
156 :
157 E : block_creator.assm()->mov(assm::eax, Immediate(referenced_block, 0));
158 E : block_creator.assm()->ret();
159 :
160 E : block_creator.ToBlock(block_graph, new_block);
161 E : ASSERT_NE(nullptr, new_block);
162 E : }
163 :
164 : // Using this function pointer type we can call our test functions.
165 : typedef int __stdcall TestFunctionType();
166 :
167 : class HotPatchingWriterTest : public testing::Test {
168 : public:
169 : HotPatchingWriterTest() : simple_block_(nullptr),
170 E : simple_proc_(nullptr) {}
171 :
172 : // Creates a simple block and writes it using the member writer. Updates the
173 : // |simple_block_| and |simple_proc_| members.
174 : // NOTE: |simple_proc_| is nullptr after the call if the write did not
175 : // succeed. This allows testing failure scenarios.
176 E : void CreateAndWriteSimpleBlock() {
177 : // Test simple block.
178 E : CreateSimpleTestBlock(4, &block_graph_, &simple_block_);
179 E : ASSERT_NE(nullptr, simple_block_);
180 :
181 : // Write the block into memory.
182 : simple_proc_ = reinterpret_cast<TestFunctionType*>(
183 E : writer_.Write(simple_block_));
184 E : }
185 :
186 : protected:
187 : // The created test blocks will be inserted into this block graph.
188 : BlockGraph block_graph_;
189 :
190 : // The block for the simple block is saved as a member so that we set
191 : // up references to it.
192 : BlockGraph::Block* simple_block_;
193 :
194 : // The pointer for the simple procedure after written by the writer.
195 : TestFunctionType* simple_proc_;
196 :
197 : // The hot patching writer used by the tests.
198 : HotPatchingWriter writer_;
199 : };
200 :
201 : } // namespace
202 :
203 E : TEST_F(HotPatchingWriterTest, SimpleBlock) {
204 : // Initialize writer with buffer that has a sufficient size.
205 i : ASSERT_TRUE(writer_.Init(kTestMemorySize));
206 :
207 : // Create and write a simple block that we will call.
208 i : ASSERT_NO_FATAL_FAILURE(CreateAndWriteSimpleBlock());
209 i : ASSERT_NE(nullptr, simple_proc_);
210 :
211 : // Call the block and test the result. Zero EAX before calling to be sure
212 : // it does not contain the right result beforehand.
213 i : __asm xor eax, eax;
214 i : int test1 = simple_proc_();
215 i : ASSERT_EQ(4, test1);
216 i : }
217 :
218 : // Test writing a block that has a PC-relative reference.
219 E : TEST_F(HotPatchingWriterTest, PCRelativeReference) {
220 : // Initialize writer with buffer that has a sufficient size.
221 E : ASSERT_TRUE(writer_.Init(kTestMemorySize));
222 :
223 : // Create and write a simple block that we can reference.
224 E : ASSERT_NO_FATAL_FAILURE(CreateAndWriteSimpleBlock());
225 E : ASSERT_NE(nullptr, simple_proc_);
226 :
227 : // Create a block with a PC-relative call.
228 E : BlockGraph::Block* block = nullptr;
229 : ASSERT_NO_FATAL_FAILURE(CreateTestBlockWithPCRelativeReference(
230 E : simple_block_, &block_graph_, &block));
231 E : ASSERT_NE(nullptr, block);
232 :
233 : // Write the block to executable memory.
234 : TestFunctionType* test_proc =
235 E : reinterpret_cast<TestFunctionType*>(writer_.Write(block));
236 E : ASSERT_NE(nullptr, test_proc);
237 :
238 : // Call the block and test the result.
239 E : int test_result = test_proc();
240 E : ASSERT_EQ(5, test_result);
241 E : }
242 :
243 : // Test writing a block that has an absolute reference.
244 E : TEST_F(HotPatchingWriterTest, AbsoluteReference) {
245 : // Initialize writer with buffer that has a sufficient size.
246 E : ASSERT_TRUE(writer_.Init(kTestMemorySize));
247 :
248 : // Create and write a simple block that we can reference.
249 E : ASSERT_NO_FATAL_FAILURE(CreateAndWriteSimpleBlock());
250 E : ASSERT_NE(nullptr, simple_proc_);
251 :
252 : // Create a block with an absolute reference.
253 E : BlockGraph::Block* block = nullptr;
254 : ASSERT_NO_FATAL_FAILURE(CreateTestBlockWithAbsoluteReference(
255 E : simple_block_, &block_graph_, &block));
256 E : ASSERT_NE(nullptr, block);
257 :
258 : // Write the block to executable memory.
259 : TestFunctionType* test_proc =
260 E : reinterpret_cast<TestFunctionType*>(writer_.Write(block));
261 E : ASSERT_NE(nullptr, test_proc);
262 :
263 : // Call the block and test the result. The expected result is the function
264 : // pointer of the simple block.
265 E : int test_result = test_proc();
266 E : ASSERT_EQ(reinterpret_cast<int>(simple_proc_), test_result);
267 E : }
268 :
269 E : TEST_F(HotPatchingWriterTest, WriteFailsIfNotEnoughSpace) {
270 : // Initialize the writer with a buffer that's not big enough to hold the
271 : // simple test block.
272 E : ASSERT_TRUE(writer_.Init(3U));
273 :
274 : // Writing the block into memory should fail.
275 E : ASSERT_NO_FATAL_FAILURE(CreateAndWriteSimpleBlock());
276 E : ASSERT_EQ(nullptr, simple_proc_);
277 E : }
278 :
279 : namespace {
280 :
281 : // A basic block transform that does not change the basic block subgraph.
282 : class IdentityBasicBlockTransform
283 : : public block_graph::transforms::NamedBasicBlockSubGraphTransformImpl<
284 : IdentityBasicBlockTransform> {
285 : public:
286 : typedef block_graph::BlockGraph BlockGraph;
287 : typedef block_graph::BasicBlockSubGraph BasicBlockSubGraph;
288 : typedef block_graph::TransformPolicyInterface TransformPolicyInterface;
289 :
290 E : IdentityBasicBlockTransform() { }
291 :
292 : // @name BasicBlockSubGraphTransformInterface method.
293 : virtual bool TransformBasicBlockSubGraph(
294 : const TransformPolicyInterface* policy,
295 : BlockGraph* block_graph,
296 E : BasicBlockSubGraph* basic_block_subgraph) override {
297 E : return true;
298 E : }
299 :
300 : static const char kTransformName[];
301 : };
302 :
303 : const char IdentityBasicBlockTransform::kTransformName[] =
304 : "IdentityBasicBlockTransform";
305 :
306 : class HotPatchingWriterTestDllTest : public testing::HotPatchingTestDllTest {
307 : };
308 :
309 : } // namespace
310 :
311 E : TEST_F(HotPatchingWriterTestDllTest, Write) {
312 E : ASSERT_NO_FATAL_FAILURE(HotPatchInstrumentTestDll());
313 :
314 : // Load hot patchable library into memory.
315 E : testing::ScopedHMODULE module;
316 E : LoadTestDll(hp_test_dll_path_, &module);
317 :
318 : // Decompose the hot patchable library.
319 E : BlockGraph block_graph;
320 E : pe::ImageLayout layout(&block_graph);
321 E : HotPatchingDecomposer decomposer(module);
322 E : decomposer.Decompose(&layout);
323 :
324 E : pe::HotPatchingWriter writer;
325 E : ASSERT_TRUE(writer.Init(kTestMemorySize));
326 :
327 : // The block map changes during the basic block transform, so save the list of
328 : // blocks to transform first.
329 E : std::vector<BlockGraph::Block*> blocks_to_transform;
330 E : for (auto& entry : block_graph.blocks_mutable()) {
331 E : BlockGraph::Block* block = &entry.second;
332 : if (block->type() == BlockGraph::CODE_BLOCK &&
333 E : !(block->attributes() & BlockGraph::BUILT_BY_UNSUPPORTED_COMPILER)) {
334 E : blocks_to_transform.push_back(block);
335 : }
336 E : }
337 :
338 E : pe::PETransformPolicy pe_policy;
339 :
340 E : bool dllmain_found = false;
341 : ASSERT_EQ(blocks_to_transform.size(),
342 E : hp_transform_.blocks_prepared().size());
343 :
344 : // NOTE: This test assumes that the blocks IDs are the same order as in the
345 : // blocks themselves in the hot patching metadata.
346 E : for (size_t i = 0; i < blocks_to_transform.size(); ++i) {
347 E : BlockGraph::Block* original_block = hp_transform_.blocks_prepared()[i];
348 E : BlockGraph::Block* block = blocks_to_transform[i];
349 E : EXPECT_EQ(block->addr(), original_block->addr());
350 :
351 : // Write the transformed block of DllMain and call the written function.
352 : // There is no sense testing the other functions as we can't call them
353 : // without knowing their calling conventions.
354 E : if (original_block->name() == "DllMain") {
355 E : dllmain_found = true;
356 E : std::vector<BlockGraph::Block*> new_blocks;
357 E : IdentityBasicBlockTransform transform;
358 :
359 E : const void* old_entry_point = block->data();
360 :
361 E : ASSERT_TRUE(pe_policy.BlockIsSafeToBasicBlockDecompose(block));
362 :
363 : // Do a basic block decomposition first, that should ruin the references
364 : // in the memory.
365 : ASSERT_TRUE(ApplyBasicBlockSubGraphTransform(&transform,
366 : &pe_policy,
367 : &block_graph,
368 : block,
369 E : &new_blocks));
370 :
371 E : ASSERT_EQ(1U, new_blocks.size());
372 E : BlockGraph::Block* transformed_block = new_blocks.front();
373 :
374 : HotPatchingWriter::FunctionPointer new_entry_point =
375 E : writer.Write(transformed_block);
376 E : ASSERT_NE(nullptr, new_entry_point);
377 E : ASSERT_NE(old_entry_point, new_entry_point);
378 :
379 : // Call the DllMain.
380 : typedef BOOL WINAPI DllMainProc(
381 : _In_ HINSTANCE hinstDLL,
382 : _In_ DWORD fdwReason,
383 : _In_ LPVOID lpvReserved
384 : );
385 : reinterpret_cast<DllMainProc*>(new_entry_point)(
386 E : nullptr, DLL_PROCESS_ATTACH, nullptr);
387 E : }
388 E : }
389 E : ASSERT_TRUE(dllmain_found);
390 E : }
391 :
392 : } // namespace pe
|