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