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/pe_image_layout_builder.h"
16 :
17 : #include <algorithm>
18 : #include <cstdlib>
19 : #include <ctime>
20 :
21 : #include "base/files/file_util.h"
22 : #include "gmock/gmock.h"
23 : #include "gtest/gtest.h"
24 : #include "syzygy/block_graph/typed_block.h"
25 : #include "syzygy/block_graph/unittest_util.h"
26 : #include "syzygy/block_graph/orderers/original_orderer.h"
27 : #include "syzygy/block_graph/orderers/random_orderer.h"
28 : #include "syzygy/core/unittest_util.h"
29 : #include "syzygy/pe/decomposer.h"
30 : #include "syzygy/pe/pe_file_writer.h"
31 : #include "syzygy/pe/pe_utils.h"
32 : #include "syzygy/pe/unittest_util.h"
33 : #include "syzygy/pe/transforms/pe_prepare_headers_transform.h"
34 :
35 : namespace pe {
36 :
37 : using block_graph::BlockGraph;
38 : using block_graph::ConstTypedBlock;
39 : using block_graph::OrderedBlockGraph;
40 : using core::AddressRange;
41 : using core::RelativeAddress;
42 :
43 : namespace {
44 :
45 : class PEImageLayoutBuilderTest : public testing::PELibUnitTest {
46 : typedef testing::PELibUnitTest Super;
47 :
48 : public:
49 E : PEImageLayoutBuilderTest()
50 : : image_layout_(&block_graph_), dos_header_block_(NULL) {
51 E : }
52 :
53 E : void SetUp() {
54 E : Super::SetUp();
55 :
56 : // Create a temporary file we can write a new image to.
57 E : base::FilePath temp_dir;
58 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir));
59 E : temp_file_ = temp_dir.Append(testing::kTestDllName);
60 :
61 : // Decompose the test DLL.
62 E : image_path_ = testing::GetExeRelativePath(testing::kTestDllName);
63 E : ASSERT_TRUE(image_file_.Init(image_path_));
64 :
65 E : Decomposer decomposer(image_file_);
66 E : ASSERT_TRUE(decomposer.Decompose(&image_layout_));
67 :
68 : dos_header_block_ =
69 E : image_layout_.blocks.GetBlockByAddress(RelativeAddress(0));
70 E : ASSERT_TRUE(dos_header_block_ != NULL);
71 E : ASSERT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
72 :
73 : // Prepare the headers. This puts our DOS stub in place.
74 E : transforms::PEPrepareHeadersTransform prep_headers;
75 E : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
76 : &prep_headers, &policy_, &block_graph_, dos_header_block_));
77 E : }
78 :
79 : protected:
80 : testing::DummyTransformPolicy policy_;
81 : base::FilePath image_path_;
82 : PEFile image_file_;
83 : BlockGraph block_graph_;
84 : ImageLayout image_layout_;
85 : BlockGraph::Block* dos_header_block_;
86 :
87 : base::FilePath temp_file_;
88 : };
89 :
90 : } // namespace
91 :
92 E : TEST_F(PEImageLayoutBuilderTest, Initialization) {
93 E : ImageLayout layout(&block_graph_);
94 E : PEImageLayoutBuilder builder(&layout);
95 :
96 E : EXPECT_EQ(&layout, builder.image_layout());
97 E : EXPECT_EQ(&block_graph_, builder.block_graph());
98 E : EXPECT_EQ(NULL, builder.dos_header_block());
99 E : EXPECT_EQ(NULL, builder.nt_headers_block());
100 E : EXPECT_EQ(0, builder.padding());
101 E : EXPECT_EQ(1, builder.code_alignment());
102 E : }
103 :
104 E : TEST_F(PEImageLayoutBuilderTest, Accessors) {
105 E : ImageLayout layout(&block_graph_);
106 E : PEImageLayoutBuilder builder(&layout);
107 :
108 E : builder.set_padding(16);
109 E : builder.set_code_alignment(8);
110 E : EXPECT_EQ(16, builder.padding());
111 E : EXPECT_EQ(8, builder.code_alignment());
112 E : }
113 :
114 E : TEST_F(PEImageLayoutBuilderTest, LayoutImageHeaders) {
115 E : ImageLayout layout(&block_graph_);
116 E : PEImageLayoutBuilder builder(&layout);
117 :
118 E : EXPECT_TRUE(builder.LayoutImageHeaders(dos_header_block_));
119 E : EXPECT_EQ(dos_header_block_, builder.dos_header_block());
120 E : EXPECT_TRUE(builder.nt_headers_block() != NULL);
121 E : }
122 :
123 E : TEST_F(PEImageLayoutBuilderTest, RewriteTestDll) {
124 E : OrderedBlockGraph obg(&block_graph_);
125 E : block_graph::orderers::OriginalOrderer orig_orderer;
126 E : ASSERT_TRUE(orig_orderer.OrderBlockGraph(&obg, dos_header_block_));
127 :
128 E : ImageLayout layout(&block_graph_);
129 E : PEImageLayoutBuilder builder(&layout);
130 E : ASSERT_TRUE(builder.LayoutImageHeaders(dos_header_block_));
131 E : EXPECT_TRUE(builder.LayoutOrderedBlockGraph(obg));
132 E : EXPECT_TRUE(builder.Finalize());
133 :
134 E : PEFileWriter writer(layout);
135 E : ASSERT_TRUE(writer.WriteImage(temp_file_));
136 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(temp_file_));
137 :
138 : // We expect all of the sections to have been placed at the same addresses,
139 : // have the same size, etc (except for relocs).
140 E : EXPECT_EQ(image_layout_.sections.size(), layout.sections.size());
141 E : for (size_t i = 0; i < image_layout_.sections.size() - 1; ++i)
142 E : EXPECT_EQ(image_layout_.sections[i], layout.sections[i]);
143 :
144 : // We expect our image to be no bigger. In fact, we are generally smaller as
145 : // we trim some cruft from the .relocs section.
146 : int64 orig_size, rewritten_size;
147 E : ASSERT_TRUE(base::GetFileSize(image_path_, &orig_size));
148 E : ASSERT_TRUE(base::GetFileSize(temp_file_, &rewritten_size));
149 E : EXPECT_LE(rewritten_size, orig_size);
150 E : }
151 :
152 E : TEST_F(PEImageLayoutBuilderTest, PadTestDll) {
153 E : OrderedBlockGraph obg(&block_graph_);
154 E : block_graph::orderers::OriginalOrderer orig_orderer;
155 E : ASSERT_TRUE(orig_orderer.OrderBlockGraph(&obg, dos_header_block_));
156 :
157 : // We modify the CV info so that the debugger doesn't try to load the
158 : // wrong symbols for this image.
159 E : ASSERT_NO_FATAL_FAILURE(testing::TwiddlePdbGuidAndPath(dos_header_block_));
160 :
161 E : ImageLayout layout(&block_graph_);
162 E : PEImageLayoutBuilder builder(&layout);
163 E : builder.set_padding(100);
164 E : ASSERT_TRUE(builder.LayoutImageHeaders(dos_header_block_));
165 E : EXPECT_TRUE(builder.LayoutOrderedBlockGraph(obg));
166 E : EXPECT_TRUE(builder.Finalize());
167 :
168 E : PEFileWriter writer(layout);
169 E : ASSERT_TRUE(writer.WriteImage(temp_file_));
170 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(temp_file_));
171 :
172 : // We expect the sections to have gotten longer by the right number of bytes.
173 E : EXPECT_EQ(image_layout_.sections.size(), layout.sections.size());
174 E : EXPECT_EQ(image_layout_.sections.size(), obg.ordered_sections().size());
175 : OrderedBlockGraph::SectionList::const_iterator obg_section_it =
176 E : obg.ordered_sections().begin();
177 E : size_t expected_file_size_increase = 0;
178 E : for (size_t i = 0; i < image_layout_.sections.size(); ++i) {
179 E : const ImageLayout::SectionInfo& old_section = image_layout_.sections[i];
180 E : const ImageLayout::SectionInfo& new_section = layout.sections[i];
181 :
182 : // All sections (except for .reloc, the last one) should only have grown
183 : // in size. As each of the non-reloc sections may now spread across more
184 : // pages than before, the .reloc section itself may have grown (it contains
185 : // a structure per page of the image). But, due to the fact that the MS
186 : // linker generally creates an overly large .reloc section, it may also have
187 : // stayed the same size or gotten smaller.
188 E : if (i + 1 < image_layout_.sections.size()) {
189 : // We expect the section to have increased in size by at least 100
190 : // in between each and every block.
191 : size_t added_bytes =
192 E : 100 * ((*obg_section_it)->ordered_blocks().size() - 1);
193 E : EXPECT_GE(new_section.size, old_section.size + added_bytes);
194 E : EXPECT_GE(new_section.data_size, old_section.data_size);
195 : }
196 :
197 : // Keep track of the total number of new bytes that should be making it
198 : // to disk.
199 : expected_file_size_increase += new_section.data_size -
200 E : old_section.data_size;
201 :
202 E : ++obg_section_it;
203 E : }
204 :
205 : int64 orig_size, rewritten_size;
206 E : ASSERT_TRUE(base::GetFileSize(image_path_, &orig_size));
207 E : ASSERT_TRUE(base::GetFileSize(temp_file_, &rewritten_size));
208 E : EXPECT_GE(rewritten_size, orig_size + expected_file_size_increase);
209 E : }
210 :
211 E : TEST_F(PEImageLayoutBuilderTest, CodeAlignmentTestDll) {
212 E : OrderedBlockGraph obg(&block_graph_);
213 E : block_graph::orderers::OriginalOrderer orig_orderer;
214 E : ASSERT_TRUE(orig_orderer.OrderBlockGraph(&obg, dos_header_block_));
215 :
216 : // We modify the CV info so that the debugger doesn't try to load the
217 : // wrong symbols for this image.
218 E : ASSERT_NO_FATAL_FAILURE(testing::TwiddlePdbGuidAndPath(dos_header_block_));
219 :
220 E : const uint32 kCodeAlignment = 8;
221 E : ImageLayout layout(&block_graph_);
222 E : PEImageLayoutBuilder builder(&layout);
223 E : builder.set_code_alignment(kCodeAlignment);
224 E : ASSERT_TRUE(builder.LayoutImageHeaders(dos_header_block_));
225 E : EXPECT_TRUE(builder.LayoutOrderedBlockGraph(obg));
226 E : EXPECT_TRUE(builder.Finalize());
227 :
228 E : PEFileWriter writer(layout);
229 E : ASSERT_TRUE(writer.WriteImage(temp_file_));
230 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(temp_file_));
231 :
232 : // Validate that code blocks are aligned correctly.
233 E : BlockGraph::AddressSpace::RangeMapConstIter iter = layout.blocks.begin();
234 E : for (; iter != layout.blocks.end(); ++iter) {
235 E : BlockGraph::Block* block = iter->second;
236 E : BlockGraph::AddressSpace::Range range = iter->first;
237 E : if (block->type() == BlockGraph::CODE_BLOCK) {
238 E : EXPECT_TRUE(range.start().IsAligned(kCodeAlignment));
239 : }
240 E : }
241 E : }
242 :
243 E : TEST_F(PEImageLayoutBuilderTest, RandomizeTestDll) {
244 E : OrderedBlockGraph obg(&block_graph_);
245 E : block_graph::orderers::RandomOrderer random_orderer(true);
246 E : ASSERT_TRUE(random_orderer.OrderBlockGraph(&obg, dos_header_block_));
247 :
248 : // We modify the CV info so that the debugger doesn't try to load the
249 : // wrong symbols for this image.
250 E : ASSERT_NO_FATAL_FAILURE(testing::TwiddlePdbGuidAndPath(dos_header_block_));
251 :
252 E : ImageLayout layout(&block_graph_);
253 E : PEImageLayoutBuilder builder(&layout);
254 E : ASSERT_TRUE(builder.LayoutImageHeaders(dos_header_block_));
255 E : EXPECT_TRUE(builder.LayoutOrderedBlockGraph(obg));
256 E : EXPECT_TRUE(builder.Finalize());
257 :
258 E : PEFileWriter writer(layout);
259 E : ASSERT_TRUE(writer.WriteImage(temp_file_));
260 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(temp_file_));
261 E : }
262 :
263 E : TEST_F(PEImageLayoutBuilderTest, ShiftTestDll) {
264 : // Create an empty section. We will place this at the beginning of the
265 : // image to ensure that everything gets shifted by a fixed amount. A loadable
266 : // module is a good indication that we properly parsed everything.
267 : BlockGraph::Section* section = block_graph_.AddSection(
268 E : ".empty", kReadOnlyDataCharacteristics);
269 : BlockGraph::Block* block = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
270 E : 10 * 1024, ".empty");
271 E : block->AllocateData(block->size());
272 E : ::memset(block->GetMutableData(), 0xcc, block->data_size());
273 E : block->set_section(section->id());
274 :
275 : // Prepare the headers (again). We need to do this to make sure that the image
276 : // headers accurately reflect the number of sections as we've added a new
277 : // one.
278 E : transforms::PEPrepareHeadersTransform prep_headers;
279 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
280 E : &prep_headers, &policy_, &block_graph_, dos_header_block_));
281 :
282 E : OrderedBlockGraph obg(&block_graph_);
283 E : block_graph::orderers::OriginalOrderer orig_orderer;
284 E : ASSERT_TRUE(orig_orderer.OrderBlockGraph(&obg, dos_header_block_));
285 :
286 : // Move the new section to the beginning of the image. This causes everything
287 : // to be shifted by a fixed amount.
288 E : obg.PlaceAtHead(section);
289 :
290 : // We modify the CV info so that the debugger doesn't try to load the
291 : // wrong symbols for this image.
292 E : ASSERT_NO_FATAL_FAILURE(testing::TwiddlePdbGuidAndPath(dos_header_block_));
293 :
294 E : ImageLayout layout(&block_graph_);
295 E : PEImageLayoutBuilder builder(&layout);
296 E : ASSERT_TRUE(builder.LayoutImageHeaders(dos_header_block_));
297 E : EXPECT_TRUE(builder.LayoutOrderedBlockGraph(obg));
298 E : EXPECT_TRUE(builder.Finalize());
299 :
300 E : PEFileWriter writer(layout);
301 E : ASSERT_TRUE(writer.WriteImage(temp_file_));
302 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(temp_file_));
303 :
304 : // Read the rewritten DLL and validate that the resources have moved.
305 E : PEFile new_image_file;
306 E : ASSERT_TRUE(new_image_file.Init(temp_file_));
307 : const IMAGE_DATA_DIRECTORY* old_data_dir =
308 E : image_file_.nt_headers()->OptionalHeader.DataDirectory;
309 : const IMAGE_DATA_DIRECTORY* new_data_dir =
310 E : new_image_file.nt_headers()->OptionalHeader.DataDirectory;
311 : ASSERT_EQ(
312 : old_data_dir[IMAGE_DIRECTORY_ENTRY_RESOURCE].Size,
313 E : new_data_dir[IMAGE_DIRECTORY_ENTRY_RESOURCE].Size);
314 : ASSERT_NE(
315 : old_data_dir[IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress,
316 E : new_data_dir[IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress);
317 E : }
318 :
319 : } // namespace pe
|