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 : #include "syzygy/pe/coff_image_layout_builder.h"
16 :
17 : #include <cstring>
18 :
19 : #include "base/command_line.h"
20 : #include "base/process_util.h"
21 : #include "gmock/gmock.h"
22 : #include "gtest/gtest.h"
23 : #include "syzygy/block_graph/typed_block.h"
24 : #include "syzygy/block_graph/orderers/original_orderer.h"
25 : #include "syzygy/common/align.h"
26 : #include "syzygy/core/random_number_generator.h"
27 : #include "syzygy/core/unittest_util.h"
28 : #include "syzygy/pe/coff_decomposer.h"
29 : #include "syzygy/pe/coff_file.h"
30 : #include "syzygy/pe/coff_file_writer.h"
31 : #include "syzygy/pe/new_decomposer.h"
32 : #include "syzygy/pe/pe_utils.h"
33 : #include "syzygy/pe/unittest_util.h"
34 : #include "syzygy/testing/toolchain.h"
35 :
36 : namespace pe {
37 : namespace {
38 :
39 : using block_graph::BlockGraph;
40 : using block_graph::ConstTypedBlock;
41 : using block_graph::OrderedBlockGraph;
42 : using core::RelativeAddress;
43 :
44 : class ShuffleOrderer : public block_graph::BlockGraphOrdererInterface {
45 : public:
46 E : explicit ShuffleOrderer(uint32 seed) : rng_(seed) {
47 E : }
48 :
49 i : virtual const char* name() const OVERRIDE {
50 i : return "ShuffleOrderer";
51 i : }
52 :
53 : // Shuffle sections while paying attention to put .debug$S sections at the
54 : // end, so that they come after the associated sections.
55 : virtual bool OrderBlockGraph(
56 : OrderedBlockGraph* ordered_graph,
57 E : BlockGraph::Block* /* headers_block */) OVERRIDE {
58 E : DCHECK(ordered_graph != NULL);
59 E : BlockGraph* graph = ordered_graph->block_graph();
60 E : DCHECK(graph != NULL);
61 :
62 E : std::vector<BlockGraph::SectionId> sections;
63 E : for (size_t i = 0; i < graph->sections().size(); ++i)
64 E : sections.push_back(i);
65 :
66 E : std::random_shuffle(sections.begin(), sections.end(), rng_);
67 E : for (size_t i = 0; i < sections.size(); ++i) {
68 E : BlockGraph::Section* section = graph->GetSectionById(sections[i]);
69 E : if (section->name() == ".debug$S")
70 E : ordered_graph->PlaceAtTail(section);
71 E : else
72 E : ordered_graph->PlaceAtHead(section);
73 E : }
74 :
75 E : return true;
76 E : }
77 :
78 : private:
79 : core::RandomNumberGenerator rng_;
80 : };
81 :
82 : // We can't rely on CommandLine to built the command-line string for us
83 : // because it doesn't maintain the order of arguments.
84 : void MakeCommandLineString(const CommandLine::StringVector& args,
85 E : CommandLine::StringType* cmd_line) {
86 E : DCHECK(cmd_line != NULL);
87 :
88 E : cmd_line->clear();
89 E : for (size_t i = 0; i < args.size(); ++i) {
90 E : if (i > 0)
91 E : cmd_line->push_back(L' ');
92 E : cmd_line->push_back(L'"');
93 E : cmd_line->append(args[i]);
94 E : cmd_line->push_back(L'"');
95 E : }
96 E : }
97 :
98 : class CoffImageLayoutBuilderTest : public testing::PELibUnitTest {
99 : public:
100 E : CoffImageLayoutBuilderTest() : image_layout_(&block_graph_) {
101 E : }
102 :
103 E : virtual void SetUp() OVERRIDE {
104 E : testing::PELibUnitTest::SetUp();
105 :
106 : test_dll_obj_path_ =
107 E : testing::GetExeTestDataRelativePath(testing::kTestDllCoffObjName);
108 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir_path_));
109 E : new_test_dll_obj_path_ = temp_dir_path_.Append(L"test_dll.obj");
110 E : new_test_dll_path_ = temp_dir_path_.Append(testing::kTestDllName);
111 E : }
112 :
113 : protected:
114 : // Decompose test_dll.coff_obj.
115 E : void DecomposeOriginal() {
116 E : ASSERT_TRUE(image_file_.Init(test_dll_obj_path_));
117 E : CoffDecomposer decomposer(image_file_);
118 E : ASSERT_TRUE(decomposer.Decompose(&image_layout_));
119 E : }
120 :
121 : // Reorder and lay out test_dll.coff_obj into a new object file, located
122 : // at new_test_dll_obj_path_.
123 E : void LayoutAndWriteNew(block_graph::BlockGraphOrdererInterface* orderer) {
124 E : DCHECK(orderer != NULL);
125 :
126 : // Fetch headers block.
127 E : ConstTypedBlock<IMAGE_FILE_HEADER> file_header;
128 : BlockGraph::Block* headers_block =
129 E : image_layout_.blocks.GetBlockByAddress(RelativeAddress(0));
130 E : ASSERT_TRUE(headers_block != NULL);
131 E : ASSERT_TRUE(file_header.Init(0, headers_block));
132 :
133 : // Reorder using the specified ordering.
134 E : OrderedBlockGraph ordered_graph(&block_graph_);
135 E : ASSERT_TRUE(orderer->OrderBlockGraph(&ordered_graph, headers_block));
136 :
137 : // Wipe references from headers, so we can remove relocation blocks
138 : // during laying out.
139 E : ASSERT_TRUE(headers_block->RemoveAllReferences());
140 :
141 : // Lay out new image.
142 E : ImageLayout new_image_layout(&block_graph_);
143 E : CoffImageLayoutBuilder layout_builder(&new_image_layout);
144 E : ASSERT_TRUE(layout_builder.LayoutImage(ordered_graph));
145 :
146 : // Write temporary image file.
147 E : CoffFileWriter writer(&new_image_layout);
148 E : ASSERT_TRUE(writer.WriteImage(new_test_dll_obj_path_));
149 E : }
150 :
151 : // Decompose, reorder, and lay out test_dll.coff_obj.
152 E : void RewriteTestDllObj(block_graph::BlockGraphOrdererInterface* orderer) {
153 E : DCHECK(orderer != NULL);
154 :
155 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
156 E : ASSERT_NO_FATAL_FAILURE(LayoutAndWriteNew(orderer));
157 E : }
158 :
159 : // Call the linker to produce a new test DLL located at
160 : // new_test_dll_obj_path_.
161 E : void LinkNewTestDll() {
162 : // Link the rewritten object file into a new DLL.
163 E : base::LaunchOptions opts;
164 E : opts.wait = true;
165 :
166 : // Build linker command line.
167 E : CommandLine::StringVector args;
168 E : args.push_back(testing::kToolchainWrapperPath);
169 E : args.push_back(L"LINK.EXE");
170 E : args.push_back(L"/NOLOGO");
171 E : args.push_back(L"/INCREMENTAL:NO");
172 E : args.push_back(L"/DEBUG");
173 E : args.push_back(L"/PROFILE");
174 E : args.push_back(L"/SAFESEH");
175 E : args.push_back(L"/LARGEADDRESSAWARE");
176 E : args.push_back(L"/NXCOMPAT");
177 E : args.push_back(L"/NODEFAULTLIB:libcmtd.lib");
178 E : args.push_back(L"/DLL");
179 E : args.push_back(L"/MACHINE:X86");
180 E : args.push_back(L"/SUBSYSTEM:CONSOLE");
181 :
182 E : args.push_back(L"/OUT:" + new_test_dll_path_.value());
183 : args.push_back(L"/IMPLIB:" +
184 E : temp_dir_path_.Append(L"test_dll.lib").value());
185 : args.push_back(L"/PDB:" +
186 E : temp_dir_path_.Append(L"test_dll.dll.pdb").value());
187 :
188 : args.push_back(L"/LIBPATH:" +
189 E : testing::GetExeTestDataRelativePath(L".").value());
190 E : args.push_back(L"ole32.lib");
191 E : args.push_back(L"export_dll.lib");
192 E : args.push_back(L"test_dll_no_private_symbols.lib");
193 :
194 : base::FilePath def_path(
195 E : testing::GetSrcRelativePath(L"syzygy\\pe\\test_dll.def"));
196 : base::FilePath label_test_func_obj_path(
197 E : testing::GetExeTestDataRelativePath(L"test_dll_label_test_func.obj"));
198 E : args.push_back(L"/DEF:" + def_path.value());
199 E : args.push_back(label_test_func_obj_path.value());
200 E : args.push_back(new_test_dll_obj_path_.value());
201 :
202 : // Link and check result.
203 E : CommandLine::StringType cmd_line;
204 E : MakeCommandLineString(args, &cmd_line);
205 E : ASSERT_TRUE(base::LaunchProcess(cmd_line, opts, NULL));
206 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(new_test_dll_path_));
207 E : }
208 :
209 : base::FilePath test_dll_obj_path_;
210 : base::FilePath new_test_dll_obj_path_;
211 : base::FilePath new_test_dll_path_;
212 : base::FilePath temp_dir_path_;
213 :
214 : // Original image details.
215 : CoffFile image_file_;
216 : BlockGraph block_graph_;
217 : ImageLayout image_layout_;
218 : };
219 :
220 E : bool IsDebugBlock(const BlockGraph::Block& block, const BlockGraph& graph) {
221 E : if (block.section() == BlockGraph::kInvalidSectionId)
222 E : return false;
223 E : const BlockGraph::Section* section = graph.GetSectionById(block.section());
224 E : if (section->name() != ".debug$S")
225 E : return false;
226 E : return true;
227 E : }
228 :
229 : } // namespace
230 :
231 E : TEST_F(CoffImageLayoutBuilderTest, Redecompose) {
232 E : block_graph::orderers::OriginalOrderer orig_orderer;
233 E : ASSERT_NO_FATAL_FAILURE(RewriteTestDllObj(&orig_orderer));
234 :
235 : // Redecompose.
236 E : CoffFile image_file;
237 E : ASSERT_TRUE(image_file.Init(new_test_dll_obj_path_));
238 :
239 E : CoffDecomposer decomposer(image_file);
240 E : block_graph::BlockGraph block_graph;
241 E : pe::ImageLayout image_layout(&block_graph);
242 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
243 :
244 : // Compare the results of the two decompositions.
245 E : ConstTypedBlock<IMAGE_FILE_HEADER> file_header;
246 : BlockGraph::Block* headers_block =
247 E : image_layout.blocks.GetBlockByAddress(RelativeAddress(0));
248 E : ASSERT_TRUE(headers_block != NULL);
249 E : ASSERT_TRUE(file_header.Init(0, headers_block));
250 :
251 E : EXPECT_EQ(image_layout_.sections.size(), image_layout.sections.size());
252 E : EXPECT_EQ(image_layout_.blocks.size(), image_layout.blocks.size());
253 :
254 : // Expect same sections in the same order, due to original ordering.
255 E : for (size_t i = 0; i < image_layout_.sections.size(); ++i) {
256 E : EXPECT_EQ(image_layout_.sections[i].name, image_layout.sections[i].name);
257 : EXPECT_EQ(image_layout_.sections[i].characteristics,
258 E : image_layout.sections[i].characteristics);
259 E : }
260 E : }
261 :
262 E : TEST_F(CoffImageLayoutBuilderTest, Shuffle) {
263 E : ShuffleOrderer shuffle_orderer(1234);
264 E : ASSERT_NO_FATAL_FAILURE(RewriteTestDllObj(&shuffle_orderer));
265 :
266 : // Save rounded-up sizes of sections to which symbols point; these should
267 : // not change even if sections are shuffled.
268 E : std::vector<size_t> symbol_ref_block_sizes;
269 : BlockGraph::BlockMap::const_iterator it =
270 E : block_graph_.blocks().begin();
271 E : for (; it != block_graph_.blocks().end(); ++it) {
272 E : if ((it->second.attributes() & BlockGraph::COFF_SYMBOL_TABLE) == 0)
273 E : continue;
274 :
275 : BlockGraph::Block::ReferenceMap::const_iterator ref_it =
276 E : it->second.references().begin();
277 E : for (; ref_it != it->second.references().end(); ++ref_it) {
278 : symbol_ref_block_sizes.push_back(
279 E : common::AlignUp(ref_it->second.referenced()->data_size(), 4));
280 E : }
281 E : }
282 :
283 : // Redecompose.
284 E : CoffFile image_file;
285 E : ASSERT_TRUE(image_file.Init(new_test_dll_obj_path_));
286 :
287 E : CoffDecomposer decomposer(image_file);
288 E : block_graph::BlockGraph block_graph;
289 E : pe::ImageLayout image_layout(&block_graph);
290 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
291 :
292 : // Compare the results of the two decompositions.
293 E : ConstTypedBlock<IMAGE_FILE_HEADER> file_header;
294 : BlockGraph::Block* headers_block =
295 E : image_layout.blocks.GetBlockByAddress(RelativeAddress(0));
296 E : ASSERT_TRUE(headers_block != NULL);
297 E : ASSERT_TRUE(file_header.Init(0, headers_block));
298 :
299 E : EXPECT_EQ(image_layout_.sections.size(), image_layout.sections.size());
300 E : EXPECT_EQ(image_layout_.blocks.size(), image_layout.blocks.size());
301 :
302 : // Expect symbols to point to the same blocks they did in the previous
303 : // graph, though they will have been shuffled around.
304 E : it = block_graph.blocks().begin();
305 E : for (; it != block_graph.blocks().end(); ++it) {
306 E : if ((it->second.attributes() & BlockGraph::COFF_SYMBOL_TABLE) == 0)
307 E : continue;
308 :
309 E : ASSERT_EQ(symbol_ref_block_sizes.size(), it->second.references().size());
310 :
311 E : size_t i = 0;
312 : BlockGraph::Block::ReferenceMap::const_iterator ref_it =
313 E : it->second.references().begin();
314 E : for (; ref_it != it->second.references().end(); ++ref_it) {
315 : EXPECT_EQ(symbol_ref_block_sizes[i++],
316 E : ref_it->second.referenced()->data_size());
317 E : }
318 E : EXPECT_EQ(symbol_ref_block_sizes.size(), i);
319 E : }
320 E : }
321 :
322 E : TEST_F(CoffImageLayoutBuilderTest, ShiftedCode) {
323 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
324 :
325 : // Store hard-coded references in debug sections.
326 E : std::vector<BlockGraph::Reference> orig_refs;
327 : BlockGraph::BlockMap::const_iterator it =
328 E : block_graph_.blocks().begin();
329 E : for (; it != block_graph_.blocks().end(); ++it) {
330 E : if (!IsDebugBlock(it->second, block_graph_))
331 E : continue;
332 : BlockGraph::Block::ReferenceMap::const_iterator ref_it =
333 E : it->second.references().begin();
334 E : for (; ref_it != it->second.references().end(); ++ref_it) {
335 : if (ref_it->second.referenced()->type() != BlockGraph::CODE_BLOCK ||
336 E : (ref_it->second.type() & BlockGraph::RELOC_REF_BIT) != 0) {
337 E : continue;
338 : }
339 E : orig_refs.push_back(ref_it->second);
340 E : }
341 E : }
342 :
343 : // Shift every code block.
344 E : BlockGraph::BlockMap& blocks = block_graph_.blocks_mutable();
345 E : BlockGraph::BlockMap::iterator mutable_it = blocks.begin();
346 E : for (; mutable_it != blocks.end(); ++mutable_it) {
347 E : if (mutable_it->second.type() != BlockGraph::CODE_BLOCK)
348 E : continue;
349 :
350 E : mutable_it->second.InsertData(0, 11, false);
351 E : uint8* data = mutable_it->second.GetMutableData();
352 E : for (size_t i = 0; i < 11; ++i) {
353 : // NOP.
354 E : data[i] = 0x90;
355 E : }
356 E : }
357 :
358 : // Check that references have been shifted.
359 E : size_t i = 0;
360 E : it = block_graph_.blocks().begin();
361 E : for (; it != block_graph_.blocks().end(); ++it) {
362 E : if (!IsDebugBlock(it->second, block_graph_))
363 E : continue;
364 : BlockGraph::Block::ReferenceMap::const_iterator ref_it =
365 E : it->second.references().begin();
366 E : for (; ref_it != it->second.references().end(); ++ref_it) {
367 : if (ref_it->second.referenced()->type() != BlockGraph::CODE_BLOCK ||
368 E : (ref_it->second.type() & BlockGraph::RELOC_REF_BIT) != 0) {
369 E : continue;
370 : }
371 E : ASSERT_EQ(orig_refs[i++].offset() + 11, ref_it->second.offset());
372 E : }
373 E : }
374 :
375 : // Write the new object.
376 E : block_graph::orderers::OriginalOrderer orig_orderer;
377 E : ASSERT_NO_FATAL_FAILURE(LayoutAndWriteNew(&orig_orderer));
378 :
379 : // Redecompose.
380 E : CoffFile image_file;
381 E : ASSERT_TRUE(image_file.Init(new_test_dll_obj_path_));
382 :
383 E : CoffDecomposer decomposer(image_file);
384 E : block_graph::BlockGraph block_graph;
385 E : pe::ImageLayout image_layout(&block_graph);
386 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
387 :
388 : // Compare references.
389 E : i = 0;
390 E : it = block_graph.blocks().begin();
391 E : for (; it != block_graph.blocks().end(); ++it) {
392 E : if (!IsDebugBlock(it->second, block_graph))
393 E : continue;
394 : BlockGraph::Block::ReferenceMap::const_iterator ref_it =
395 E : it->second.references().begin();
396 E : for (; ref_it != it->second.references().end(); ++ref_it) {
397 : if (ref_it->second.referenced()->type() != BlockGraph::CODE_BLOCK ||
398 E : (ref_it->second.type() & BlockGraph::RELOC_REF_BIT) != 0) {
399 E : continue;
400 : }
401 E : EXPECT_EQ(orig_refs[i++].offset() + 11, ref_it->second.offset());
402 E : }
403 E : }
404 E : }
405 :
406 E : TEST_F(CoffImageLayoutBuilderTest, RedecomposePE) {
407 E : block_graph::orderers::OriginalOrderer orig_orderer;
408 E : ASSERT_NO_FATAL_FAILURE(RewriteTestDllObj(&orig_orderer));
409 E : ASSERT_NO_FATAL_FAILURE(LinkNewTestDll());
410 :
411 E : PEFile pe_file;
412 E : ASSERT_TRUE(pe_file.Init(new_test_dll_path_));
413 :
414 E : NewDecomposer pe_decomposer(pe_file);
415 E : block_graph::BlockGraph pe_block_graph;
416 E : pe::ImageLayout pe_image_layout(&pe_block_graph);
417 E : ASSERT_TRUE(pe_decomposer.Decompose(&pe_image_layout));
418 E : }
419 :
420 E : TEST_F(CoffImageLayoutBuilderTest, RedecomposeRandom) {
421 E : ShuffleOrderer shuffle_orderer(1234);
422 E : ASSERT_NO_FATAL_FAILURE(RewriteTestDllObj(&shuffle_orderer));
423 :
424 : // Redecompose.
425 E : CoffFile image_file;
426 E : ASSERT_TRUE(image_file.Init(new_test_dll_obj_path_));
427 :
428 E : CoffDecomposer decomposer(image_file);
429 E : block_graph::BlockGraph block_graph;
430 E : pe::ImageLayout image_layout(&block_graph);
431 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
432 :
433 : // Compare the results of the two decompositions.
434 E : ConstTypedBlock<IMAGE_FILE_HEADER> file_header;
435 : BlockGraph::Block* headers_block =
436 E : image_layout.blocks.GetBlockByAddress(RelativeAddress(0));
437 E : ASSERT_TRUE(headers_block != NULL);
438 E : ASSERT_TRUE(file_header.Init(0, headers_block));
439 :
440 E : EXPECT_EQ(image_layout_.sections.size(), image_layout.sections.size());
441 E : EXPECT_EQ(image_layout_.blocks.size(), image_layout.blocks.size());
442 E : }
443 :
444 E : TEST_F(CoffImageLayoutBuilderTest, RedecomposePERandom) {
445 E : ShuffleOrderer shuffle_orderer(1234);
446 E : ASSERT_NO_FATAL_FAILURE(RewriteTestDllObj(&shuffle_orderer));
447 E : ASSERT_NO_FATAL_FAILURE(LinkNewTestDll());
448 :
449 E : PEFile pe_file;
450 E : ASSERT_TRUE(pe_file.Init(new_test_dll_path_));
451 :
452 E : NewDecomposer pe_decomposer(pe_file);
453 E : block_graph::BlockGraph pe_block_graph;
454 E : pe::ImageLayout pe_image_layout(&pe_block_graph);
455 E : ASSERT_TRUE(pe_decomposer.Decompose(&pe_image_layout));
456 E : }
457 :
458 : } // namespace pe
|