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/strings/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_.set_image_format(BlockGraph::PE_IMAGE);
43 E : block_graph_.AddSection(kCodeSectionName, kCodeCharacteristics);
44 : block_graph_.AddSection(kReadOnlyDataSectionName,
45 E : kReadOnlyDataCharacteristics);
46 : block_graph_.AddSection(kReadWriteDataSectionName,
47 E : kReadWriteDataCharacteristics);
48 E : block_graph_.AddSection(kTlsSectionName, kReadOnlyDataCharacteristics);
49 : block_graph_.AddSection(kResourceSectionName,
50 E : kReadOnlyDataCharacteristics);
51 E : block_graph_.AddSection(kRelocSectionName, kRelocCharacteristics);
52 :
53 E : expected_dos_header_size_ = sizeof(IMAGE_DOS_HEADER);
54 : expected_nt_headers_size_ = sizeof(IMAGE_NT_HEADERS) +
55 E : block_graph_.sections().size() * sizeof(IMAGE_SECTION_HEADER);
56 E : }
57 :
58 : // Builds a set of dummy DOS and NT headers.
59 E : void BuildHeaders(size_t extra_dos_header_bytes, size_t section_count) {
60 : dos_header_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
61 : sizeof(IMAGE_DOS_HEADER) + extra_dos_header_bytes,
62 E : "Dos Header");
63 E : ASSERT_TRUE(dos_header_block_ != NULL);
64 E : dos_header_block_->AllocateData(dos_header_block_->size());
65 :
66 : nt_headers_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
67 : sizeof(IMAGE_NT_HEADERS) + section_count * sizeof(IMAGE_SECTION_HEADER),
68 E : "Nt Headers");
69 E : ASSERT_TRUE(nt_headers_block_ != NULL);
70 E : nt_headers_block_->AllocateData(nt_headers_block_->size());
71 :
72 E : TypedBlock<IMAGE_DOS_HEADER> dos_header;
73 E : ASSERT_TRUE(dos_header.Init(0, dos_header_block_));
74 : dos_header.SetReference(BlockGraph::RELATIVE_REF,
75 : dos_header->e_lfanew,
76 : nt_headers_block_,
77 E : 0, 0);
78 :
79 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
80 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
81 E : nt_headers->OptionalHeader.FileAlignment = 512;
82 E : }
83 :
84 : size_t expected_dos_header_size_;
85 : size_t expected_nt_headers_size_;
86 :
87 : testing::DummyTransformPolicy policy_;
88 : BlockGraph block_graph_;
89 : BlockGraph::Block* dos_header_block_;
90 : BlockGraph::Block* nt_headers_block_;
91 : };
92 :
93 : } // namespace
94 :
95 E : TEST_F(PEPrepareHeadersTransformTest, ShrinkHeaders) {
96 E : BuildHeaders(100, block_graph_.sections().size() + 2);
97 :
98 E : PEPrepareHeadersTransform tx;
99 : EXPECT_TRUE(tx.TransformBlockGraph(
100 E : &policy_, &block_graph_, dos_header_block_));
101 :
102 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
103 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
104 :
105 E : EXPECT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
106 E : EXPECT_EQ(expected_nt_headers_size_, nt_headers_block_->size());
107 E : EXPECT_EQ(expected_nt_headers_size_, nt_headers_block_->data_size());
108 : EXPECT_EQ(block_graph_.sections().size(),
109 E : nt_headers->FileHeader.NumberOfSections);
110 E : }
111 :
112 E : TEST_F(PEPrepareHeadersTransformTest, GrowHeaders) {
113 E : size_t section_count = block_graph_.sections().size() - 2;
114 E : ASSERT_LT(section_count, block_graph_.sections().size());
115 E : BuildHeaders(0, section_count);
116 :
117 E : PEPrepareHeadersTransform tx;
118 : EXPECT_TRUE(tx.TransformBlockGraph(
119 E : &policy_, &block_graph_, dos_header_block_));
120 :
121 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
122 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
123 :
124 E : EXPECT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
125 E : EXPECT_EQ(expected_nt_headers_size_, nt_headers_block_->size());
126 E : EXPECT_EQ(expected_nt_headers_size_, nt_headers_block_->data_size());
127 : EXPECT_EQ(block_graph_.sections().size(),
128 E : nt_headers->FileHeader.NumberOfSections);
129 E : }
130 :
131 : } // namespace transforms
132 : } // namespace pe
|