1 : // Copyright 2012 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/relink/relink_app.h"
16 :
17 : #include "base/stringprintf.h"
18 : #include "gmock/gmock.h"
19 : #include "gtest/gtest.h"
20 : #include "syzygy/block_graph/basic_block_decomposer.h"
21 : #include "syzygy/block_graph/typed_block.h"
22 : #include "syzygy/block_graph/unittest_util.h"
23 : #include "syzygy/common/unittest_util.h"
24 : #include "syzygy/core/unittest_util.h"
25 : #include "syzygy/pe/decomposer.h"
26 : #include "syzygy/pe/pe_utils.h"
27 : #include "syzygy/pe/unittest_util.h"
28 : #include "syzygy/reorder/reorderer.h"
29 :
30 : namespace relink {
31 :
32 : using block_graph::BlockGraph;
33 : using common::Application;
34 : using core::RelativeAddress;
35 : using ::testing::ScopedLogLevelSaver;
36 :
37 : namespace {
38 :
39 : class TestRelinkApp : public RelinkApp {
40 : public:
41 : using RelinkApp::input_image_path_;
42 : using RelinkApp::input_pdb_path_;
43 : using RelinkApp::output_image_path_;
44 : using RelinkApp::output_pdb_path_;
45 : using RelinkApp::order_file_path_;
46 : using RelinkApp::seed_;
47 : using RelinkApp::padding_;
48 : using RelinkApp::no_augment_pdb_;
49 : using RelinkApp::compress_pdb_;
50 : using RelinkApp::no_strip_strings_;
51 : using RelinkApp::output_metadata_;
52 : using RelinkApp::overwrite_;
53 : };
54 :
55 : typedef common::Application<TestRelinkApp> TestApp;
56 :
57 : class RelinkAppTest : public testing::PELibUnitTest {
58 : public:
59 : typedef testing::PELibUnitTest Super;
60 :
61 : RelinkAppTest()
62 : : cmd_line_(FilePath(L"relink.exe")),
63 : test_impl_(test_app_.implementation()),
64 : seed_(1234567),
65 : padding_(32),
66 : no_augment_pdb_(false),
67 : compress_pdb_(false),
68 : no_strip_strings_(false),
69 : output_metadata_(false),
70 E : overwrite_(false) {
71 E : }
72 :
73 E : void SetUp() {
74 E : Super::SetUp();
75 :
76 : // Several of the tests generate progress and (deliberate) error messages
77 : // that would otherwise clutter the unittest output.
78 E : logging::SetMinLogLevel(logging::LOG_FATAL);
79 :
80 : // Setup the IO streams.
81 E : CreateTemporaryDir(&temp_dir_);
82 E : stdin_path_ = temp_dir_.Append(L"NUL");
83 E : stdout_path_ = temp_dir_.Append(L"stdout.txt");
84 E : stderr_path_ = temp_dir_.Append(L"stderr.txt");
85 E : InitStreams(stdin_path_, stdout_path_, stderr_path_);
86 :
87 : // Initialize the (potential) input and output path values.
88 E : abs_input_image_path_ = testing::GetExeRelativePath(testing::kTestDllName);
89 E : input_image_path_ = testing::GetRelativePath(abs_input_image_path_);
90 E : abs_input_pdb_path_ = testing::GetExeRelativePath(testing::kTestDllPdbName);
91 E : input_pdb_path_ = testing::GetRelativePath(abs_input_pdb_path_);
92 E : output_image_path_ = temp_dir_.Append(input_image_path_.BaseName());
93 E : output_pdb_path_ = temp_dir_.Append(input_pdb_path_.BaseName());
94 E : order_file_path_ = temp_dir_.Append(L"order.json");
95 :
96 : // Point the application at the test's command-line and IO streams.
97 E : test_app_.set_command_line(&cmd_line_);
98 E : test_app_.set_in(in());
99 E : test_app_.set_out(out());
100 E : test_app_.set_err(err());
101 E : }
102 :
103 E : void GetDllMain(const pe::ImageLayout& layout, BlockGraph::Block** dll_main) {
104 E : ASSERT_TRUE(dll_main != NULL);
105 : BlockGraph::Block* dos_header_block = layout.blocks.GetBlockByAddress(
106 E : RelativeAddress(0));
107 E : ASSERT_TRUE(dos_header_block != NULL);
108 : BlockGraph::Block* nt_headers_block =
109 E : pe::GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
110 E : ASSERT_TRUE(nt_headers_block != NULL);
111 E : block_graph::ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
112 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block));
113 E : BlockGraph::Reference dll_main_ref;
114 E : ASSERT_TRUE(nt_headers_block->GetReference(
115 : nt_headers.OffsetOf(nt_headers->OptionalHeader.AddressOfEntryPoint),
116 : &dll_main_ref));
117 E : ASSERT_EQ(0u, dll_main_ref.offset());
118 E : ASSERT_EQ(0u, dll_main_ref.base());
119 E : *dll_main = dll_main_ref.referenced();
120 E : }
121 :
122 : // Stashes the current log-level before each test instance and restores it
123 : // after each test completes.
124 : ScopedLogLevelSaver log_level_saver;
125 :
126 : // @name The application under test.
127 : // @{
128 : TestApp test_app_;
129 : TestApp::Implementation& test_impl_;
130 : FilePath temp_dir_;
131 : FilePath stdin_path_;
132 : FilePath stdout_path_;
133 : FilePath stderr_path_;
134 : // @}
135 :
136 : // @name Command-line and parameters.
137 : // @{
138 : CommandLine cmd_line_;
139 : FilePath input_image_path_;
140 : FilePath input_pdb_path_;
141 : FilePath output_image_path_;
142 : FilePath output_pdb_path_;
143 : FilePath order_file_path_;
144 : uint32 seed_;
145 : size_t padding_;
146 : bool no_augment_pdb_;
147 : bool compress_pdb_;
148 : bool no_strip_strings_;
149 : bool output_metadata_;
150 : bool overwrite_;
151 : // @}
152 :
153 : // @name Expected final values of input parameters.
154 : // @{
155 : FilePath abs_input_image_path_;
156 : FilePath abs_input_pdb_path_;
157 : // @}
158 : };
159 :
160 : } // namespace
161 :
162 E : TEST_F(RelinkAppTest, GetHelp) {
163 E : cmd_line_.AppendSwitch("help");
164 E : ASSERT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
165 E : }
166 :
167 E : TEST_F(RelinkAppTest, EmptyCommandLineFails) {
168 E : ASSERT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
169 E : }
170 :
171 E : TEST_F(RelinkAppTest, ParseWithNeitherInputNorOrderFails) {
172 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
173 :
174 E : ASSERT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
175 E : }
176 :
177 E : TEST_F(RelinkAppTest, ParseWithSeedAndOrderFails) {
178 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
179 E : cmd_line_.AppendSwitchASCII("seed", base::StringPrintf("%d", seed_));
180 E : cmd_line_.AppendSwitchPath("order_file", order_file_path_);
181 :
182 E : ASSERT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
183 E : }
184 :
185 E : TEST_F(RelinkAppTest, ParseWithEmptySeedFails) {
186 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
187 E : cmd_line_.AppendSwitch("seed");
188 :
189 E : ASSERT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
190 E : }
191 :
192 E : TEST_F(RelinkAppTest, ParseWithInvalidSeedFails) {
193 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
194 E : cmd_line_.AppendSwitchASCII("seed", "hello");
195 :
196 E : ASSERT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
197 E : }
198 :
199 E : TEST_F(RelinkAppTest, ParseWithEmptyPaddingFails) {
200 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
201 E : cmd_line_.AppendSwitch("padding");
202 :
203 E : ASSERT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
204 E : }
205 :
206 E : TEST_F(RelinkAppTest, ParseWithInvalidPaddingFails) {
207 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
208 E : cmd_line_.AppendSwitchASCII("padding", "hello");
209 :
210 E : ASSERT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
211 E : }
212 :
213 E : TEST_F(RelinkAppTest, ParseMinimalCommandLineWithInputDll) {
214 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
215 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
216 :
217 E : EXPECT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
218 E : EXPECT_TRUE(test_impl_.SetUp());
219 E : }
220 :
221 E : TEST_F(RelinkAppTest, ParseMinimalCommandLineWithOrderFile) {
222 : // The order file doesn't actually exist, so setup should fail to infer the
223 : // input dll.
224 E : cmd_line_.AppendSwitchPath("order-file", order_file_path_);
225 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
226 :
227 E : EXPECT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
228 E : EXPECT_FALSE(test_impl_.SetUp());
229 E : }
230 :
231 E : TEST_F(RelinkAppTest, ParseFullCommandLineWithOrderFile) {
232 : // Note that we specify the no-metadata flag, so we expect false below
233 : // for the output_metadata_ member. Also note that neither seed nor padding
234 : // are given, and should default to 0.
235 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
236 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
237 E : cmd_line_.AppendSwitchPath("order-file", order_file_path_);
238 E : cmd_line_.AppendSwitch("no-augment-pdb");
239 E : cmd_line_.AppendSwitch("compress-pdb");
240 E : cmd_line_.AppendSwitch("no-strip-strings");
241 E : cmd_line_.AppendSwitch("no-metadata");
242 E : cmd_line_.AppendSwitch("overwrite");
243 :
244 E : EXPECT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
245 E : EXPECT_TRUE(test_impl_.input_image_path_.empty());
246 E : EXPECT_TRUE(test_impl_.input_pdb_path_.empty());
247 E : EXPECT_EQ(output_image_path_, test_impl_.output_image_path_);
248 E : EXPECT_EQ(output_pdb_path_, test_impl_.output_pdb_path_);
249 E : EXPECT_EQ(order_file_path_, test_impl_.order_file_path_);
250 E : EXPECT_EQ(0, test_impl_.seed_);
251 E : EXPECT_EQ(0, test_impl_.padding_);
252 E : EXPECT_TRUE(test_impl_.no_augment_pdb_);
253 E : EXPECT_TRUE(test_impl_.compress_pdb_);
254 E : EXPECT_TRUE(test_impl_.no_strip_strings_);
255 E : EXPECT_FALSE(test_impl_.output_metadata_);
256 E : EXPECT_TRUE(test_impl_.overwrite_);
257 :
258 : // The order file doesn't actually exist, so setup should fail to infer the
259 : // input dll.
260 E : EXPECT_FALSE(test_impl_.SetUp());
261 E : }
262 :
263 E : TEST_F(RelinkAppTest, ParseFullCommandLineWithInputSeedAndMetadata) {
264 : // Note that we omit the no-metadata flag, so we expect true below for the
265 : // output_metadata_ member.
266 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
267 E : cmd_line_.AppendSwitchPath("input-pdb", input_pdb_path_);
268 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
269 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
270 E : cmd_line_.AppendSwitchASCII("seed", base::StringPrintf("%d", seed_));
271 E : cmd_line_.AppendSwitchASCII("padding", base::StringPrintf("%d", padding_));
272 E : cmd_line_.AppendSwitch("no-augment-pdb");
273 E : cmd_line_.AppendSwitch("compress-pdb");
274 E : cmd_line_.AppendSwitch("no-strip-strings");
275 E : cmd_line_.AppendSwitch("overwrite");
276 :
277 E : EXPECT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
278 E : EXPECT_EQ(abs_input_image_path_, test_impl_.input_image_path_);
279 E : EXPECT_EQ(abs_input_pdb_path_, test_impl_.input_pdb_path_);
280 E : EXPECT_EQ(output_image_path_, test_impl_.output_image_path_);
281 E : EXPECT_EQ(output_pdb_path_, test_impl_.output_pdb_path_);
282 E : EXPECT_TRUE(test_impl_.order_file_path_.empty());
283 E : EXPECT_EQ(seed_, test_impl_.seed_);
284 E : EXPECT_EQ(padding_, test_impl_.padding_);
285 E : EXPECT_TRUE(test_impl_.no_augment_pdb_);
286 E : EXPECT_TRUE(test_impl_.compress_pdb_);
287 E : EXPECT_TRUE(test_impl_.no_strip_strings_);
288 E : EXPECT_TRUE(test_impl_.output_metadata_);
289 E : EXPECT_TRUE(test_impl_.overwrite_);
290 :
291 : // SetUp() has nothing else to infer so it should succeed.
292 E : EXPECT_TRUE(test_impl_.SetUp());
293 E : }
294 :
295 E : TEST_F(RelinkAppTest, DeprecatedFlagsSucceeds) {
296 E : cmd_line_.AppendSwitchPath("input-dll", input_image_path_);
297 E : cmd_line_.AppendSwitchPath("output-dll", output_image_path_);
298 E : EXPECT_TRUE(test_impl_.ParseCommandLine(&cmd_line_));
299 :
300 E : EXPECT_EQ(abs_input_image_path_, test_impl_.input_image_path_);
301 E : EXPECT_EQ(output_image_path_, test_impl_.output_image_path_);
302 E : }
303 :
304 E : TEST_F(RelinkAppTest, DeprecatedFlagsConflictingInputsFail) {
305 E : cmd_line_.AppendSwitchPath("input-dll", input_image_path_);
306 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
307 E : cmd_line_.AppendSwitchPath("output-dll", output_image_path_);
308 E : EXPECT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
309 E : }
310 :
311 E : TEST_F(RelinkAppTest, DeprecatedFlagsConflictingOutputsFail) {
312 E : cmd_line_.AppendSwitchPath("input-dll", input_image_path_);
313 E : cmd_line_.AppendSwitchPath("output-dll", output_image_path_);
314 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
315 E : EXPECT_FALSE(test_impl_.ParseCommandLine(&cmd_line_));
316 E : }
317 :
318 E : TEST_F(RelinkAppTest, RandomRelink) {
319 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
320 E : cmd_line_.AppendSwitchPath("input-pdb", input_pdb_path_);
321 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
322 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
323 E : cmd_line_.AppendSwitchASCII("seed", base::StringPrintf("%d", seed_));
324 E : cmd_line_.AppendSwitchASCII("padding", base::StringPrintf("%d", padding_));
325 E : cmd_line_.AppendSwitch("overwrite");
326 :
327 E : ASSERT_EQ(0, test_app_.Run());
328 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(output_image_path_));
329 E : }
330 :
331 E : TEST_F(RelinkAppTest, RandomRelinkBasicBlocks) {
332 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
333 E : cmd_line_.AppendSwitchPath("input-pdb", input_pdb_path_);
334 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
335 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
336 E : cmd_line_.AppendSwitchASCII("seed", base::StringPrintf("%d", seed_));
337 E : cmd_line_.AppendSwitchASCII("padding", base::StringPrintf("%d", padding_));
338 E : cmd_line_.AppendSwitch("overwrite");
339 E : cmd_line_.AppendSwitch("basic-blocks");
340 E : cmd_line_.AppendSwitch("exclude-bb-padding");
341 :
342 E : ASSERT_EQ(0, test_app_.Run());
343 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(output_image_path_));
344 E : }
345 :
346 E : TEST_F(RelinkAppTest, RelinkBlockOrder) {
347 E : pe::PEFile pe_file;
348 E : ASSERT_TRUE(pe_file.Init(input_image_path_));
349 :
350 E : BlockGraph bg;
351 E : pe::ImageLayout layout(&bg);
352 E : pe::Decomposer decomposer(pe_file);
353 E : ASSERT_TRUE(decomposer.Decompose(&layout));
354 :
355 : // Get the DLL main entry point.
356 E : BlockGraph::Block* dll_main_block = NULL;
357 E : ASSERT_NO_FATAL_FAILURE(GetDllMain(layout, &dll_main_block));
358 :
359 : // Build a block-level ordering by placing the DLL main entry point at the
360 : // beginning of its section.
361 E : BlockGraph::Section* text_section = bg.FindSection(".text");
362 E : ASSERT_TRUE(text_section != NULL);
363 E : reorder::Reorderer::Order order;
364 E : order.sections.resize(1);
365 E : order.sections[0].id = text_section->id();
366 E : order.sections[0].name = text_section->name();
367 E : order.sections[0].characteristics = text_section->characteristics();
368 E : order.sections[0].blocks.resize(1);
369 E : order.sections[0].blocks[0].block = dll_main_block;
370 :
371 : // Serialize the order file.
372 E : ASSERT_TRUE(order.SerializeToJSON(pe_file, order_file_path_, false));
373 :
374 : // Order the test DLL using the order file we just created.
375 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
376 E : cmd_line_.AppendSwitchPath("input-pdb", input_pdb_path_);
377 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
378 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
379 E : cmd_line_.AppendSwitchPath("order-file", order_file_path_);
380 E : ASSERT_EQ(0, test_app_.Run());
381 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(output_image_path_));
382 E : }
383 :
384 E : TEST_F(RelinkAppTest, RelinkTestDataBlockOrder) {
385 : // Try a block-level reordering using an actual order file generated by our
386 : // test_data dependency.
387 :
388 : FilePath test_dll_order_json =
389 E : testing::GetExeTestDataRelativePath(L"test_dll_order.json");
390 :
391 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
392 E : cmd_line_.AppendSwitchPath("input-pdb", input_pdb_path_);
393 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
394 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
395 E : cmd_line_.AppendSwitchPath("order-file", test_dll_order_json);
396 E : ASSERT_EQ(0, test_app_.Run());
397 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(output_image_path_));
398 E : }
399 :
400 E : TEST_F(RelinkAppTest, RelinkBasicBlockOrder) {
401 E : pe::PEFile pe_file;
402 E : ASSERT_TRUE(pe_file.Init(input_image_path_));
403 :
404 E : BlockGraph bg;
405 E : pe::ImageLayout layout(&bg);
406 E : pe::Decomposer decomposer(pe_file);
407 E : ASSERT_TRUE(decomposer.Decompose(&layout));
408 :
409 : // Get the DLL main entry point.
410 E : BlockGraph::Block* dll_main_block = NULL;
411 E : ASSERT_NO_FATAL_FAILURE(GetDllMain(layout, &dll_main_block));
412 :
413 : // Build a block-level ordering by splitting the DLL main block into two
414 : // blocks, each half in a different section.
415 E : BlockGraph::Section* text_section = bg.FindSection(".text");
416 E : ASSERT_TRUE(text_section != NULL);
417 E : reorder::Reorderer::Order order;
418 E : order.sections.resize(2);
419 E : order.sections[0].id = text_section->id();
420 E : order.sections[0].name = text_section->name();
421 E : order.sections[0].characteristics = text_section->characteristics();
422 E : order.sections[0].blocks.resize(1);
423 E : order.sections[0].blocks[0].block = dll_main_block;
424 E : order.sections[1].id = reorder::Reorderer::Order::SectionSpec::kNewSectionId;
425 E : order.sections[1].name = ".text2";
426 E : order.sections[1].characteristics = text_section->characteristics();
427 E : order.sections[1].blocks.resize(1);
428 E : order.sections[1].blocks[0].block = dll_main_block;
429 :
430 : // Decompose the block. Iterate over its basic-blocks and take turns placing
431 : // them into each of the two above blocks.
432 E : block_graph::BasicBlockSubGraph bbsg;
433 E : block_graph::BasicBlockDecomposer bb_decomposer(dll_main_block, &bbsg);
434 E : ASSERT_TRUE(bb_decomposer.Decompose());
435 E : ASSERT_LE(2u, bbsg.basic_blocks().size());
436 : block_graph::BasicBlockSubGraph::BBCollection::const_iterator bb_it =
437 E : bbsg.basic_blocks().begin();
438 E : for (size_t i = 0; bb_it != bbsg.basic_blocks().end(); ++bb_it, i ^= 1) {
439 : order.sections[i].blocks[0].basic_block_offsets.push_back(
440 E : (*bb_it)->offset());
441 E : }
442 :
443 : // Serialize the order file.
444 E : ASSERT_TRUE(order.SerializeToJSON(pe_file, order_file_path_, false));
445 :
446 : // Order the test DLL using the order file we just created.
447 E : cmd_line_.AppendSwitchPath("input-image", input_image_path_);
448 E : cmd_line_.AppendSwitchPath("input-pdb", input_pdb_path_);
449 E : cmd_line_.AppendSwitchPath("output-image", output_image_path_);
450 E : cmd_line_.AppendSwitchPath("output-pdb", output_pdb_path_);
451 E : cmd_line_.AppendSwitchPath("order-file", order_file_path_);
452 E : ASSERT_EQ(0, test_app_.Run());
453 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(output_image_path_));
454 E : }
455 :
456 : } // namespace relink
|