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