1 : // Copyright 2013 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/pe_coff_image_layout_builder.h"
16 :
17 : #include "gtest/gtest.h"
18 : #include "syzygy/common/align.h"
19 :
20 : namespace pe {
21 :
22 : using block_graph::BlockGraph;
23 : using core::RelativeAddress;
24 :
25 : namespace {
26 :
27 : class TestImageLayoutBuilder : public PECoffImageLayoutBuilder {
28 : public:
29 : // Make the builder publicly constructible.
30 E : TestImageLayoutBuilder(ImageLayout* image_layout,
31 : size_t section_alignment,
32 : size_t file_alignment)
33 : : PECoffImageLayoutBuilder(image_layout) {
34 E : PECoffImageLayoutBuilder::Init(section_alignment, file_alignment);
35 :
36 : // Advance cursor to simulate headers having been written.
37 E : cursor_ += 1;
38 E : }
39 : };
40 :
41 : class PECoffImageLayoutBuilderTest : public testing::Test {
42 : public:
43 E : PECoffImageLayoutBuilderTest() {
44 E : }
45 :
46 : protected:
47 : BlockGraph block_graph_;
48 :
49 : private:
50 : DISALLOW_COPY_AND_ASSIGN(PECoffImageLayoutBuilderTest);
51 : };
52 :
53 : } // namespace
54 :
55 E : TEST_F(PECoffImageLayoutBuilderTest, Initialization) {
56 E : ImageLayout layout(&block_graph_);
57 E : TestImageLayoutBuilder builder(&layout, 1, 1);
58 :
59 E : EXPECT_EQ(&layout, builder.image_layout());
60 E : EXPECT_EQ(&block_graph_, builder.block_graph());
61 E : }
62 :
63 E : TEST_F(PECoffImageLayoutBuilderTest, AddSection) {
64 E : ImageLayout layout(&block_graph_);
65 E : TestImageLayoutBuilder builder(&layout, 1, 1);
66 :
67 : // Create a few dummy blocks for populating our sections.
68 : BlockGraph::Block* b1 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
69 E : 0x1234, "b1");
70 : BlockGraph::Block* b2 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
71 E : 0x1234, "b2");
72 E : b1->AllocateData(0x1000);
73 E : b2->AllocateData(0x1000);
74 E : memset(b1->GetMutableData(), 0xCC, 0x1000);
75 E : memset(b2->GetMutableData(), 0xCC, 0x1000);
76 :
77 E : const uint32 kCharacteristics = IMAGE_SCN_CNT_CODE;
78 E : EXPECT_TRUE(builder.OpenSection("foo", kCharacteristics));
79 E : EXPECT_TRUE(builder.LayoutBlock(b1));
80 E : EXPECT_TRUE(builder.CloseSection());
81 :
82 E : EXPECT_TRUE(builder.OpenSection("bar", kCharacteristics));
83 E : EXPECT_TRUE(builder.LayoutBlock(b2));
84 E : EXPECT_TRUE(builder.CloseSection());
85 :
86 : // Check sections.
87 : const std::vector<ImageLayout::SectionInfo>& sections =
88 E : builder.image_layout()->sections;
89 :
90 E : EXPECT_EQ("foo", sections[0].name);
91 E : EXPECT_EQ(RelativeAddress(0x1), sections[0].addr);
92 E : EXPECT_EQ(0x1234, sections[0].size);
93 E : EXPECT_EQ(0x1000, sections[0].data_size);
94 E : EXPECT_EQ(kCharacteristics, sections[0].characteristics);
95 :
96 E : EXPECT_EQ("bar", sections[1].name);
97 E : EXPECT_EQ(sections[0].addr + sections[0].size, sections[1].addr);
98 E : EXPECT_EQ(0x1234, sections[1].size);
99 E : EXPECT_EQ(0x1000, sections[1].data_size);
100 E : EXPECT_EQ(kCharacteristics, sections[1].characteristics);
101 E : }
102 :
103 E : TEST_F(PECoffImageLayoutBuilderTest, Alignment) {
104 E : ImageLayout layout(&block_graph_);
105 :
106 E : const size_t kSectionAlignment = 300;
107 E : const size_t kFileAlignment = 150;
108 E : TestImageLayoutBuilder builder(&layout, kSectionAlignment, kFileAlignment);
109 :
110 : // Create a few dummy blocks for populating our sections.
111 : BlockGraph::Block* b1 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
112 E : 0x1234, "b1");
113 : BlockGraph::Block* b2 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
114 E : 0x1234, "b2");
115 E : b1->AllocateData(0x1000);
116 E : b2->AllocateData(0x1000);
117 E : memset(b1->GetMutableData(), 0xCC, 0x1000);
118 E : memset(b2->GetMutableData(), 0xCC, 0x1000);
119 :
120 E : const uint32 kCharacteristics = IMAGE_SCN_CNT_CODE;
121 E : EXPECT_TRUE(builder.OpenSection("foo", kCharacteristics));
122 E : EXPECT_TRUE(builder.LayoutBlock(b1));
123 E : EXPECT_TRUE(builder.CloseSection());
124 :
125 E : EXPECT_TRUE(builder.OpenSection("bar", kCharacteristics));
126 E : EXPECT_TRUE(builder.LayoutBlock(b2));
127 E : EXPECT_TRUE(builder.CloseSection());
128 :
129 : // Check sections; section addresses should have been rounded up, as well
130 : // as raw data sizes. Virtual sizes should be untouched.
131 : const std::vector<ImageLayout::SectionInfo>& sections =
132 E : builder.image_layout()->sections;
133 :
134 E : EXPECT_EQ("foo", sections[0].name);
135 E : EXPECT_EQ(RelativeAddress(0x1).AlignUp(kSectionAlignment), sections[0].addr);
136 E : EXPECT_EQ(0x1234, sections[0].size);
137 E : EXPECT_EQ(common::AlignUp(0x1000, kFileAlignment), sections[0].data_size);
138 E : EXPECT_EQ(kCharacteristics, sections[0].characteristics);
139 :
140 E : EXPECT_EQ("bar", sections[1].name);
141 : EXPECT_EQ((sections[0].addr + sections[0].size).AlignUp(kSectionAlignment),
142 E : sections[1].addr);
143 E : EXPECT_EQ(0x1234, sections[1].size);
144 E : EXPECT_EQ(common::AlignUp(0x1000, kFileAlignment), sections[1].data_size);
145 E : EXPECT_EQ(kCharacteristics, sections[1].characteristics);
146 E : }
147 :
148 E : TEST_F(PECoffImageLayoutBuilderTest, Padding) {
149 E : ImageLayout layout(&block_graph_);
150 E : TestImageLayoutBuilder builder(&layout, 1, 1);
151 :
152 E : const size_t kPadding = 100;
153 E : builder.set_padding(kPadding);
154 :
155 : // Create a few dummy blocks for populating our sections.
156 : BlockGraph::Block* b1 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
157 E : 0x1234, "b1");
158 : BlockGraph::Block* b2 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
159 E : 0x1234, "b2");
160 : BlockGraph::Block* b3 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
161 E : 0x123, "b3");
162 E : b1->AllocateData(0x1000);
163 E : b2->AllocateData(0x1000);
164 E : b3->AllocateData(0x100);
165 E : memset(b1->GetMutableData(), 0xCC, 0x1000);
166 E : memset(b2->GetMutableData(), 0xCC, 0x1000);
167 E : memset(b3->GetMutableData(), 0xCC, 0x100);
168 :
169 E : const uint32 kCharacteristics = IMAGE_SCN_CNT_CODE;
170 E : EXPECT_TRUE(builder.OpenSection("foo", kCharacteristics));
171 E : EXPECT_TRUE(builder.LayoutBlock(b1));
172 E : EXPECT_TRUE(builder.LayoutBlock(b3));
173 E : EXPECT_TRUE(builder.CloseSection());
174 :
175 E : EXPECT_TRUE(builder.OpenSection("bar", kCharacteristics));
176 E : EXPECT_TRUE(builder.LayoutBlock(b2));
177 E : EXPECT_TRUE(builder.CloseSection());
178 :
179 : // Check sections. Only last block can be trimmed; any non-last block is
180 : // written up to its virtual size, before any padding is added.
181 : const std::vector<ImageLayout::SectionInfo>& sections =
182 E : builder.image_layout()->sections;
183 :
184 E : EXPECT_EQ("foo", sections[0].name);
185 E : EXPECT_EQ(RelativeAddress(0x1), sections[0].addr);
186 E : EXPECT_EQ(0x1234 + kPadding + 0x123, sections[0].size);
187 E : EXPECT_EQ(0x1234 + kPadding + 0x100, sections[0].data_size);
188 E : EXPECT_EQ(kCharacteristics, sections[0].characteristics);
189 :
190 E : EXPECT_EQ("bar", sections[1].name);
191 E : EXPECT_EQ(sections[0].addr + sections[0].size, sections[1].addr);
192 E : EXPECT_EQ(0x1234, sections[1].size);
193 E : EXPECT_EQ(0x1000, sections[1].data_size);
194 E : EXPECT_EQ(kCharacteristics, sections[1].characteristics);
195 E : }
196 :
197 : } // namespace pe
|