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