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