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