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/pe/pe_relinker.h"
16 :
17 : #include "base/file_util.h"
18 : #include "gmock/gmock.h"
19 : #include "gtest/gtest.h"
20 : #include "syzygy/common/defs.h"
21 : #include "syzygy/core/serialization.h"
22 : #include "syzygy/core/unittest_util.h"
23 : #include "syzygy/pdb/pdb_byte_stream.h"
24 : #include "syzygy/pdb/pdb_file.h"
25 : #include "syzygy/pdb/pdb_reader.h"
26 : #include "syzygy/pdb/pdb_util.h"
27 : #include "syzygy/pe/find.h"
28 : #include "syzygy/pe/metadata.h"
29 : #include "syzygy/pe/pdb_info.h"
30 : #include "syzygy/pe/unittest_util.h"
31 :
32 : namespace pe {
33 :
34 : namespace {
35 :
36 : using block_graph::BlockGraphOrdererInterface;
37 : using block_graph::BlockGraphTransformInterface;
38 : using block_graph::OrderedBlockGraph;
39 : using pdb::PdbFile;
40 : using pdb::PdbMutatorInterface;
41 : using testing::_;
42 : using testing::Return;
43 : using testing::StrictMock;
44 :
45 : class TestPERelinker : public PERelinker {
46 : public:
47 : using PERelinker::transforms_;
48 : using PERelinker::orderers_;
49 : using PERelinker::pdb_mutators_;
50 : };
51 :
52 : class PERelinkerTest : public testing::PELibUnitTest {
53 : typedef testing::PELibUnitTest Super;
54 :
55 : public:
56 E : void SetUp() {
57 E : Super::SetUp();
58 :
59 E : input_dll_ = testing::GetExeRelativePath(testing::kTestDllName);
60 E : input_pdb_ = testing::GetExeRelativePath(testing::kTestDllPdbName);
61 :
62 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir_));
63 E : temp_dll_ = temp_dir_.Append(testing::kTestDllName);
64 E : temp_pdb_ = temp_dir_.Append(testing::kTestDllPdbName);
65 E : }
66 :
67 : base::FilePath input_dll_;
68 : base::FilePath input_pdb_;
69 : base::FilePath temp_dir_;
70 : base::FilePath temp_dll_;
71 : base::FilePath temp_pdb_;
72 : };
73 :
74 : class MockTransform : public BlockGraphTransformInterface {
75 : public:
76 E : const char* name() const { return "MockTransform"; }
77 E : MOCK_METHOD2(TransformBlockGraph, bool(BlockGraph*, BlockGraph::Block*));
78 : };
79 :
80 : class MockOrderer : public BlockGraphOrdererInterface {
81 : public:
82 E : const char* name() const { return "MockOrderer"; }
83 E : MOCK_METHOD2(OrderBlockGraph, bool(OrderedBlockGraph*, BlockGraph::Block*));
84 : };
85 :
86 : class MockPdbMutator : public PdbMutatorInterface {
87 : public:
88 E : const char* name() const { return "MockPdbMutator"; }
89 E : MOCK_METHOD1(MutatePdb, bool(PdbFile*));
90 : };
91 :
92 : } // namespace
93 :
94 E : TEST_F(PERelinkerTest, Properties) {
95 E : TestPERelinker relinker;
96 E : base::FilePath dummy_path(L"foo");
97 :
98 E : EXPECT_EQ(base::FilePath(), relinker.input_path());
99 E : relinker.set_input_path(dummy_path);
100 E : EXPECT_EQ(dummy_path, relinker.input_path());
101 :
102 E : EXPECT_EQ(base::FilePath(), relinker.input_pdb_path());
103 E : relinker.set_input_pdb_path(dummy_path);
104 E : EXPECT_EQ(dummy_path, relinker.input_pdb_path());
105 :
106 E : EXPECT_EQ(base::FilePath(), relinker.output_path());
107 E : relinker.set_output_path(dummy_path);
108 E : EXPECT_EQ(dummy_path, relinker.output_path());
109 :
110 E : EXPECT_EQ(base::FilePath(), relinker.output_pdb_path());
111 E : relinker.set_output_pdb_path(dummy_path);
112 E : EXPECT_EQ(dummy_path, relinker.output_pdb_path());
113 :
114 E : EXPECT_TRUE(relinker.add_metadata());
115 E : relinker.set_add_metadata(false);
116 E : EXPECT_FALSE(relinker.add_metadata());
117 E : relinker.set_add_metadata(TRUE);
118 E : EXPECT_TRUE(relinker.add_metadata());
119 :
120 E : EXPECT_FALSE(relinker.allow_overwrite());
121 E : relinker.set_allow_overwrite(true);
122 E : EXPECT_TRUE(relinker.allow_overwrite());
123 E : relinker.set_allow_overwrite(false);
124 E : EXPECT_FALSE(relinker.allow_overwrite());
125 :
126 E : EXPECT_TRUE(relinker.augment_pdb());
127 E : relinker.set_augment_pdb(false);
128 E : EXPECT_FALSE(relinker.augment_pdb());
129 E : relinker.set_augment_pdb(true);
130 E : EXPECT_TRUE(relinker.augment_pdb());
131 :
132 E : EXPECT_FALSE(relinker.compress_pdb());
133 E : relinker.set_compress_pdb(true);
134 E : EXPECT_TRUE(relinker.compress_pdb());
135 E : relinker.set_compress_pdb(false);
136 E : EXPECT_FALSE(relinker.compress_pdb());
137 :
138 E : EXPECT_TRUE(relinker.parse_debug_info());
139 E : relinker.set_parse_debug_info(false);
140 E : EXPECT_FALSE(relinker.parse_debug_info());
141 E : relinker.set_parse_debug_info(true);
142 E : EXPECT_TRUE(relinker.parse_debug_info());
143 :
144 E : EXPECT_FALSE(relinker.strip_strings());
145 E : relinker.set_strip_strings(true);
146 E : EXPECT_TRUE(relinker.strip_strings());
147 E : relinker.set_strip_strings(false);
148 E : EXPECT_FALSE(relinker.strip_strings());
149 :
150 E : EXPECT_FALSE(relinker.use_new_decomposer());
151 E : relinker.set_use_new_decomposer(true);
152 E : EXPECT_TRUE(relinker.use_new_decomposer());
153 E : relinker.set_use_new_decomposer(false);
154 E : EXPECT_FALSE(relinker.use_new_decomposer());
155 :
156 E : EXPECT_EQ(0u, relinker.padding());
157 E : relinker.set_padding(10);
158 E : EXPECT_EQ(10u, relinker.padding());
159 E : relinker.set_padding(0);
160 E : EXPECT_EQ(0u, relinker.padding());
161 E : }
162 :
163 E : TEST_F(PERelinkerTest, AppendTransforms) {
164 E : TestPERelinker relinker;
165 :
166 E : MockTransform transform1, transform2;
167 E : std::vector<BlockGraphTransformInterface*> transforms;
168 E : transforms.push_back(&transform2);
169 :
170 E : relinker.AppendTransform(&transform1);
171 E : relinker.AppendTransforms(transforms);
172 :
173 E : std::vector<BlockGraphTransformInterface*> expected;
174 E : expected.push_back(&transform1);
175 E : expected.push_back(&transform2);
176 :
177 E : EXPECT_EQ(expected, relinker.transforms_);
178 E : }
179 :
180 E : TEST_F(PERelinkerTest, AppendOrderers) {
181 E : TestPERelinker relinker;
182 :
183 E : MockOrderer orderer1, orderer2;
184 E : std::vector<BlockGraphOrdererInterface*> orderers;
185 E : orderers.push_back(&orderer2);
186 :
187 E : relinker.AppendOrderer(&orderer1);
188 E : relinker.AppendOrderers(orderers);
189 :
190 E : std::vector<BlockGraphOrdererInterface*> expected;
191 E : expected.push_back(&orderer1);
192 E : expected.push_back(&orderer2);
193 :
194 E : EXPECT_EQ(expected, relinker.orderers_);
195 E : }
196 :
197 E : TEST_F(PERelinkerTest, AppendPdbMutators) {
198 E : TestPERelinker relinker;
199 :
200 E : MockPdbMutator pdb_mutator1, pdb_mutator2;
201 E : std::vector<PdbMutatorInterface*> pdb_mutators;
202 E : pdb_mutators.push_back(&pdb_mutator2);
203 :
204 E : relinker.AppendPdbMutator(&pdb_mutator1);
205 E : relinker.AppendPdbMutators(pdb_mutators);
206 :
207 E : std::vector<PdbMutatorInterface*> expected;
208 E : expected.push_back(&pdb_mutator1);
209 E : expected.push_back(&pdb_mutator2);
210 :
211 E : EXPECT_EQ(expected, relinker.pdb_mutators_);
212 E : }
213 :
214 E : TEST_F(PERelinkerTest, InitFailsOnUnspecifiedInput) {
215 E : TestPERelinker relinker;
216 :
217 E : relinker.set_output_path(temp_dll_);
218 E : EXPECT_FALSE(relinker.Init());
219 E : }
220 :
221 E : TEST_F(PERelinkerTest, InitFailsOnUnspecifiedOutput) {
222 E : TestPERelinker relinker;
223 :
224 E : relinker.set_input_path(input_dll_);
225 E : EXPECT_FALSE(relinker.Init());
226 E : }
227 :
228 E : TEST_F(PERelinkerTest, InitFailsOnNonexistentInput) {
229 E : TestPERelinker relinker;
230 :
231 E : relinker.set_input_path(temp_dir_.Append(L"nonexistent.dll"));
232 E : relinker.set_output_path(temp_dll_);
233 E : EXPECT_FALSE(relinker.Init());
234 E : }
235 :
236 E : TEST_F(PERelinkerTest, InitFailsOnDisallowedOverwrite) {
237 E : TestPERelinker relinker;
238 :
239 : // Copy the image in case the test actually does overwrite the input; this
240 : // way we don't accidentally turf our test data.
241 E : file_util::CopyFile(input_dll_, temp_dll_);
242 :
243 E : relinker.set_input_path(temp_dll_);
244 E : relinker.set_output_path(temp_dll_);
245 :
246 E : relinker.set_allow_overwrite(false);
247 E : EXPECT_FALSE(relinker.Init());
248 E : }
249 :
250 E : TEST_F(PERelinkerTest, InitSucceeds) {
251 E : TestPERelinker relinker;
252 :
253 E : relinker.set_input_path(input_dll_);
254 E : relinker.set_output_path(temp_dll_);
255 :
256 E : EXPECT_TRUE(relinker.Init());
257 E : }
258 :
259 E : TEST_F(PERelinkerTest, IntermediateAccessors) {
260 E : TestPERelinker relinker;
261 :
262 E : relinker.set_input_path(input_dll_);
263 E : relinker.set_output_path(temp_dll_);
264 :
265 E : EXPECT_TRUE(relinker.Init());
266 :
267 E : EXPECT_EQ(input_dll_, relinker.input_pe_file().path());
268 E : EXPECT_TRUE(relinker.dos_header_block() != NULL);
269 E : }
270 :
271 E : TEST_F(PERelinkerTest, FailsWhenTransformFails) {
272 E : TestPERelinker relinker;
273 E : StrictMock<MockTransform> transform;
274 :
275 E : EXPECT_CALL(transform, TransformBlockGraph(_, _)).WillOnce(Return(false));
276 :
277 E : relinker.AppendTransform(&transform);
278 E : relinker.set_input_path(input_dll_);
279 E : relinker.set_output_path(temp_dll_);
280 E : EXPECT_TRUE(relinker.Init());
281 E : EXPECT_FALSE(relinker.Relink());
282 E : }
283 :
284 E : TEST_F(PERelinkerTest, FailsWhenOrdererFails) {
285 E : TestPERelinker relinker;
286 E : StrictMock<MockOrderer> orderer;
287 :
288 E : EXPECT_CALL(orderer, OrderBlockGraph(_, _)).WillOnce(Return(false));
289 :
290 E : relinker.AppendOrderer(&orderer);
291 E : relinker.set_input_path(input_dll_);
292 E : relinker.set_output_path(temp_dll_);
293 E : EXPECT_TRUE(relinker.Init());
294 E : EXPECT_FALSE(relinker.Relink());
295 E : }
296 :
297 E : TEST_F(PERelinkerTest, FailsWhenPdbMutatorFails) {
298 E : TestPERelinker relinker;
299 E : StrictMock<MockPdbMutator> pdb_mutator;
300 :
301 E : EXPECT_CALL(pdb_mutator, MutatePdb(_)).WillOnce(Return(false));
302 :
303 E : relinker.AppendPdbMutator(&pdb_mutator);
304 E : relinker.set_input_path(input_dll_);
305 E : relinker.set_output_path(temp_dll_);
306 E : EXPECT_TRUE(relinker.Init());
307 E : EXPECT_FALSE(relinker.Relink());
308 E : }
309 :
310 E : TEST_F(PERelinkerTest, Success) {
311 E : TestPERelinker relinker;
312 E : StrictMock<MockTransform> transform;
313 E : StrictMock<MockOrderer> orderer;
314 E : StrictMock<MockPdbMutator> pdb_mutator;
315 :
316 E : EXPECT_CALL(transform, TransformBlockGraph(_, _)).WillOnce(Return(true));
317 E : EXPECT_CALL(orderer, OrderBlockGraph(_, _)).WillOnce(Return(true));
318 E : EXPECT_CALL(pdb_mutator, MutatePdb(_)).WillOnce(Return(true));
319 :
320 E : relinker.AppendTransform(&transform);
321 E : relinker.AppendOrderer(&orderer);
322 E : relinker.AppendPdbMutator(&pdb_mutator);
323 :
324 E : relinker.set_input_path(input_dll_);
325 E : relinker.set_output_path(temp_dll_);
326 :
327 E : EXPECT_TRUE(relinker.Init());
328 E : EXPECT_TRUE(relinker.Relink());
329 E : }
330 :
331 E : TEST_F(PERelinkerTest, IdentityRelink) {
332 E : TestPERelinker relinker;
333 :
334 E : relinker.set_input_path(input_dll_);
335 E : relinker.set_output_path(temp_dll_);
336 :
337 : // We let the relinker infer the PDB output. The mechanism should cause it
338 : // to produce a PDB file in the temporary directory with the same basename
339 : // as the input PDB.
340 E : EXPECT_TRUE(relinker.Init());
341 E : EXPECT_TRUE(relinker.Relink());
342 E : EXPECT_EQ(temp_pdb_, relinker.output_pdb_path());
343 :
344 E : EXPECT_TRUE(file_util::PathExists(relinker.output_path()));
345 E : EXPECT_TRUE(file_util::PathExists(relinker.output_pdb_path()));
346 :
347 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(relinker.output_path()));
348 :
349 E : PEFile orig_pe_file;
350 E : PEFile::Signature orig_pe_sig;
351 E : ASSERT_TRUE(orig_pe_file.Init(input_dll_));
352 E : orig_pe_file.GetSignature(&orig_pe_sig);
353 :
354 : // Ensure that the produced binary contains a metadata section. This
355 : // confirms that the AddMetadataTransform has run.
356 E : PEFile new_pe_file;
357 E : ASSERT_TRUE(new_pe_file.Init(temp_dll_));
358 : ASSERT_NE(kInvalidSection,
359 E : new_pe_file.GetSectionIndex(common::kSyzygyMetadataSectionName));
360 E : Metadata metadata;
361 E : ASSERT_TRUE(metadata.LoadFromPE(new_pe_file));
362 E : EXPECT_TRUE(metadata.IsConsistent(orig_pe_sig));
363 :
364 : // Ensure that the PDB file can be found from the module. This confirms that
365 : // the AddPdbInfoTransform has run.
366 :
367 E : PdbInfo pdb_info;
368 E : ASSERT_TRUE(pdb_info.Init(relinker.output_path()));
369 E : EXPECT_EQ(pdb_info.pdb_file_name(), relinker.output_pdb_path());
370 :
371 E : base::FilePath pdb_path;
372 E : ASSERT_TRUE(FindPdbForModule(relinker.output_path(), &pdb_path));
373 E : EXPECT_EQ(pdb_path, relinker.output_pdb_path());
374 E : }
375 :
376 E : TEST_F(PERelinkerTest, IdentityRelinkNewDecomposer) {
377 E : TestPERelinker relinker;
378 :
379 E : relinker.set_input_path(input_dll_);
380 E : relinker.set_output_path(temp_dll_);
381 E : relinker.set_parse_debug_info(false);
382 E : relinker.set_use_new_decomposer(true);
383 :
384 : // We let the relinker infer the PDB output. The mechanism should cause it
385 : // to produce a PDB file in the temporary directory with the same basename
386 : // as the input PDB.
387 E : EXPECT_TRUE(relinker.Init());
388 E : EXPECT_TRUE(relinker.Relink());
389 E : EXPECT_EQ(temp_pdb_, relinker.output_pdb_path());
390 :
391 E : EXPECT_TRUE(file_util::PathExists(relinker.output_path()));
392 E : EXPECT_TRUE(file_util::PathExists(relinker.output_pdb_path()));
393 :
394 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(relinker.output_path()));
395 :
396 E : PEFile orig_pe_file;
397 E : PEFile::Signature orig_pe_sig;
398 E : ASSERT_TRUE(orig_pe_file.Init(input_dll_));
399 E : orig_pe_file.GetSignature(&orig_pe_sig);
400 :
401 : // Ensure that the produced binary contains a metadata section. This
402 : // confirms that the AddMetadataTransform has run.
403 E : PEFile new_pe_file;
404 E : ASSERT_TRUE(new_pe_file.Init(temp_dll_));
405 : ASSERT_NE(kInvalidSection,
406 E : new_pe_file.GetSectionIndex(common::kSyzygyMetadataSectionName));
407 E : Metadata metadata;
408 E : ASSERT_TRUE(metadata.LoadFromPE(new_pe_file));
409 E : EXPECT_TRUE(metadata.IsConsistent(orig_pe_sig));
410 :
411 : // Ensure that the PDB file can be found from the module. This confirms that
412 : // the AddPdbInfoTransform has run.
413 :
414 E : PdbInfo pdb_info;
415 E : ASSERT_TRUE(pdb_info.Init(relinker.output_path()));
416 E : EXPECT_EQ(pdb_info.pdb_file_name(), relinker.output_pdb_path());
417 :
418 E : base::FilePath pdb_path;
419 E : ASSERT_TRUE(FindPdbForModule(relinker.output_path(), &pdb_path));
420 E : EXPECT_EQ(pdb_path, relinker.output_pdb_path());
421 E : }
422 :
423 E : TEST_F(PERelinkerTest, BlockGraphStreamIsCreated) {
424 E : TestPERelinker relinker;
425 :
426 E : relinker.set_input_path(input_dll_);
427 E : relinker.set_output_path(temp_dll_);
428 E : relinker.set_augment_pdb(true);
429 E : EXPECT_TRUE(relinker.augment_pdb());
430 :
431 E : EXPECT_TRUE(relinker.Init());
432 E : EXPECT_TRUE(relinker.Relink());
433 E : EXPECT_EQ(temp_pdb_, relinker.output_pdb_path());
434 :
435 : // Ensure that the block-graph stream has been written to the PDB. The
436 : // content of the stream is not validated, we only check that the named
437 : // stream exists in the generated PDB file.
438 E : pdb::PdbFile pdb_file;
439 E : pdb::PdbReader pdb_reader;
440 E : EXPECT_TRUE(pdb_reader.Read(temp_pdb_, &pdb_file));
441 E : pdb::PdbInfoHeader70 pdb_header = {0};
442 E : pdb::NameStreamMap name_stream_map;
443 : EXPECT_TRUE(ReadHeaderInfoStream(
444 : pdb_file.GetStream(pdb::kPdbHeaderInfoStream),
445 : &pdb_header,
446 E : &name_stream_map));
447 : pdb::NameStreamMap::const_iterator name_it = name_stream_map.find(
448 E : pdb::kSyzygyBlockGraphStreamName);
449 E : ASSERT_TRUE(name_it != name_stream_map.end());
450 E : scoped_refptr<pdb::PdbStream> stream = pdb_file.GetStream(name_it->second);
451 E : ASSERT_TRUE(stream.get() != NULL);
452 E : ASSERT_GT(stream->length(), 0u);
453 E : }
454 :
455 E : TEST_F(PERelinkerTest, BlockGraphStreamVersionIsTheCurrentOne) {
456 E : TestPERelinker relinker;
457 :
458 E : relinker.set_input_path(input_dll_);
459 E : relinker.set_output_path(temp_dll_);
460 E : relinker.set_augment_pdb(true);
461 E : EXPECT_TRUE(relinker.augment_pdb());
462 :
463 E : EXPECT_TRUE(relinker.Init());
464 E : EXPECT_TRUE(relinker.Relink());
465 E : EXPECT_EQ(temp_pdb_, relinker.output_pdb_path());
466 :
467 : // Looks for the block-graph stream in the PDB.
468 E : pdb::PdbFile pdb_file;
469 E : pdb::PdbReader pdb_reader;
470 E : EXPECT_TRUE(pdb_reader.Read(temp_pdb_, &pdb_file));
471 E : pdb::PdbInfoHeader70 pdb_header = {0};
472 E : pdb::NameStreamMap name_stream_map;
473 : EXPECT_TRUE(ReadHeaderInfoStream(
474 : pdb_file.GetStream(pdb::kPdbHeaderInfoStream),
475 : &pdb_header,
476 E : &name_stream_map));
477 : pdb::NameStreamMap::const_iterator name_it = name_stream_map.find(
478 E : pdb::kSyzygyBlockGraphStreamName);
479 E : ASSERT_TRUE(name_it != name_stream_map.end());
480 E : scoped_refptr<pdb::PdbStream> stream = pdb_file.GetStream(name_it->second);
481 E : ASSERT_TRUE(stream.get() != NULL);
482 E : ASSERT_LT(0U, stream->length());
483 :
484 E : scoped_refptr<pdb::PdbByteStream> byte_stream = new pdb::PdbByteStream();
485 E : EXPECT_TRUE(byte_stream->Init(stream.get()));
486 E : EXPECT_TRUE(byte_stream.get() != NULL);
487 E : core::ScopedInStreamPtr in_stream;
488 : in_stream.reset(core::CreateByteInStream(byte_stream->data(),
489 E : byte_stream->data() + byte_stream->length()));
490 E : core::NativeBinaryInArchive in_archive(in_stream.get());
491 :
492 : // Ensure that the version of the stream is the current one.
493 E : uint32 stream_version = 0;
494 E : EXPECT_TRUE(in_archive.Load(&stream_version));
495 E : ASSERT_EQ(stream_version, pdb::kSyzygyBlockGraphStreamVersion);
496 E : }
497 :
498 : } // namespace pe
|