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/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 : // Exposes the protected methods for testing.
47 : class TestDecomposer: public Decomposer {
48 : public:
49 E : explicit TestDecomposer(const PEFile& image_file)
50 : : Decomposer(image_file) {
51 E : }
52 :
53 : // Expose as public for testing.
54 : using Decomposer::LoadBlockGraphFromPdbStream;
55 : using Decomposer::GetBlockGraphStreamFromPdb;
56 : };
57 :
58 : class DecomposerTest : public testing::PELibUnitTest {
59 : typedef testing::PELibUnitTest Super;
60 :
61 : public:
62 E : void SetUp() {
63 E : Super::SetUp();
64 :
65 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir_));
66 E : }
67 :
68 : FilePath temp_dir_;
69 : };
70 :
71 : } // namespace
72 :
73 E : TEST_F(DecomposerTest, Decompose) {
74 E : FilePath image_path(testing::GetExeRelativePath(kDllName));
75 E : PEFile image_file;
76 :
77 E : ASSERT_TRUE(image_file.Init(image_path));
78 :
79 : // Decompose the test image and look at the result.
80 E : Decomposer decomposer(image_file);
81 E : EXPECT_TRUE(decomposer.pdb_path().empty());
82 :
83 E : BlockGraph block_graph;
84 E : ImageLayout image_layout(&block_graph);
85 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
86 E : EXPECT_FALSE(decomposer.pdb_path().empty());
87 :
88 : // Retrieve and validate the DOS header.
89 : BlockGraph::Block* dos_header_block =
90 E : image_layout.blocks.GetBlockByAddress(RelativeAddress(0));
91 E : ASSERT_TRUE(dos_header_block != NULL);
92 E : ASSERT_TRUE(IsValidDosHeaderBlock(dos_header_block));
93 :
94 : // Retrieve and validate the NT header.
95 : BlockGraph::Block* nt_headers_block =
96 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
97 E : ASSERT_TRUE(nt_headers_block != NULL);
98 E : ASSERT_TRUE(IsValidNtHeadersBlock(nt_headers_block));
99 :
100 : // There should be some blocks in the graph and in the layout.
101 E : EXPECT_NE(0U, block_graph.blocks().size());
102 E : EXPECT_NE(0U, image_layout.blocks.address_space_impl().size());
103 :
104 : // All the blocks in the graph should be represented in the address space.
105 : EXPECT_EQ(block_graph.blocks().size(),
106 E : image_layout.blocks.address_space_impl().size());
107 :
108 E : ASSERT_EQ(6, image_layout.sections.size());
109 :
110 E : EXPECT_EQ(".text", image_layout.sections[0].name);
111 E : EXPECT_NE(0U, image_layout.sections[0].addr.value());
112 E : EXPECT_NE(0U, image_layout.sections[0].size);
113 E : EXPECT_NE(0U, image_layout.sections[0].data_size);
114 : EXPECT_EQ(IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
115 E : image_layout.sections[0].characteristics);
116 :
117 E : EXPECT_EQ(".rdata", image_layout.sections[1].name);
118 E : EXPECT_NE(0U, image_layout.sections[1].addr.value());
119 E : EXPECT_NE(0U, image_layout.sections[1].size);
120 E : EXPECT_NE(0U, image_layout.sections[1].data_size);
121 : EXPECT_EQ(IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ,
122 E : image_layout.sections[1].characteristics);
123 :
124 E : EXPECT_EQ(".data", image_layout.sections[2].name);
125 E : EXPECT_NE(0U, image_layout.sections[2].addr.value());
126 E : EXPECT_NE(0U, image_layout.sections[2].size);
127 E : EXPECT_NE(0U, image_layout.sections[2].data_size);
128 : EXPECT_EQ(
129 : IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE,
130 E : image_layout.sections[2].characteristics);
131 :
132 E : EXPECT_EQ(".tls", image_layout.sections[3].name);
133 E : EXPECT_NE(0U, image_layout.sections[3].addr.value());
134 E : EXPECT_NE(0U, image_layout.sections[3].size);
135 E : EXPECT_NE(0U, image_layout.sections[3].data_size);
136 : EXPECT_EQ(
137 : IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE,
138 E : image_layout.sections[3].characteristics);
139 :
140 E : EXPECT_EQ(".rsrc", image_layout.sections[4].name);
141 E : EXPECT_NE(0U, image_layout.sections[4].addr.value());
142 E : EXPECT_NE(0U, image_layout.sections[4].size);
143 E : EXPECT_NE(0U, image_layout.sections[4].data_size);
144 : EXPECT_EQ(IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ,
145 E : image_layout.sections[4].characteristics);
146 :
147 E : EXPECT_EQ(".reloc", image_layout.sections[5].name);
148 E : EXPECT_NE(0U, image_layout.sections[5].addr.value());
149 E : EXPECT_NE(0U, image_layout.sections[5].size);
150 E : EXPECT_NE(0U, image_layout.sections[5].data_size);
151 : EXPECT_EQ(IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_DISCARDABLE |
152 E : IMAGE_SCN_MEM_READ, image_layout.sections[5].characteristics);
153 :
154 : // We expect the ImageLayout sections to agree with the BlockGraph sections
155 : // in number, id, name and characteristics.
156 E : EXPECT_EQ(block_graph.sections().size(), image_layout.sections.size());
157 E : for (size_t i = 0; i < image_layout.sections.size(); ++i) {
158 : const BlockGraph::Section* section =
159 E : block_graph.GetSectionById(i);
160 E : ASSERT_TRUE(section != NULL);
161 E : EXPECT_EQ(section->id(), i);
162 E : EXPECT_EQ(section->name(), image_layout.sections[i].name);
163 : EXPECT_EQ(section->characteristics(),
164 E : image_layout.sections[i].characteristics);
165 E : }
166 :
167 : // We expect every block to be associated with a section, and only two blocks
168 : // should not be assigned to a section--the two header blocks.
169 E : size_t non_section_blocks = 0;
170 : BlockGraph::BlockMap::const_iterator it =
171 E : block_graph.blocks().begin();
172 E : for (; it != block_graph.blocks().end(); ++it) {
173 E : const BlockGraph::Block& block = it->second;
174 E : if (block.section() == BlockGraph::kInvalidSectionId) {
175 E : ++non_section_blocks;
176 E : } else {
177 : // If this is not a header block, it should refer to a valid section id.
178 E : EXPECT_LE(0u, block.section());
179 E : EXPECT_LT(block.section(), block_graph.sections().size());
180 : }
181 E : }
182 E : EXPECT_EQ(2u, non_section_blocks);
183 E : }
184 :
185 E : TEST_F(DecomposerTest, DecomposeFailsWithNonexistentPdb) {
186 E : FilePath image_path(testing::GetExeRelativePath(kDllName));
187 E : PEFile image_file;
188 :
189 E : ASSERT_TRUE(image_file.Init(image_path));
190 :
191 E : Decomposer decomposer(image_file);
192 E : decomposer.set_pdb_path(testing::GetExeRelativePath(L"nonexistent.pdb"));
193 :
194 E : BlockGraph block_graph;
195 E : ImageLayout image_layout(&block_graph);
196 E : EXPECT_FALSE(decomposer.Decompose(&image_layout));
197 E : }
198 :
199 E : TEST_F(DecomposerTest, LabelsAndAttributes) {
200 E : FilePath image_path(testing::GetExeRelativePath(kDllName));
201 E : PEFile image_file;
202 :
203 E : ASSERT_TRUE(image_file.Init(image_path));
204 :
205 : // Decompose the test image and look at the result.
206 E : Decomposer decomposer(image_file);
207 E : BlockGraph block_graph;
208 E : ImageLayout image_layout(&block_graph);
209 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
210 :
211 : // Locate various specific function blocks in the block-graph.
212 E : const BlockGraph::Block* dll_main_block = NULL;
213 E : const BlockGraph::Block* func_with_inl_asm_block = NULL;
214 E : const BlockGraph::Block* strchr_block = NULL;
215 : {
216 : BlockGraph::BlockMap::const_iterator it =
217 E : block_graph.blocks().begin();
218 E : for (; it != block_graph.blocks().end(); ++it) {
219 E : if (it->second.name() == "DllMain") {
220 E : ASSERT_TRUE(dll_main_block == NULL);
221 E : dll_main_block = &it->second;
222 E : } else if (it->second.name() == "FunctionWithInlineAssembly") {
223 E : ASSERT_TRUE(func_with_inl_asm_block == NULL);
224 E : func_with_inl_asm_block = &it->second;
225 E : } else if (it->second.name() == "found_bx") {
226 : // This corresponds to the section contribution "strchr.obj".
227 E : ASSERT_TRUE(strchr_block == NULL);
228 E : strchr_block = &it->second;
229 : }
230 E : }
231 E : }
232 :
233 : // Validate that the FunctionWithInlineAssembly block has the appropriate
234 : // attributes.
235 E : ASSERT_FALSE(func_with_inl_asm_block == NULL);
236 : ASSERT_TRUE(func_with_inl_asm_block->attributes() &
237 E : BlockGraph::HAS_INLINE_ASSEMBLY);
238 :
239 : // Validate that the strchr block has the appropriate attributes.
240 E : ASSERT_FALSE(strchr_block == NULL);
241 : ASSERT_TRUE(strchr_block->attributes() &
242 E : BlockGraph::BUILT_BY_UNSUPPORTED_COMPILER);
243 :
244 : // Validate that the DllMain block has the expected population of labels.
245 E : ASSERT_FALSE(dll_main_block == NULL);
246 :
247 : static const size_t kNumDebugEndLabels = (_MSC_VER <= 1500) ? 1 : 0;
248 E : EXPECT_EQ(18 + kNumDebugEndLabels, dll_main_block->labels().size());
249 :
250 E : std::map<BlockGraph::LabelAttributes, size_t> label_attr_counts;
251 : {
252 : BlockGraph::Block::LabelMap::const_iterator it =
253 E : dll_main_block->labels().begin();
254 E : for (; it != dll_main_block->labels().end(); ++it) {
255 E : BlockGraph::LabelAttributes attr_mask = 1;
256 E : for (; attr_mask != BlockGraph::LABEL_ATTRIBUTES_MAX; attr_mask <<= 1) {
257 E : if (it->second.has_attributes(attr_mask))
258 E : label_attr_counts[attr_mask]++;
259 E : }
260 E : }
261 E : }
262 :
263 E : EXPECT_EQ(13, label_attr_counts[BlockGraph::CODE_LABEL]);
264 E : EXPECT_EQ(4, label_attr_counts[BlockGraph::DATA_LABEL]);
265 E : EXPECT_EQ(2, label_attr_counts[BlockGraph::JUMP_TABLE_LABEL]);
266 E : EXPECT_EQ(2, label_attr_counts[BlockGraph::CASE_TABLE_LABEL]);
267 E : EXPECT_EQ(1, label_attr_counts[BlockGraph::DEBUG_START_LABEL]);
268 E : EXPECT_EQ(kNumDebugEndLabels, label_attr_counts[BlockGraph::DEBUG_END_LABEL]);
269 E : }
270 :
271 : namespace {
272 :
273 : void GetNtHeadersBlock(const BlockGraph::Block* dos_header_block,
274 E : BlockGraph::Block** out_nt_headers_block) {
275 E : DCHECK(out_nt_headers_block != NULL);
276 :
277 E : ConstTypedBlock<IMAGE_DOS_HEADER> dos_header;
278 E : ASSERT_TRUE(dos_header.Init(0, dos_header_block));
279 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
280 E : ASSERT_TRUE(dos_header.Dereference(dos_header->e_lfanew, &nt_headers));
281 E : ASSERT_TRUE(nt_headers.block() != NULL);
282 E : *out_nt_headers_block = const_cast<BlockGraph::Block*>(nt_headers.block());
283 E : }
284 :
285 : // This test fixture class contains all the tests that need files generated by
286 : // the relinker (the new image and its corresponding PDB).
287 : class DecomposerAfterRelinkTest : public DecomposerTest {
288 : typedef DecomposerTest Super;
289 :
290 : public:
291 E : virtual void SetUp() OVERRIDE {
292 E : Super::SetUp();
293 E : }
294 :
295 E : void Relink(bool compress_pdb) {
296 : // Initialize a relinker and generate a pdb that contains a block-graph
297 : // stream.
298 E : relinked_dll_ = temp_dir_.Append(kDllName);
299 E : relinked_pdb_ = temp_dir_.Append(kDllPdbName);
300 :
301 E : relinker_.set_input_path(testing::GetExeRelativePath(kDllName));
302 E : relinker_.set_input_pdb_path(testing::GetExeRelativePath(kDllPdbName));
303 E : relinker_.set_allow_overwrite(true);
304 E : relinker_.set_augment_pdb(true);
305 E : relinker_.set_compress_pdb(compress_pdb);
306 E : relinker_.set_output_path(relinked_dll_);
307 E : relinker_.set_output_pdb_path(relinked_pdb_);
308 E : ASSERT_TRUE(relinker_.Init());
309 E : ASSERT_TRUE(relinker_.Relink());
310 E : }
311 :
312 E : void ReconcileNtHeaders(ImageLayout* image_layout) {
313 E : DCHECK(image_layout != NULL);
314 :
315 E : BlockGraph::Block* nt1 = NULL;
316 E : ASSERT_NO_FATAL_FAILURE(GetNtHeadersBlock(relinker_.dos_header_block(),
317 : &nt1));
318 E : ASSERT_TRUE(nt1 != NULL);
319 :
320 : BlockGraph::Block* dos_header_block =
321 E : image_layout->blocks.GetBlockByAddress(core::RelativeAddress(0));
322 E : ASSERT_TRUE(dos_header_block != NULL);
323 E : BlockGraph::Block* nt2 = NULL;
324 E : ASSERT_NO_FATAL_FAILURE(GetNtHeadersBlock(dos_header_block, &nt2));
325 E : ASSERT_TRUE(nt2 != NULL);
326 :
327 : // The NT headers don't compare equal because things like the timestamp and
328 : // checksum are filled out post transform.
329 E : ASSERT_EQ(nt1->data_size(), nt2->data_size());
330 E : nt1->SetData(nt2->data(), nt2->data_size());
331 E : }
332 :
333 E : void LoadRedecompositionData(bool compressed) {
334 E : ASSERT_NO_FATAL_FAILURE(Relink(compressed));
335 :
336 E : PEFile image_file;
337 E : ASSERT_TRUE(image_file.Init(relinked_dll_));
338 :
339 : // Decompose the test image and look at the result.
340 E : Decomposer decomposer(image_file);
341 E : BlockGraph block_graph;
342 E : ImageLayout image_layout(&block_graph);
343 :
344 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
345 :
346 : // Certain data is written to the NT headers post-transform (checksum), so
347 : // it's not reflected in the relinker's block-graph. We reconcile the
348 : // headers prior to doing the comparison.
349 E : ASSERT_NO_FATAL_FAILURE(ReconcileNtHeaders(&image_layout));
350 :
351 : // Ensure that the post-relink block-graph and the deserialized one from the
352 : // PDB are the same.
353 E : block_graph::BlockGraphSerializer bgs;
354 E : ASSERT_TRUE(::testing::BlockGraphsEqual(relinker_.block_graph(),
355 : block_graph,
356 : bgs));
357 E : }
358 :
359 : PERelinker relinker_;
360 : FilePath relinked_dll_;
361 : FilePath relinked_pdb_;
362 : };
363 :
364 : } // namespace
365 :
366 E : TEST_F(DecomposerAfterRelinkTest, LoadRedecompositionDataUncompressed) {
367 E : ASSERT_NO_FATAL_FAILURE(LoadRedecompositionData(false));
368 E : }
369 :
370 E : TEST_F(DecomposerAfterRelinkTest, LoadRedecompositionDataCompressed) {
371 E : ASSERT_NO_FATAL_FAILURE(LoadRedecompositionData(true));
372 E : }
373 :
374 E : TEST_F(DecomposerAfterRelinkTest, FailToLoadBlockGraphWithInvalidVersion) {
375 E : ASSERT_NO_FATAL_FAILURE(Relink(true));
376 :
377 : // Get the block-graph stream from the PDB and change the version of it.
378 :
379 : // Get the stream.
380 E : PEFile image_file;
381 E : ASSERT_TRUE(image_file.Init(relinked_dll_));
382 E : TestDecomposer decomposer(image_file);
383 E : pdb::PdbFile pdb_file;
384 E : pdb::PdbReader pdb_reader;
385 E : pdb_reader.Read(relinked_pdb_, &pdb_file);
386 : scoped_refptr<pdb::PdbStream> block_graph_stream =
387 E : decomposer.GetBlockGraphStreamFromPdb(&pdb_file);
388 :
389 : // Create a copy of the stream. We need to do this to have a stream that we
390 : // can modify.
391 E : scoped_refptr<pdb::PdbByteStream> new_stream = new pdb::PdbByteStream();
392 E : ASSERT_TRUE(new_stream->Init(block_graph_stream.get()));
393 E : block_graph_stream = new_stream.get();
394 : scoped_refptr<pdb::WritablePdbStream> block_graph_writer =
395 E : block_graph_stream->GetWritablePdbStream();
396 E : ASSERT_TRUE(block_graph_writer.get() != NULL);
397 :
398 : // Change the version of the stream.
399 E : block_graph_writer->set_pos(0);
400 E : block_graph_writer->Write(pdb::kSyzygyBlockGraphStreamVersion + 1);
401 :
402 E : BlockGraph block_graph;
403 E : ImageLayout image_layout(&block_graph);
404 : // We've invalided the version previously so this test should fail.
405 : ASSERT_FALSE(decomposer.LoadBlockGraphFromPdbStream(image_file,
406 : block_graph_stream.get(),
407 E : &image_layout));
408 E : }
409 :
410 : } // namespace pe
|