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 E : TEST_F(PECoffImageLayoutBuilderTest, BlockPadding) {
198 E : ImageLayout layout(&block_graph_);
199 E : TestImageLayoutBuilder builder(&layout, 1, 1);
200 :
201 E : const size_t kBlockPadding = 7;
202 :
203 : // Create a few dummy blocks for populating our sections.
204 : BlockGraph::Block* b1 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
205 E : 0x123, "b1");
206 : BlockGraph::Block* b2 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
207 E : 0x123, "b2");
208 : BlockGraph::Block* b3 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
209 E : 0x123, "b3");
210 E : b1->AllocateData(0x100);
211 E : b2->AllocateData(0x100);
212 E : b3->AllocateData(0x100);
213 E : memset(b1->GetMutableData(), 0xCC, 0x100);
214 E : memset(b2->GetMutableData(), 0xCC, 0x100);
215 E : memset(b3->GetMutableData(), 0xCC, 0x100);
216 :
217 : // Set block paddings.
218 E : b2->set_padding_before(kBlockPadding);
219 E : b3->set_padding_before(kBlockPadding);
220 :
221 E : const uint32 kCharacteristics = IMAGE_SCN_CNT_CODE;
222 E : EXPECT_TRUE(builder.OpenSection("foo", kCharacteristics));
223 E : EXPECT_TRUE(builder.LayoutBlock(b1));
224 E : EXPECT_TRUE(builder.LayoutBlock(b2));
225 E : EXPECT_TRUE(builder.CloseSection());
226 :
227 E : EXPECT_TRUE(builder.OpenSection("bar", kCharacteristics));
228 E : EXPECT_TRUE(builder.LayoutBlock(b3));
229 E : EXPECT_TRUE(builder.CloseSection());
230 :
231 : // Check sections. Only last block can be trimmed; any non-last block is
232 : // written up to its virtual size, before any padding is added.
233 : const std::vector<ImageLayout::SectionInfo>& sections =
234 E : builder.image_layout()->sections;
235 :
236 E : EXPECT_EQ("foo", sections[0].name);
237 E : EXPECT_EQ(RelativeAddress(0x1), sections[0].addr);
238 E : EXPECT_EQ(0x123 + kBlockPadding + 0x123, sections[0].size);
239 E : EXPECT_EQ(0x123 + kBlockPadding + 0x100, sections[0].data_size);
240 E : EXPECT_EQ(kCharacteristics, sections[0].characteristics);
241 :
242 : // Padding is applied to the first block in a section as well.
243 E : EXPECT_EQ("bar", sections[1].name);
244 E : EXPECT_EQ(sections[0].addr + sections[0].size, sections[1].addr);
245 E : EXPECT_EQ(kBlockPadding + 0x123, sections[1].size);
246 E : EXPECT_EQ(kBlockPadding + 0x100, sections[1].data_size);
247 E : EXPECT_EQ(kCharacteristics, sections[1].characteristics);
248 E : }
249 :
250 E : TEST_F(PECoffImageLayoutBuilderTest, PaddingAndBlockPadding) {
251 E : ImageLayout layout(&block_graph_);
252 E : TestImageLayoutBuilder builder(&layout, 1, 1);
253 :
254 E : const size_t kPadding = 5;
255 E : builder.set_padding(kPadding);
256 :
257 : // Test a smaller and a bigger value than kPadding.
258 E : const size_t kBlockPaddingSmall = 3;
259 E : const size_t kBlockPaddingBig = 7;
260 :
261 : // Create a few dummy blocks for populating our sections.
262 : BlockGraph::Block* b1 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
263 E : 0x123, "b1");
264 : BlockGraph::Block* b2 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
265 E : 0x123, "b2");
266 : BlockGraph::Block* b3 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
267 E : 0x123, "b3");
268 : BlockGraph::Block* b4 = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
269 E : 0x123, "b4");
270 :
271 E : b1->AllocateData(0x100);
272 E : b2->AllocateData(0x100);
273 E : b3->AllocateData(0x100);
274 E : b4->AllocateData(0x100);
275 E : memset(b1->GetMutableData(), 0xCC, 0x100);
276 E : memset(b2->GetMutableData(), 0xCC, 0x100);
277 E : memset(b3->GetMutableData(), 0xCC, 0x100);
278 E : memset(b4->GetMutableData(), 0xCC, 0x100);
279 :
280 : // Set block paddings.
281 E : b2->set_padding_before(kBlockPaddingSmall);
282 E : b4->set_padding_before(kBlockPaddingBig);
283 :
284 E : const uint32 kCharacteristics = IMAGE_SCN_CNT_CODE;
285 E : EXPECT_TRUE(builder.OpenSection("foo", kCharacteristics));
286 E : EXPECT_TRUE(builder.LayoutBlock(b1));
287 E : EXPECT_TRUE(builder.LayoutBlock(b2));
288 E : EXPECT_TRUE(builder.CloseSection());
289 :
290 E : EXPECT_TRUE(builder.OpenSection("bar", kCharacteristics));
291 E : EXPECT_TRUE(builder.LayoutBlock(b3));
292 E : EXPECT_TRUE(builder.LayoutBlock(b4));
293 E : EXPECT_TRUE(builder.CloseSection());
294 :
295 : const std::vector<ImageLayout::SectionInfo>& sections =
296 E : builder.image_layout()->sections;
297 :
298 : // Inter-block padding is bigger, that should be in effect.
299 E : EXPECT_EQ("foo", sections[0].name);
300 E : EXPECT_EQ(RelativeAddress(0x1), sections[0].addr);
301 E : EXPECT_EQ(0x123 + kPadding + 0x123, sections[0].size);
302 E : EXPECT_EQ(0x123 + kPadding + 0x100, sections[0].data_size);
303 E : EXPECT_EQ(kCharacteristics, sections[0].characteristics);
304 :
305 : // Block's own padding is bigger, that should be in effect.
306 E : EXPECT_EQ("bar", sections[1].name);
307 E : EXPECT_EQ(sections[0].addr + sections[0].size, sections[1].addr);
308 E : EXPECT_EQ(0x123 + kBlockPaddingBig + 0x123, sections[1].size);
309 E : EXPECT_EQ(0x123 + kBlockPaddingBig + 0x100, sections[1].data_size);
310 E : EXPECT_EQ(kCharacteristics, sections[1].characteristics);
311 E : }
312 :
313 E : TEST_F(PECoffImageLayoutBuilderTest, Align) {
314 E : ImageLayout layout(&block_graph_);
315 E : TestImageLayoutBuilder builder(&layout, 1, 1);
316 :
317 E : const size_t kAlignment = 16U;
318 E : const size_t kBlockSize = 17U;
319 E : const BlockGraph::Offset kOffsetMin = -1;
320 E : const BlockGraph::Offset kOffsetMax = 100;
321 :
322 : // Create aligned blocks with different alignment offsets.
323 E : std::vector<BlockGraph::Block*> blocks;
324 E : for (BlockGraph::Offset i = kOffsetMin; i < kOffsetMax; ++i) {
325 : BlockGraph::Block* block = block_graph_.AddBlock(BlockGraph::CODE_BLOCK,
326 : kBlockSize,
327 E : "b" + std::to_string(i));
328 E : block->AllocateData(kBlockSize);
329 E : memset(block->GetMutableData(), 0xCC, kBlockSize);
330 :
331 E : block->set_alignment(kAlignment);
332 E : block->set_alignment_offset(i);
333 :
334 E : blocks.push_back(block);
335 E : }
336 :
337 E : const uint32 kCharacteristics = IMAGE_SCN_CNT_CODE;
338 E : EXPECT_TRUE(builder.OpenSection("foo", kCharacteristics));
339 E : for (BlockGraph::Block* block : blocks) {
340 E : EXPECT_TRUE(builder.LayoutBlock(block));
341 E : }
342 E : EXPECT_TRUE(builder.CloseSection());
343 :
344 : const std::vector<ImageLayout::SectionInfo>& sections =
345 E : builder.image_layout()->sections;
346 E : EXPECT_EQ("foo", sections[0].name);
347 E : EXPECT_EQ(RelativeAddress(0x1), sections[0].addr);
348 :
349 : // Check if each block is placed at an address that respects its alignment
350 : // and that the blocks do not overlap, nor are they placed too far away from
351 : // each other.
352 : // This test uses the fact that Block::addr_ is populated upon layout.
353 E : BlockGraph::RelativeAddress last_address;
354 E : bool first = true;
355 E : for (const BlockGraph::Block* block : blocks) {
356 E : BlockGraph::RelativeAddress curr_address = block->addr();
357 E : BlockGraph::Offset curr_offset = block->alignment_offset();
358 :
359 : // Test proper alignment.
360 E : EXPECT_TRUE((curr_address + curr_offset).IsAligned(kAlignment));
361 :
362 E : if (first) {
363 E : first = false;
364 :
365 : // This is true because kOffsetMin is negative.
366 E : EXPECT_EQ(static_cast<uint32>(-kOffsetMin), curr_address.value());
367 E : } else {
368 : // The space between the blocks is the difference of the addresses minus
369 : // the data size.
370 : int space_between_blocks = curr_address - last_address -
371 E : static_cast<int>(kBlockSize);
372 :
373 : // Check that blocks do not overlap.
374 E : EXPECT_GE(space_between_blocks, 0);
375 :
376 : // If the space is bigger then kAlignment bytes then the block could have
377 : // been placed kAlignment bytes ahead.
378 E : EXPECT_LT(space_between_blocks, static_cast<int>(kAlignment));
379 : }
380 E : last_address = curr_address;
381 E : }
382 E : }
383 :
384 : } // namespace pe
|