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/decomposer.h"
16 :
17 : #include <set>
18 :
19 : #include "base/file_util.h"
20 : #include "base/path_service.h"
21 : #include "base/string_util.h"
22 : #include "gmock/gmock.h"
23 : #include "gtest/gtest.h"
24 : #include "syzygy/block_graph/block_graph_serializer.h"
25 : #include "syzygy/block_graph/typed_block.h"
26 : #include "syzygy/block_graph/unittest_util.h"
27 : #include "syzygy/core/unittest_util.h"
28 : #include "syzygy/pdb/pdb_byte_stream.h"
29 : #include "syzygy/pdb/pdb_file.h"
30 : #include "syzygy/pdb/pdb_reader.h"
31 : #include "syzygy/pdb/pdb_stream.h"
32 : #include "syzygy/pdb/pdb_util.h"
33 : #include "syzygy/pdb/pdb_writer.h"
34 : #include "syzygy/pe/pe_relinker.h"
35 : #include "syzygy/pe/pe_utils.h"
36 : #include "syzygy/pe/unittest_util.h"
37 :
38 : namespace pe {
39 :
40 : using block_graph::BlockGraph;
41 : using block_graph::ConstTypedBlock;
42 : using core::RelativeAddress;
43 :
44 : namespace {
45 :
46 : const size_t kPointerSize = BlockGraph::Reference::kMaximumSize;
47 :
48 : // Exposes the protected methods for testing.
49 : class TestDecomposer: public Decomposer {
50 : public:
51 E : explicit TestDecomposer(const PEFile& image_file)
52 : : Decomposer(image_file) {
53 E : }
54 :
55 : // Expose as public for testing.
56 : using Decomposer::LoadBlockGraphFromPdbStream;
57 : using Decomposer::GetBlockGraphStreamFromPdb;
58 : };
59 :
60 : class DecomposerTest : public testing::PELibUnitTest {
61 : typedef testing::PELibUnitTest Super;
62 :
63 : public:
64 E : void SetUp() {
65 E : Super::SetUp();
66 :
67 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir_));
68 E : }
69 :
70 : base::FilePath temp_dir_;
71 : };
72 :
73 : } // namespace
74 :
75 E : TEST_F(DecomposerTest, Decompose) {
76 E : base::FilePath image_path(testing::GetExeRelativePath(testing::kTestDllName));
77 E : PEFile image_file;
78 :
79 E : ASSERT_TRUE(image_file.Init(image_path));
80 :
81 : // Decompose the test image and look at the result.
82 E : Decomposer decomposer(image_file);
83 E : EXPECT_TRUE(decomposer.pdb_path().empty());
84 :
85 E : BlockGraph block_graph;
86 E : ImageLayout image_layout(&block_graph);
87 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
88 E : EXPECT_FALSE(decomposer.pdb_path().empty());
89 :
90 : // Retrieve and validate the DOS header.
91 : BlockGraph::Block* dos_header_block =
92 E : image_layout.blocks.GetBlockByAddress(RelativeAddress(0));
93 E : ASSERT_TRUE(dos_header_block != NULL);
94 E : ASSERT_TRUE(IsValidDosHeaderBlock(dos_header_block));
95 :
96 : // Retrieve and validate the NT header.
97 : BlockGraph::Block* nt_headers_block =
98 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
99 E : ASSERT_TRUE(nt_headers_block != NULL);
100 E : ASSERT_TRUE(IsValidNtHeadersBlock(nt_headers_block));
101 :
102 : // There should be some blocks in the graph and in the layout.
103 E : EXPECT_NE(0U, block_graph.blocks().size());
104 E : EXPECT_NE(0U, image_layout.blocks.address_space_impl().size());
105 :
106 : // All the blocks in the graph should be represented in the address space.
107 : EXPECT_EQ(block_graph.blocks().size(),
108 E : image_layout.blocks.address_space_impl().size());
109 :
110 E : ASSERT_EQ(6, image_layout.sections.size());
111 :
112 E : EXPECT_EQ(".text", image_layout.sections[0].name);
113 E : EXPECT_NE(0U, image_layout.sections[0].addr.value());
114 E : EXPECT_NE(0U, image_layout.sections[0].size);
115 E : EXPECT_NE(0U, image_layout.sections[0].data_size);
116 : EXPECT_EQ(IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
117 E : image_layout.sections[0].characteristics);
118 :
119 E : EXPECT_EQ(".rdata", image_layout.sections[1].name);
120 E : EXPECT_NE(0U, image_layout.sections[1].addr.value());
121 E : EXPECT_NE(0U, image_layout.sections[1].size);
122 E : EXPECT_NE(0U, image_layout.sections[1].data_size);
123 : EXPECT_EQ(IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ,
124 E : image_layout.sections[1].characteristics);
125 :
126 E : EXPECT_EQ(".data", image_layout.sections[2].name);
127 E : EXPECT_NE(0U, image_layout.sections[2].addr.value());
128 E : EXPECT_NE(0U, image_layout.sections[2].size);
129 E : EXPECT_NE(0U, image_layout.sections[2].data_size);
130 : EXPECT_EQ(
131 : IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE,
132 E : image_layout.sections[2].characteristics);
133 :
134 E : EXPECT_EQ(".tls", image_layout.sections[3].name);
135 E : EXPECT_NE(0U, image_layout.sections[3].addr.value());
136 E : EXPECT_NE(0U, image_layout.sections[3].size);
137 E : EXPECT_NE(0U, image_layout.sections[3].data_size);
138 : EXPECT_EQ(
139 : IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE,
140 E : image_layout.sections[3].characteristics);
141 :
142 E : EXPECT_EQ(".rsrc", image_layout.sections[4].name);
143 E : EXPECT_NE(0U, image_layout.sections[4].addr.value());
144 E : EXPECT_NE(0U, image_layout.sections[4].size);
145 E : EXPECT_NE(0U, image_layout.sections[4].data_size);
146 : EXPECT_EQ(IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ,
147 E : image_layout.sections[4].characteristics);
148 :
149 E : EXPECT_EQ(".reloc", image_layout.sections[5].name);
150 E : EXPECT_NE(0U, image_layout.sections[5].addr.value());
151 E : EXPECT_NE(0U, image_layout.sections[5].size);
152 E : EXPECT_NE(0U, image_layout.sections[5].data_size);
153 : EXPECT_EQ(IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_DISCARDABLE |
154 E : IMAGE_SCN_MEM_READ, image_layout.sections[5].characteristics);
155 :
156 : // We expect the ImageLayout sections to agree with the BlockGraph sections
157 : // in number, id, name and characteristics.
158 E : EXPECT_EQ(block_graph.sections().size(), image_layout.sections.size());
159 E : for (size_t i = 0; i < image_layout.sections.size(); ++i) {
160 : const BlockGraph::Section* section =
161 E : block_graph.GetSectionById(i);
162 E : ASSERT_TRUE(section != NULL);
163 E : EXPECT_EQ(section->id(), i);
164 E : EXPECT_EQ(section->name(), image_layout.sections[i].name);
165 : EXPECT_EQ(section->characteristics(),
166 E : image_layout.sections[i].characteristics);
167 E : }
168 :
169 : // We expect every block to be associated with a section, and only two blocks
170 : // should not be assigned to a section--the two header blocks.
171 E : size_t non_section_blocks = 0;
172 : BlockGraph::BlockMap::const_iterator it =
173 E : block_graph.blocks().begin();
174 E : for (; it != block_graph.blocks().end(); ++it) {
175 E : const BlockGraph::Block& block = it->second;
176 E : if (block.section() == BlockGraph::kInvalidSectionId) {
177 E : ++non_section_blocks;
178 E : } else {
179 : // If this is not a header block, it should refer to a valid section id.
180 E : EXPECT_LE(0u, block.section());
181 E : EXPECT_LT(block.section(), block_graph.sections().size());
182 : }
183 E : }
184 E : EXPECT_EQ(2u, non_section_blocks);
185 :
186 : // Make sure that all bracketed COFF groups have been parsed. There are 8
187 : // of them that we currently know of:
188 : // .CRT$XCA -> .CRT$XCZ: C initializers
189 : // .CRT$XIA -> .CRT$XLZ: C++ initializers
190 : // .CRT$XLA -> .CRT$XLZ: TLS callbacks
191 : // .CRT$XPA -> .CRT$XPZ: CRT pre-termination functions.
192 : // .CRT$XTA -> .CRT$XTZ: CRT termination functions.
193 : // .rtc$IAA -> .rtc$IZZ: Run-time checking initializers.
194 : // .rtc$TAA -> .rtc$TZZ: Run-time checking termination functions.
195 : // .tls -> .tls$ZZZ: TLS data.
196 E : size_t coff_group_blocks = 0;
197 E : it = block_graph.blocks().begin();
198 E : for (; it != block_graph.blocks().end(); ++it) {
199 E : const BlockGraph::Block& block = it->second;
200 E : if (block.attributes() & BlockGraph::COFF_GROUP)
201 E : ++coff_group_blocks;
202 E : }
203 E : EXPECT_EQ(8u, coff_group_blocks);
204 E : }
205 :
206 E : TEST_F(DecomposerTest, DecomposeFailsWithNonexistentPdb) {
207 E : base::FilePath image_path(testing::GetExeRelativePath(testing::kTestDllName));
208 E : PEFile image_file;
209 :
210 E : ASSERT_TRUE(image_file.Init(image_path));
211 :
212 E : Decomposer decomposer(image_file);
213 E : decomposer.set_pdb_path(testing::GetExeRelativePath(L"nonexistent.pdb"));
214 :
215 E : BlockGraph block_graph;
216 E : ImageLayout image_layout(&block_graph);
217 E : EXPECT_FALSE(decomposer.Decompose(&image_layout));
218 E : }
219 :
220 E : TEST_F(DecomposerTest, LabelsAndAttributes) {
221 E : base::FilePath image_path(testing::GetExeRelativePath(testing::kTestDllName));
222 E : PEFile image_file;
223 :
224 E : ASSERT_TRUE(image_file.Init(image_path));
225 :
226 : // Decompose the test image and look at the result.
227 E : Decomposer decomposer(image_file);
228 E : BlockGraph block_graph;
229 E : ImageLayout image_layout(&block_graph);
230 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
231 :
232 : // Locate various specific function blocks in the block-graph.
233 E : const BlockGraph::Block* dll_main_block = NULL;
234 E : const BlockGraph::Block* func_with_inl_asm_block = NULL;
235 E : const BlockGraph::Block* strchr_block = NULL;
236 E : const BlockGraph::Block* imp_load_block = NULL;
237 E : const BlockGraph::Block* no_private_symbols_block = NULL;
238 :
239 : {
240 : typedef std::map<std::string, const BlockGraph::Block**> TestBlockMap;
241 :
242 E : TestBlockMap test_blocks;
243 E : test_blocks.insert(std::make_pair("DllMain", &dll_main_block));
244 : test_blocks.insert(std::make_pair("FunctionWithInlineAssembly",
245 E : &func_with_inl_asm_block));
246 E : test_blocks.insert(std::make_pair("found_bx", &strchr_block));
247 : test_blocks.insert(std::make_pair("__imp_load_CoCreateGuid",
248 E : &imp_load_block));
249 : test_blocks.insert(std::make_pair("TestFunctionWithNoPrivateSymbols",
250 E : &no_private_symbols_block));
251 :
252 E : BlockGraph::BlockMap::const_iterator it = block_graph.blocks().begin();
253 E : for (; it != block_graph.blocks().end(); ++it) {
254 E : const BlockGraph::Block& block = it->second;
255 :
256 E : TestBlockMap::const_iterator test_it = test_blocks.find(block.name());
257 E : if (test_it == test_blocks.end())
258 E : continue;
259 :
260 E : ASSERT_TRUE(*test_it->second == NULL);
261 E : *test_it->second = █
262 E : }
263 E : }
264 :
265 : // The block with no private symbols should be marked as ERRORED_DISASSEMBLY,
266 : // and only have a single public symbol label.
267 E : ASSERT_FALSE(no_private_symbols_block == NULL);
268 : EXPECT_TRUE(no_private_symbols_block->attributes() &
269 E : BlockGraph::ERRORED_DISASSEMBLY);
270 E : EXPECT_EQ(1u, no_private_symbols_block->labels().size());
271 : BlockGraph::Block::LabelMap::const_iterator label_it =
272 E : no_private_symbols_block->labels().begin();
273 E : EXPECT_EQ(0, label_it->first);
274 E : EXPECT_EQ(BlockGraph::PUBLIC_SYMBOL_LABEL, label_it->second.attributes());
275 :
276 : // The __imp_load__ block should be a thunk.
277 E : ASSERT_FALSE(imp_load_block == NULL);
278 E : EXPECT_NE(0UL, imp_load_block->attributes() & BlockGraph::THUNK);
279 :
280 : // Validate that the FunctionWithInlineAssembly block has the appropriate
281 : // attributes.
282 E : ASSERT_FALSE(func_with_inl_asm_block == NULL);
283 : ASSERT_TRUE(func_with_inl_asm_block->attributes() &
284 E : BlockGraph::HAS_INLINE_ASSEMBLY);
285 :
286 : // Validate that the strchr block has the appropriate attributes.
287 E : ASSERT_FALSE(strchr_block == NULL);
288 : ASSERT_TRUE(strchr_block->attributes() &
289 E : BlockGraph::BUILT_BY_UNSUPPORTED_COMPILER);
290 :
291 : #ifdef OFFICIAL_BUILD
292 : static const size_t kDllMainLabelCount = 41;
293 : static const size_t kCallSiteLabelCount = 26;
294 : #else
295 : static const size_t kDllMainLabelCount = 31;
296 : static const size_t kCallSiteLabelCount = 10;
297 : #endif
298 :
299 : // Validate that the DllMain block has the expected population of labels.
300 E : ASSERT_FALSE(dll_main_block == NULL);
301 E : EXPECT_EQ(kDllMainLabelCount, dll_main_block->labels().size());
302 :
303 : // DllMain has a jump table so it should have pointer alignment.
304 E : ASSERT_EQ(kPointerSize, dll_main_block->alignment());
305 :
306 E : std::map<BlockGraph::LabelAttributes, size_t> label_attr_counts;
307 : {
308 : BlockGraph::Block::LabelMap::const_iterator it =
309 E : dll_main_block->labels().begin();
310 E : for (; it != dll_main_block->labels().end(); ++it) {
311 E : BlockGraph::LabelAttributes attr_mask = 1;
312 E : for (; attr_mask != BlockGraph::LABEL_ATTRIBUTES_MAX; attr_mask <<= 1) {
313 E : if (it->second.has_attributes(attr_mask))
314 E : label_attr_counts[attr_mask]++;
315 E : }
316 E : }
317 : }
318 :
319 E : EXPECT_EQ(19, label_attr_counts[BlockGraph::CODE_LABEL]);
320 : EXPECT_EQ(kCallSiteLabelCount,
321 E : label_attr_counts[BlockGraph::CALL_SITE_LABEL]);
322 E : EXPECT_EQ(5, label_attr_counts[BlockGraph::DATA_LABEL]);
323 E : EXPECT_EQ(3, label_attr_counts[BlockGraph::JUMP_TABLE_LABEL]);
324 E : EXPECT_EQ(2, label_attr_counts[BlockGraph::CASE_TABLE_LABEL]);
325 E : EXPECT_EQ(1, label_attr_counts[BlockGraph::DEBUG_START_LABEL]);
326 E : }
327 :
328 : namespace {
329 :
330 : void GetNtHeadersBlock(const BlockGraph::Block* dos_header_block,
331 E : BlockGraph::Block** out_nt_headers_block) {
332 E : DCHECK(out_nt_headers_block != NULL);
333 :
334 E : ConstTypedBlock<IMAGE_DOS_HEADER> dos_header;
335 E : ASSERT_TRUE(dos_header.Init(0, dos_header_block));
336 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
337 E : ASSERT_TRUE(dos_header.Dereference(dos_header->e_lfanew, &nt_headers));
338 E : ASSERT_TRUE(nt_headers.block() != NULL);
339 E : *out_nt_headers_block = const_cast<BlockGraph::Block*>(nt_headers.block());
340 E : }
341 :
342 : // This test fixture class contains all the tests that need files generated by
343 : // the relinker (the new image and its corresponding PDB).
344 : class DecomposerAfterRelinkTest : public DecomposerTest {
345 : typedef DecomposerTest Super;
346 :
347 : public:
348 E : virtual void SetUp() OVERRIDE {
349 E : Super::SetUp();
350 E : }
351 :
352 E : void Relink(bool compress_pdb) {
353 : // Initialize a relinker and generate a pdb that contains a block-graph
354 : // stream.
355 E : relinked_dll_ = temp_dir_.Append(testing::kTestDllName);
356 E : relinked_pdb_ = temp_dir_.Append(testing::kTestDllPdbName);
357 :
358 : relinker_.set_input_path(testing::GetExeRelativePath(
359 E : testing::kTestDllName));
360 : relinker_.set_input_pdb_path(testing::GetExeRelativePath(
361 E : testing::kTestDllPdbName));
362 E : relinker_.set_allow_overwrite(true);
363 E : relinker_.set_augment_pdb(true);
364 E : relinker_.set_compress_pdb(compress_pdb);
365 E : relinker_.set_output_path(relinked_dll_);
366 E : relinker_.set_output_pdb_path(relinked_pdb_);
367 E : ASSERT_TRUE(relinker_.Init());
368 E : ASSERT_TRUE(relinker_.Relink());
369 E : }
370 :
371 E : void ReconcileNtHeaders(ImageLayout* image_layout) {
372 E : DCHECK(image_layout != NULL);
373 :
374 E : BlockGraph::Block* nt1 = NULL;
375 E : ASSERT_NO_FATAL_FAILURE(GetNtHeadersBlock(relinker_.dos_header_block(),
376 : &nt1));
377 E : ASSERT_TRUE(nt1 != NULL);
378 :
379 : BlockGraph::Block* dos_header_block =
380 E : image_layout->blocks.GetBlockByAddress(core::RelativeAddress(0));
381 E : ASSERT_TRUE(dos_header_block != NULL);
382 E : BlockGraph::Block* nt2 = NULL;
383 E : ASSERT_NO_FATAL_FAILURE(GetNtHeadersBlock(dos_header_block, &nt2));
384 E : ASSERT_TRUE(nt2 != NULL);
385 :
386 : // The NT headers don't compare equal because things like the timestamp and
387 : // checksum are filled out post transform.
388 E : ASSERT_EQ(nt1->data_size(), nt2->data_size());
389 E : nt1->SetData(nt2->data(), nt2->data_size());
390 E : }
391 :
392 E : void LoadRedecompositionData(bool compressed) {
393 E : ASSERT_NO_FATAL_FAILURE(Relink(compressed));
394 :
395 E : PEFile image_file;
396 E : ASSERT_TRUE(image_file.Init(relinked_dll_));
397 :
398 : // Decompose the test image and look at the result.
399 E : Decomposer decomposer(image_file);
400 E : BlockGraph block_graph;
401 E : ImageLayout image_layout(&block_graph);
402 :
403 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
404 :
405 : // Certain data is written to the NT headers post-transform (checksum), so
406 : // it's not reflected in the relinker's block-graph. We reconcile the
407 : // headers prior to doing the comparison.
408 E : ASSERT_NO_FATAL_FAILURE(ReconcileNtHeaders(&image_layout));
409 :
410 : // Ensure that the post-relink block-graph and the deserialized one from the
411 : // PDB are the same.
412 E : block_graph::BlockGraphSerializer bgs;
413 E : ASSERT_TRUE(::testing::BlockGraphsEqual(relinker_.block_graph(),
414 : block_graph,
415 : bgs));
416 E : }
417 :
418 : PERelinker relinker_;
419 : base::FilePath relinked_dll_;
420 : base::FilePath relinked_pdb_;
421 : };
422 :
423 : } // namespace
424 :
425 E : TEST_F(DecomposerAfterRelinkTest, LoadRedecompositionDataUncompressed) {
426 E : ASSERT_NO_FATAL_FAILURE(LoadRedecompositionData(false));
427 E : }
428 :
429 E : TEST_F(DecomposerAfterRelinkTest, LoadRedecompositionDataCompressed) {
430 E : ASSERT_NO_FATAL_FAILURE(LoadRedecompositionData(true));
431 E : }
432 :
433 E : TEST_F(DecomposerAfterRelinkTest, FailToLoadBlockGraphWithInvalidVersion) {
434 E : ASSERT_NO_FATAL_FAILURE(Relink(true));
435 :
436 : // Get the block-graph stream from the PDB and change the version of it.
437 :
438 : // Get the stream.
439 E : PEFile image_file;
440 E : ASSERT_TRUE(image_file.Init(relinked_dll_));
441 E : TestDecomposer decomposer(image_file);
442 E : pdb::PdbFile pdb_file;
443 E : pdb::PdbReader pdb_reader;
444 E : pdb_reader.Read(relinked_pdb_, &pdb_file);
445 : scoped_refptr<pdb::PdbStream> block_graph_stream =
446 E : decomposer.GetBlockGraphStreamFromPdb(&pdb_file);
447 :
448 : // Create a copy of the stream. We need to do this to have a stream that we
449 : // can modify.
450 E : scoped_refptr<pdb::PdbByteStream> new_stream = new pdb::PdbByteStream();
451 E : ASSERT_TRUE(new_stream->Init(block_graph_stream.get()));
452 E : block_graph_stream = new_stream.get();
453 : scoped_refptr<pdb::WritablePdbStream> block_graph_writer =
454 E : block_graph_stream->GetWritablePdbStream();
455 E : ASSERT_TRUE(block_graph_writer.get() != NULL);
456 :
457 : // Change the version of the stream.
458 E : block_graph_writer->set_pos(0);
459 E : block_graph_writer->Write(pdb::kSyzygyBlockGraphStreamVersion + 1);
460 :
461 E : BlockGraph block_graph;
462 E : ImageLayout image_layout(&block_graph);
463 : // We've invalided the version previously so this test should fail.
464 : ASSERT_FALSE(decomposer.LoadBlockGraphFromPdbStream(image_file,
465 : block_graph_stream.get(),
466 E : &image_layout));
467 E : }
468 :
469 : } // namespace pe
|