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/transforms/pe_prepare_headers_transform.h"
16 :
17 : #include "base/stringprintf.h"
18 : #include "gtest/gtest.h"
19 : #include "syzygy/block_graph/typed_block.h"
20 : #include "syzygy/block_graph/unittest_util.h"
21 : #include "syzygy/pe/pe_utils.h"
22 :
23 : namespace pe {
24 : namespace transforms {
25 :
26 : using block_graph::BlockGraph;
27 : using block_graph::ConstTypedBlock;
28 : using block_graph::TypedBlock;
29 :
30 : namespace {
31 :
32 : class PEPrepareHeadersTransformTest : public testing::Test {
33 : public:
34 : PEPrepareHeadersTransformTest()
35 : : expected_dos_header_size_(0),
36 : expected_nt_headers_size_(0),
37 : dos_header_block_(NULL),
38 E : nt_headers_block_(NULL) {
39 E : }
40 :
41 E : virtual void SetUp() {
42 E : block_graph_.AddSection(kCodeSectionName, kCodeCharacteristics);
43 : block_graph_.AddSection(kReadOnlyDataSectionName,
44 E : kReadOnlyDataCharacteristics);
45 : block_graph_.AddSection(kReadWriteDataSectionName,
46 E : kReadWriteDataCharacteristics);
47 E : block_graph_.AddSection(kTlsSectionName, kReadOnlyDataCharacteristics);
48 : block_graph_.AddSection(kResourceSectionName,
49 E : kReadOnlyDataCharacteristics);
50 E : block_graph_.AddSection(kRelocSectionName, kRelocCharacteristics);
51 :
52 E : expected_dos_header_size_ = sizeof(IMAGE_DOS_HEADER);
53 : expected_nt_headers_size_ = sizeof(IMAGE_NT_HEADERS) +
54 E : block_graph_.sections().size() * sizeof(IMAGE_SECTION_HEADER);
55 E : }
56 :
57 : // Builds a set of dummy DOS and NT headers.
58 E : void BuildHeaders(size_t extra_dos_header_bytes, size_t section_count) {
59 : dos_header_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
60 : sizeof(IMAGE_DOS_HEADER) + extra_dos_header_bytes,
61 E : "Dos Header");
62 E : ASSERT_TRUE(dos_header_block_ != NULL);
63 E : dos_header_block_->AllocateData(dos_header_block_->size());
64 :
65 : nt_headers_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
66 : sizeof(IMAGE_NT_HEADERS) + section_count * sizeof(IMAGE_SECTION_HEADER),
67 E : "Nt Headers");
68 E : ASSERT_TRUE(nt_headers_block_ != NULL);
69 E : nt_headers_block_->AllocateData(nt_headers_block_->size());
70 :
71 E : TypedBlock<IMAGE_DOS_HEADER> dos_header;
72 E : ASSERT_TRUE(dos_header.Init(0, dos_header_block_));
73 : dos_header.SetReference(BlockGraph::RELATIVE_REF,
74 : dos_header->e_lfanew,
75 : nt_headers_block_,
76 E : 0, 0);
77 :
78 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
79 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
80 E : nt_headers->OptionalHeader.FileAlignment = 512;
81 E : }
82 :
83 : size_t expected_dos_header_size_;
84 : size_t expected_nt_headers_size_;
85 :
86 : testing::DummyTransformPolicy policy_;
87 : BlockGraph block_graph_;
88 : BlockGraph::Block* dos_header_block_;
89 : BlockGraph::Block* nt_headers_block_;
90 : };
91 :
92 : } // namespace
93 :
94 E : TEST_F(PEPrepareHeadersTransformTest, ShrinkHeaders) {
95 E : BuildHeaders(100, block_graph_.sections().size() + 2);
96 :
97 E : PEPrepareHeadersTransform tx;
98 : EXPECT_TRUE(tx.TransformBlockGraph(
99 E : &policy_, &block_graph_, dos_header_block_));
100 :
101 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
102 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
103 :
104 E : EXPECT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
105 E : EXPECT_EQ(expected_nt_headers_size_, nt_headers_block_->size());
106 E : EXPECT_EQ(expected_nt_headers_size_, nt_headers_block_->data_size());
107 : EXPECT_EQ(block_graph_.sections().size(),
108 E : nt_headers->FileHeader.NumberOfSections);
109 E : }
110 :
111 E : TEST_F(PEPrepareHeadersTransformTest, GrowHeaders) {
112 E : size_t section_count = block_graph_.sections().size() - 2;
113 E : ASSERT_LT(section_count, block_graph_.sections().size());
114 E : BuildHeaders(0, section_count);
115 :
116 E : PEPrepareHeadersTransform tx;
117 : EXPECT_TRUE(tx.TransformBlockGraph(
118 E : &policy_, &block_graph_, dos_header_block_));
119 :
120 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
121 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
122 :
123 E : EXPECT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
124 E : EXPECT_EQ(expected_nt_headers_size_, nt_headers_block_->size());
125 E : EXPECT_EQ(expected_nt_headers_size_, nt_headers_block_->data_size());
126 : EXPECT_EQ(block_graph_.sections().size(),
127 E : nt_headers->FileHeader.NumberOfSections);
128 E : }
129 :
130 : } // namespace transforms
131 : } // namespace pe
|