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 { return "ShuffleOrderer"; }
51 :
52 : // Shuffle sections while paying attention to put .debug$S sections at the
53 : // end, so that they come after the associated sections.
54 : virtual bool OrderBlockGraph(
55 : OrderedBlockGraph* ordered_graph,
56 E : BlockGraph::Block* /* headers_block */) override {
57 E : DCHECK(ordered_graph != NULL);
58 E : BlockGraph* graph = ordered_graph->block_graph();
59 E : DCHECK(graph != NULL);
60 :
61 E : std::vector<BlockGraph::SectionId> sections;
62 E : for (size_t i = 0; i < graph->sections().size(); ++i)
63 E : sections.push_back(i);
64 :
65 E : std::random_shuffle(sections.begin(), sections.end(), rng_);
66 E : for (size_t i = 0; i < sections.size(); ++i) {
67 E : BlockGraph::Section* section = graph->GetSectionById(sections[i]);
68 E : if (section->name() == ".debug$S")
69 E : ordered_graph->PlaceAtTail(section);
70 E : else
71 E : ordered_graph->PlaceAtHead(section);
72 E : }
73 :
74 E : return true;
75 E : }
76 :
77 : private:
78 : core::RandomNumberGenerator rng_;
79 : };
80 :
81 : // We can't rely on base::CommandLine to built the command-line string for us
82 : // because it doesn't maintain the order of arguments.
83 : void MakeCommandLineString(const base::CommandLine::StringVector& args,
84 E : base::CommandLine::StringType* cmd_line) {
85 E : DCHECK(cmd_line != NULL);
86 :
87 E : cmd_line->clear();
88 E : for (size_t i = 0; i < args.size(); ++i) {
89 E : if (i > 0)
90 E : cmd_line->push_back(L' ');
91 E : cmd_line->push_back(L'"');
92 E : cmd_line->append(args[i]);
93 E : cmd_line->push_back(L'"');
94 E : }
95 E : }
96 :
97 : class CoffImageLayoutBuilderTest : public testing::PELibUnitTest {
98 : public:
99 E : CoffImageLayoutBuilderTest() : image_layout_(&block_graph_) {
100 E : }
101 :
102 E : virtual void SetUp() override {
103 E : testing::PELibUnitTest::SetUp();
104 :
105 : test_dll_obj_path_ =
106 E : testing::GetExeTestDataRelativePath(testing::kTestDllCoffObjName);
107 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir_path_));
108 E : new_test_dll_obj_path_ = temp_dir_path_.Append(L"test_dll.obj");
109 E : new_test_dll_path_ = temp_dir_path_.Append(testing::kTestDllName);
110 E : }
111 :
112 : protected:
113 : // Decompose test_dll.coff_obj.
114 E : void DecomposeOriginal() {
115 E : ASSERT_TRUE(image_file_.Init(test_dll_obj_path_));
116 E : CoffDecomposer decomposer(image_file_);
117 E : ASSERT_TRUE(decomposer.Decompose(&image_layout_));
118 E : }
119 :
120 : // Reorder and lay out test_dll.coff_obj into a new object file, located
121 : // at new_test_dll_obj_path_.
122 E : void LayoutAndWriteNew(block_graph::BlockGraphOrdererInterface* orderer) {
123 E : DCHECK(orderer != NULL);
124 :
125 : // Fetch headers block.
126 E : ConstTypedBlock<IMAGE_FILE_HEADER> file_header;
127 : BlockGraph::Block* headers_block =
128 E : image_layout_.blocks.GetBlockByAddress(RelativeAddress(0));
129 E : ASSERT_TRUE(headers_block != NULL);
130 E : ASSERT_TRUE(file_header.Init(0, headers_block));
131 :
132 : // Reorder using the specified ordering.
133 E : OrderedBlockGraph ordered_graph(&block_graph_);
134 E : ASSERT_TRUE(orderer->OrderBlockGraph(&ordered_graph, headers_block));
135 :
136 : // Wipe references from headers, so we can remove relocation blocks
137 : // during laying out.
138 E : ASSERT_TRUE(headers_block->RemoveAllReferences());
139 :
140 : // Lay out new image.
141 E : ImageLayout new_image_layout(&block_graph_);
142 E : CoffImageLayoutBuilder layout_builder(&new_image_layout);
143 E : ASSERT_TRUE(layout_builder.LayoutImage(ordered_graph));
144 :
145 : // Write temporary image file.
146 E : CoffFileWriter writer(&new_image_layout);
147 E : ASSERT_TRUE(writer.WriteImage(new_test_dll_obj_path_));
148 E : }
149 :
150 : // Decompose, reorder, and lay out test_dll.coff_obj.
151 E : void RewriteTestDllObj(block_graph::BlockGraphOrdererInterface* orderer) {
152 E : DCHECK(orderer != NULL);
153 :
154 E : ASSERT_NO_FATAL_FAILURE(DecomposeOriginal());
155 E : ASSERT_NO_FATAL_FAILURE(LayoutAndWriteNew(orderer));
156 E : }
157 :
158 : // Call the linker to produce a new test DLL located at
159 : // new_test_dll_obj_path_.
160 E : void LinkNewTestDll() {
161 : // Link the rewritten object file into a new DLL.
162 E : base::LaunchOptions opts;
163 E : opts.wait = true;
164 :
165 : // Build linker command line.
166 E : base::CommandLine::StringVector args;
167 E : args.push_back(testing::kToolchainWrapperPath);
168 E : args.push_back(L"LINK.EXE");
169 E : args.push_back(L"/NOLOGO");
170 E : args.push_back(L"/INCREMENTAL:NO");
171 E : args.push_back(L"/DEBUG");
172 E : args.push_back(L"/PROFILE");
173 E : args.push_back(L"/SAFESEH");
174 E : args.push_back(L"/LARGEADDRESSAWARE");
175 E : args.push_back(L"/NXCOMPAT");
176 E : args.push_back(L"/NODEFAULTLIB:libcmtd.lib");
177 E : args.push_back(L"/DLL");
178 E : args.push_back(L"/MACHINE:X86");
179 E : args.push_back(L"/SUBSYSTEM:CONSOLE");
180 :
181 E : args.push_back(L"/OUT:" + new_test_dll_path_.value());
182 : args.push_back(L"/IMPLIB:" +
183 E : temp_dir_path_.Append(L"test_dll.dll.lib").value());
184 : args.push_back(L"/PDB:" +
185 E : temp_dir_path_.Append(L"test_dll.dll.pdb").value());
186 :
187 : args.push_back(L"/LIBPATH:" +
188 E : testing::GetExeTestDataRelativePath(L".").value());
189 E : args.push_back(L"ole32.lib");
190 E : args.push_back(L"export_dll.dll.lib");
191 E : args.push_back(L"test_dll_no_private_symbols.lib");
192 :
193 : base::FilePath def_path(
194 E : testing::GetSrcRelativePath(L"syzygy\\pe\\test_dll.def"));
195 : base::FilePath label_test_func_obj_path(
196 E : testing::GetExeTestDataRelativePath(L"test_dll_label_test_func.obj"));
197 E : args.push_back(L"/DEF:" + def_path.value());
198 E : args.push_back(label_test_func_obj_path.value());
199 E : args.push_back(new_test_dll_obj_path_.value());
200 :
201 : // Link and check result.
202 E : base::CommandLine::StringType cmd_line;
203 E : MakeCommandLineString(args, &cmd_line);
204 E : base::Process process = base::LaunchProcess(cmd_line, opts);
205 E : ASSERT_TRUE(process.IsValid());
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] = testing::kNop1[0];
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 : Decomposer 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 : Decomposer 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
|