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 : // Unittests for pe::orderers::PEOrderer.
16 :
17 : #include "syzygy/pe/orderers/pe_orderer.h"
18 :
19 : #include <algorithm>
20 :
21 : #include "gmock/gmock.h"
22 : #include "gtest/gtest.h"
23 : #include "syzygy/block_graph/typed_block.h"
24 : #include "syzygy/core/unittest_util.h"
25 : #include "syzygy/pe/decomposer.h"
26 : #include "syzygy/pe/image_layout_builder.h"
27 : #include "syzygy/pe/pe_file_writer.h"
28 : #include "syzygy/pe/pe_utils.h"
29 : #include "syzygy/pe/unittest_util.h"
30 :
31 : namespace pe {
32 : namespace orderers {
33 :
34 : namespace {
35 :
36 : using block_graph::BlockGraph;
37 : using block_graph::BlockVector;
38 : using block_graph::ConstBlockVector;
39 : using block_graph::ConstTypedBlock;
40 : using block_graph::OrderedBlockGraph;
41 : using block_graph::TypedBlock;
42 : using core::RelativeAddress;
43 :
44 : const BlockGraph::Section* FindSection(const BlockGraph* block_graph,
45 E : const char* section_name) {
46 E : DCHECK(block_graph != NULL);
47 E : DCHECK(section_name != NULL);
48 :
49 : BlockGraph::SectionMap::const_iterator section_it =
50 E : block_graph->sections().begin();
51 E : for (; section_it != block_graph->sections().end(); ++section_it) {
52 E : if (section_it->second.name() == section_name)
53 E : return §ion_it->second;
54 E : }
55 :
56 i : return NULL;
57 E : }
58 :
59 : const BlockGraph::Block* GetDataDirEntryBlock(
60 : const ConstTypedBlock<IMAGE_NT_HEADERS>& nt_headers,
61 E : size_t data_dir_index) {
62 E : ConstTypedBlock<char> data_dir;
63 : if (!nt_headers.Dereference(
64 : nt_headers->OptionalHeader.DataDirectory[data_dir_index],
65 E : &data_dir)) {
66 E : return NULL;
67 : }
68 : // We expect the block to be at zero offset.
69 i : CHECK_EQ(0, data_dir.offset());
70 i : return data_dir.block();
71 E : }
72 :
73 : // TODO(chrisha): Move all of the PE-layout validation code to somewhere public
74 : // in pe_lib or pe_unittest_utils_lib. It is useful elsewhere. Similarly,
75 : // the routines for generating valid PE structures should live in
76 : // pe_unittest_utils_lib.
77 :
78 : void VerifySectionStartsWith(
79 : const OrderedBlockGraph* obg,
80 : const BlockGraph::Section* section,
81 E : const ConstBlockVector& blocks) {
82 E : ASSERT_TRUE(obg != NULL);
83 :
84 : const OrderedBlockGraph::OrderedSection& ordered_section =
85 E : obg->ordered_section(section);
86 :
87 : const OrderedBlockGraph::BlockList& block_list(
88 E : ordered_section.ordered_blocks());
89 : OrderedBlockGraph::BlockList::const_iterator block_it =
90 E : block_list.begin();
91 E : size_t i = 0;
92 E : for (; i < blocks.size() && block_it != block_list.end(); ++block_it, ++i)
93 E : ASSERT_EQ(*block_it, blocks[i]);
94 :
95 E : ASSERT_EQ(i, blocks.size());
96 E : }
97 :
98 : // Verifies that the generated layout is valid.
99 : void VerifyValidLayout(const OrderedBlockGraph* obg,
100 E : const BlockGraph::Block* dos_header_block) {
101 E : ASSERT_TRUE(obg != NULL);
102 E : ASSERT_TRUE(dos_header_block != NULL);
103 :
104 E : ConstTypedBlock<IMAGE_DOS_HEADER> dos_header;
105 E : ASSERT_TRUE(dos_header.Init(0, dos_header_block));
106 :
107 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
108 E : ASSERT_TRUE(dos_header.Dereference(dos_header->e_lfanew, &nt_headers));
109 :
110 : // Ensure the headers are in the right order.
111 E : ConstBlockVector header_blocks;
112 E : header_blocks.push_back(dos_header.block());
113 E : header_blocks.push_back(nt_headers.block());
114 E : ASSERT_NO_FATAL_FAILURE(VerifySectionStartsWith(obg, NULL, header_blocks));
115 :
116 : // Check that the resources are in the right section.
117 : const BlockGraph::Section* rsrc_section =
118 E : FindSection(obg->block_graph(), kResourceSectionName);
119 : const BlockGraph::Block* rsrc_data_dir =
120 E : GetDataDirEntryBlock(nt_headers, IMAGE_DIRECTORY_ENTRY_RESOURCE);
121 :
122 E : if (rsrc_data_dir != NULL) {
123 i : ASSERT_TRUE(rsrc_section != NULL);
124 i : ASSERT_EQ(rsrc_section->id(), rsrc_data_dir->section());
125 : }
126 :
127 : // Check that the relocs are in the right section.
128 : const BlockGraph::Section* reloc_section =
129 E : FindSection(obg->block_graph(), kRelocSectionName);
130 : const BlockGraph::Block* reloc_data_dir =
131 E : GetDataDirEntryBlock(nt_headers, IMAGE_DIRECTORY_ENTRY_BASERELOC);
132 :
133 E : if (reloc_data_dir != NULL) {
134 i : ASSERT_TRUE(reloc_section != NULL);
135 i : ASSERT_EQ(reloc_section->id(), reloc_data_dir->section());
136 : }
137 :
138 : // Ensure that .rsrc and .reloc are the last two sections.
139 : OrderedBlockGraph::SectionList::const_iterator section_it =
140 E : obg->ordered_sections().end();
141 E : --section_it;
142 E : if (reloc_section != NULL) {
143 E : ASSERT_EQ((*section_it)->section(), reloc_section);
144 E : --section_it;
145 : }
146 E : if (rsrc_section != NULL) {
147 E : ASSERT_EQ((*section_it)->section(), rsrc_section);
148 : }
149 E : }
150 :
151 : class PEOrdererTest : public testing::PELibUnitTest {
152 : public:
153 E : PEOrdererTest() : dos_header_block_(NULL) { }
154 :
155 E : void InitOrderedBlockGraph() {
156 E : if (ordered_block_graph_.get() == NULL)
157 E : ordered_block_graph_.reset(new OrderedBlockGraph(&block_graph_));
158 E : }
159 :
160 E : void RandomizeOrderedBlockGraph() {
161 E : InitOrderedBlockGraph();
162 :
163 : // We randomize *everything*, including the orders of sections and to
164 : // which sections blocks belong. This is more general than
165 : // block_graph::orderers::RandomOrderer.
166 :
167 E : std::vector<BlockGraph::Section*> sections;
168 : BlockGraph::SectionMap::iterator section_it =
169 E : block_graph_.sections_mutable().begin();
170 E : for (; section_it != block_graph_.sections_mutable().end(); ++section_it)
171 E : sections.push_back(§ion_it->second);
172 :
173 E : BlockVector blocks;
174 : BlockGraph::BlockMap::iterator block_it =
175 E : block_graph_.blocks_mutable().begin();
176 E : for (; block_it != block_graph_.blocks_mutable().end(); ++block_it)
177 E : blocks.push_back(&block_it->second);
178 :
179 E : std::random_shuffle(sections.begin(), sections.end());
180 E : std::random_shuffle(blocks.begin(), blocks.end());
181 :
182 E : for (size_t i = 0; i < sections.size(); ++i)
183 E : ordered_block_graph_->PlaceAtTail(sections[i]);
184 :
185 E : for (size_t i = 0; i < blocks.size(); ++i) {
186 : // We randomly place some blocks in the 'header' section as well.
187 E : size_t j = rand() % (sections.size() + 1);
188 : ordered_block_graph_->PlaceAtTail(
189 : j == sections.size() ? NULL : sections[j],
190 E : blocks[i]);
191 E : }
192 E : }
193 :
194 : // This generates a dummy image with all of the PE features we wish to test,
195 : // but it will not result in a loadable/runnable module if written. It is
196 : // significantly more lightweight than test_dll, however.
197 E : void GenerateDummyImage() {
198 : // Create the standard assortment of sections.
199 : BlockGraph::Section* text = block_graph_.AddSection(
200 E : kCodeSectionName, kCodeCharacteristics);
201 : BlockGraph::Section* rdata = block_graph_.AddSection(
202 E : kReadOnlyDataSectionName, kReadOnlyDataCharacteristics);
203 : BlockGraph::Section* data = block_graph_.AddSection(
204 E : kReadWriteDataSectionName, kReadWriteDataCharacteristics);
205 : BlockGraph::Section* rsrc = block_graph_.AddSection(
206 E : kResourceSectionName, kReadOnlyDataCharacteristics);
207 : BlockGraph::Section* reloc = block_graph_.AddSection(
208 E : kRelocSectionName, kRelocCharacteristics);
209 :
210 : // Create one dummy block per section. This just ensures they have something
211 : // in them.
212 : BlockGraph::Block* text_block = AddBlock(
213 E : BlockGraph::CODE_BLOCK, 16, "text block", text);
214 :
215 : BlockGraph::Block* rdata_block = AddBlock(
216 E : BlockGraph::DATA_BLOCK, 16, "rdata block", rdata);
217 :
218 : BlockGraph::Block* data_block = AddBlock(
219 E : BlockGraph::DATA_BLOCK, 16, "data block", data);
220 :
221 : BlockGraph::Block* rsrc_block = AddBlock(BlockGraph::DATA_BLOCK,
222 E : sizeof(IMAGE_RESOURCE_DIRECTORY), "rsrc block", rsrc);
223 :
224 : BlockGraph::Block* reloc_block = AddBlock(BlockGraph::DATA_BLOCK,
225 E : sizeof(IMAGE_BASE_RELOCATION), "reloc block", reloc);
226 :
227 : // Create and initialize the headers.
228 : dos_header_block_ = block_graph_.AddBlock(
229 E : BlockGraph::DATA_BLOCK, sizeof(IMAGE_DOS_HEADER), "Dos Headers");
230 E : dos_header_block_->ResizeData(dos_header_block_->size());
231 : BlockGraph::Block* nt_headers_block = block_graph_.AddBlock(
232 : BlockGraph::DATA_BLOCK,
233 : sizeof(IMAGE_NT_HEADERS) +
234 : block_graph_.sections().size() * sizeof(IMAGE_SECTION_HEADER),
235 E : "Nt Headers");
236 E : nt_headers_block->ResizeData(nt_headers_block->size());
237 :
238 E : size_t header_size = dos_header_block_->size() + nt_headers_block->size();
239 :
240 E : TypedBlock<IMAGE_DOS_HEADER> dos_header;
241 E : ASSERT_TRUE(dos_header.Init(0, dos_header_block_));
242 : dos_header.SetReference(BlockGraph::RELATIVE_REF,
243 : dos_header->e_lfanew,
244 : nt_headers_block,
245 E : 0, 0);
246 E : ASSERT_TRUE(UpdateDosHeader(dos_header_block_));
247 :
248 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
249 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block));
250 E : nt_headers->FileHeader.NumberOfSections = block_graph_.sections().size();
251 E : nt_headers->FileHeader.SizeOfOptionalHeader = sizeof(IMAGE_OPTIONAL_HEADER);
252 E : nt_headers->OptionalHeader.Magic = IMAGE_NT_OPTIONAL_HDR_MAGIC;
253 E : nt_headers->Signature = IMAGE_NT_SIGNATURE;
254 :
255 : // Set up the relocs data directory.
256 E : TypedBlock<IMAGE_DATA_DIRECTORY> data_dir;
257 E : ASSERT_TRUE(data_dir.Init(nt_headers.OffsetOf(
258 : nt_headers->OptionalHeader.DataDirectory[
259 : IMAGE_DIRECTORY_ENTRY_BASERELOC]),
260 : nt_headers.block()));
261 : data_dir.SetReference(BlockGraph::RELATIVE_REF,
262 : data_dir->VirtualAddress,
263 : reloc_block,
264 E : 0, 0);
265 E : data_dir->Size = reloc_block->size();
266 :
267 : // Set up the resources data directory.
268 E : ASSERT_TRUE(data_dir.Init(nt_headers.OffsetOf(
269 : nt_headers->OptionalHeader.DataDirectory[
270 : IMAGE_DIRECTORY_ENTRY_RESOURCE]),
271 : nt_headers.block()));
272 : data_dir.SetReference(BlockGraph::RELATIVE_REF,
273 : data_dir->VirtualAddress,
274 : rsrc_block,
275 E : 0, 0);
276 E : data_dir->Size = rsrc_block->size();
277 :
278 E : ASSERT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
279 E : ASSERT_TRUE(IsValidNtHeadersBlock(nt_headers_block));
280 E : }
281 :
282 E : void DecomposeTestDll() {
283 : base::FilePath image_path(
284 E : testing::GetExeRelativePath(testing::kTestDllName));
285 :
286 E : ASSERT_TRUE(pe_file_.Init(image_path));
287 :
288 : // Decompose the test image and look at the result.
289 E : ImageLayout image_layout(&block_graph_);
290 E : Decomposer decomposer(pe_file_);
291 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
292 :
293 : // Retrieve and validate the DOS header.
294 : dos_header_block_ =
295 E : image_layout.blocks.GetBlockByAddress(RelativeAddress(0));
296 E : ASSERT_TRUE(dos_header_block_ != NULL);
297 E : ASSERT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
298 E : }
299 :
300 : BlockGraph::Block* AddBlock(BlockGraph::BlockType type,
301 : size_t size,
302 : const char* name,
303 E : const BlockGraph::Section* section) {
304 E : DCHECK(name != NULL);
305 E : BlockGraph::Block* block = block_graph_.AddBlock(type, size, name);
306 E : block->ResizeData(size);
307 E : if (section != NULL)
308 E : block->set_section(section->id());
309 E : return block;
310 E : }
311 :
312 : PEFile pe_file_;
313 : BlockGraph block_graph_;
314 : BlockGraph::Block* dos_header_block_;
315 : scoped_ptr<OrderedBlockGraph> ordered_block_graph_;
316 : };
317 :
318 : } // namespace
319 :
320 E : TEST_F(PEOrdererTest, SucceedsWithDummyImage) {
321 E : ASSERT_NO_FATAL_FAILURE(GenerateDummyImage());
322 E : ASSERT_NO_FATAL_FAILURE(RandomizeOrderedBlockGraph());
323 :
324 E : PEOrderer pe_orderer;
325 : EXPECT_TRUE(pe_orderer.OrderBlockGraph(ordered_block_graph_.get(),
326 E : dos_header_block_));
327 : ASSERT_NO_FATAL_FAILURE(VerifyValidLayout(ordered_block_graph_.get(),
328 E : dos_header_block_));
329 E : }
330 :
331 E : TEST_F(PEOrdererTest, FailsWithInvalidHeaders) {
332 E : ASSERT_NO_FATAL_FAILURE(GenerateDummyImage());
333 E : dos_header_block_->ResizeData(sizeof(IMAGE_DOS_HEADER) - 1);
334 E : ASSERT_NO_FATAL_FAILURE(InitOrderedBlockGraph());
335 :
336 E : PEOrderer pe_orderer;
337 : EXPECT_FALSE(pe_orderer.OrderBlockGraph(ordered_block_graph_.get(),
338 E : dos_header_block_));
339 E : }
340 :
341 E : TEST_F(PEOrdererTest, FailsOnMultipleRsrcSections) {
342 E : ASSERT_NO_FATAL_FAILURE(GenerateDummyImage());
343 E : block_graph_.AddSection(kResourceSectionName, kReadOnlyDataCharacteristics);
344 E : ASSERT_NO_FATAL_FAILURE(InitOrderedBlockGraph());
345 :
346 E : PEOrderer pe_orderer;
347 : EXPECT_FALSE(pe_orderer.OrderBlockGraph(ordered_block_graph_.get(),
348 E : dos_header_block_));
349 E : }
350 :
351 E : TEST_F(PEOrdererTest, FailsWithRsrcDataDirButNoRsrcSection) {
352 E : ASSERT_NO_FATAL_FAILURE(GenerateDummyImage());
353 : block_graph_.RemoveSection(
354 E : block_graph_.FindOrAddSection(kResourceSectionName, 0));
355 E : ASSERT_NO_FATAL_FAILURE(InitOrderedBlockGraph());
356 :
357 E : PEOrderer pe_orderer;
358 : EXPECT_FALSE(pe_orderer.OrderBlockGraph(ordered_block_graph_.get(),
359 E : dos_header_block_));
360 E : }
361 :
362 E : TEST_F(PEOrdererTest, FailsOnMultipleRelocSections) {
363 E : ASSERT_NO_FATAL_FAILURE(GenerateDummyImage());
364 E : block_graph_.AddSection(kRelocSectionName, kRelocCharacteristics);
365 E : ASSERT_NO_FATAL_FAILURE(InitOrderedBlockGraph());
366 :
367 E : PEOrderer pe_orderer;
368 : EXPECT_FALSE(pe_orderer.OrderBlockGraph(ordered_block_graph_.get(),
369 E : dos_header_block_));
370 E : }
371 :
372 E : TEST_F(PEOrdererTest, FailsWithRelocDataDirButNoRelocSection) {
373 E : ASSERT_NO_FATAL_FAILURE(GenerateDummyImage());
374 : block_graph_.RemoveSection(
375 E : block_graph_.FindOrAddSection(kRelocSectionName, 0));
376 E : ASSERT_NO_FATAL_FAILURE(InitOrderedBlockGraph());
377 :
378 E : PEOrderer pe_orderer;
379 : EXPECT_FALSE(pe_orderer.OrderBlockGraph(ordered_block_graph_.get(),
380 E : dos_header_block_));
381 E : }
382 :
383 E : TEST_F(PEOrdererTest, SucceedsWithTestDll) {
384 E : ASSERT_NO_FATAL_FAILURE(DecomposeTestDll());
385 E : InitOrderedBlockGraph();
386 :
387 : // NOTE: The eventual goal is to continue to enhance PEOrderer until it
388 : // is able to take a completely scrambled block-graph and generate a
389 : // working image from it. We're not quite there yet, hence this test is
390 : // a little simplistic.
391 :
392 E : PEOrderer pe_orderer;
393 : EXPECT_TRUE(pe_orderer.OrderBlockGraph(ordered_block_graph_.get(),
394 E : dos_header_block_));
395 : ASSERT_NO_FATAL_FAILURE(VerifyValidLayout(ordered_block_graph_.get(),
396 E : dos_header_block_));
397 :
398 E : ImageLayout layout(&block_graph_);
399 E : ImageLayoutBuilder builder(&layout);
400 E : ASSERT_TRUE(builder.LayoutImageHeaders(dos_header_block_));
401 E : ASSERT_TRUE(builder.LayoutOrderedBlockGraph(*ordered_block_graph_.get()));
402 E : ASSERT_TRUE(builder.Finalize());
403 :
404 : // Create a temporary file we can write a new image to.
405 E : base::FilePath temp_dir;
406 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir));
407 E : base::FilePath temp_file = temp_dir.Append(testing::kTestDllName);
408 :
409 E : PEFileWriter writer(layout);
410 E : ASSERT_TRUE(writer.WriteImage(temp_file));
411 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(temp_file));
412 E : }
413 :
414 : } // namespace orderers
415 : } // namespace pe
|