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