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