1 : // Copyright 2011 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_file_writer.h"
16 :
17 : #include "base/path_service.h"
18 : #include "base/files/file_util.h"
19 : #include "gmock/gmock.h"
20 : #include "gtest/gtest.h"
21 : #include "syzygy/core/unittest_util.h"
22 : #include "syzygy/pe/decomposer.h"
23 : #include "syzygy/pe/pe_file.h"
24 : #include "syzygy/pe/pe_utils.h"
25 : #include "syzygy/pe/unittest_util.h"
26 :
27 : namespace pe {
28 :
29 : namespace {
30 :
31 : using block_graph::BlockGraph;
32 : using core::RelativeAddress;
33 :
34 : class PEFileWriterTest: public testing::PELibUnitTest {
35 : // Add customizations here.
36 : };
37 :
38 : } // namespace
39 :
40 E : TEST_F(PEFileWriterTest, LoadOriginalImage) {
41 : // This test baselines the other test(s) that operate on mutated, copied
42 : // versions of the DLLs.
43 E : base::FilePath image_path(testing::GetExeRelativePath(testing::kTestDllName));
44 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(image_path));
45 E : }
46 :
47 E : TEST_F(PEFileWriterTest, RewriteAndLoadImage) {
48 : // Create a temporary file we can write the new image to.
49 E : base::FilePath temp_dir;
50 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir));
51 E : base::FilePath temp_file = temp_dir.Append(testing::kTestDllName);
52 :
53 : // Decompose the original test image.
54 E : PEFile image_file;
55 E : base::FilePath image_path(testing::GetExeRelativePath(testing::kTestDllName));
56 E : ASSERT_TRUE(image_file.Init(image_path));
57 :
58 E : Decomposer decomposer(image_file);
59 E : block_graph::BlockGraph block_graph;
60 E : pe::ImageLayout image_layout(&block_graph);
61 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
62 :
63 E : PEFileWriter writer(image_layout);
64 :
65 E : ASSERT_TRUE(writer.WriteImage(temp_file));
66 E : ASSERT_NO_FATAL_FAILURE(CheckTestDll(temp_file));
67 E : }
68 :
69 E : TEST_F(PEFileWriterTest, UpdateFileChecksum) {
70 E : base::FilePath temp_dir;
71 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir));
72 :
73 : // Verify that the function fails on non-existent paths.
74 E : base::FilePath executable = temp_dir.Append(L"executable_file.exe");
75 E : EXPECT_FALSE(PEFileWriter::UpdateFileChecksum(executable));
76 :
77 : // Verify that the function fails for non-image files.
78 E : base::ScopedFILE file(base::OpenFile(executable, "wb"));
79 : // Grow the file to 16K.
80 E : ASSERT_EQ(0, fseek(file.get(), 16 * 1024, SEEK_SET));
81 E : file.reset();
82 E : EXPECT_FALSE(PEFileWriter::UpdateFileChecksum(executable));
83 :
84 : // Make a copy of our test DLL and check that we work on that.
85 E : base::FilePath input_path(testing::GetExeRelativePath(testing::kTestDllName));
86 E : base::FilePath image_path(temp_dir.Append(testing::kTestDllName));
87 E : EXPECT_TRUE(base::CopyFile(input_path, image_path));
88 E : EXPECT_TRUE(PEFileWriter::UpdateFileChecksum(image_path));
89 E : }
90 :
91 : namespace {
92 :
93 : bool WriteImageLayout(const ImageLayout& image_layout,
94 E : const base::FilePath& path) {
95 E : PEFileWriter pe_file_writer(image_layout);
96 E : return pe_file_writer.WriteImage(path);
97 E : }
98 :
99 : } // namespace
100 :
101 E : TEST_F(PEFileWriterTest, FailsForInconsistentImage) {
102 E : base::FilePath temp_dir;
103 E : ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir));
104 E : base::FilePath temp_file = temp_dir.Append(L"foo.dll");
105 :
106 E : PEFile image_file;
107 E : base::FilePath image_path(testing::GetExeRelativePath(testing::kTestDllName));
108 E : ASSERT_TRUE(image_file.Init(image_path));
109 :
110 E : Decomposer decomposer(image_file);
111 E : block_graph::BlockGraph block_graph;
112 E : pe::ImageLayout image_layout(&block_graph);
113 E : ASSERT_TRUE(decomposer.Decompose(&image_layout));
114 :
115 : BlockGraph::Block* dos_header_block =
116 E : image_layout.blocks.GetBlockByAddress(RelativeAddress(0));
117 : BlockGraph::Block* nt_headers_block =
118 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
119 E : ASSERT_TRUE(nt_headers_block != NULL);
120 :
121 : IMAGE_DOS_HEADER* dos_header =
122 E : reinterpret_cast<IMAGE_DOS_HEADER*>(dos_header_block->GetMutableData());
123 : IMAGE_NT_HEADERS* nt_headers =
124 E : reinterpret_cast<IMAGE_NT_HEADERS*>(nt_headers_block->GetMutableData());
125 E : IMAGE_SECTION_HEADER* section_headers = IMAGE_FIRST_SECTION(nt_headers);
126 :
127 E : ASSERT_TRUE(nt_headers != NULL);
128 :
129 : // To start with, the image should be perfectly writable.
130 E : EXPECT_TRUE(WriteImageLayout(image_layout, temp_file));
131 :
132 : // Invalid DOS header.
133 E : dos_header->e_magic ^= 0xF00D;
134 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
135 E : dos_header->e_magic ^= 0xF00D;
136 :
137 : // Section count mismatch.
138 E : --nt_headers->FileHeader.NumberOfSections;
139 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
140 E : ++nt_headers->FileHeader.NumberOfSections;
141 :
142 : // Inconsistent section headers and image layout.
143 E : section_headers[0].SizeOfRawData += 10 * 1024;
144 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
145 :
146 : // Make the section headers and image layout consistent again, while making
147 : // the first section overlap the second on disk.
148 E : image_layout.sections[0].data_size += 10 * 1024;
149 :
150 : // Overlapping sections on disk.
151 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
152 E : section_headers[0].SizeOfRawData -= 10 * 1024;
153 E : image_layout.sections[0].data_size -= 10 * 1024;
154 :
155 : // Overlapping sections in memory.
156 E : section_headers[0].Misc.VirtualSize += 10 * 1024;
157 E : image_layout.sections[0].size += 10 * 1024;
158 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
159 E : section_headers[0].Misc.VirtualSize -= 10 * 1024;
160 E : image_layout.sections[0].size -= 10 * 1024;
161 :
162 : // Unaligned section start on disk.
163 E : section_headers[0].PointerToRawData++;
164 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
165 E : section_headers[0].PointerToRawData--;
166 :
167 : // Unaligned section start in memory.
168 E : section_headers[0].VirtualAddress++;
169 E : image_layout.sections[0].addr += 1;
170 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
171 E : section_headers[0].VirtualAddress--;
172 E : image_layout.sections[0].addr -= 1;
173 :
174 E : size_t last_section_id = nt_headers->FileHeader.NumberOfSections - 1;
175 E : size_t file_alignment = nt_headers->OptionalHeader.FileAlignment;
176 : RelativeAddress last_section_start(
177 E : section_headers[last_section_id].VirtualAddress);
178 : RelativeAddress last_section_end = last_section_start +
179 E : section_headers[last_section_id].Misc.VirtualSize;
180 E : RelativeAddress new_block_addr = last_section_end.AlignUp(file_alignment);
181 :
182 : BlockGraph::Block* new_block = block_graph.AddBlock(
183 E : BlockGraph::DATA_BLOCK, 2 * file_alignment, "new block");
184 E : new_block->set_section(last_section_id);
185 :
186 : // Block that is outside its section entirely.
187 E : ASSERT_TRUE(image_layout.blocks.InsertBlock(new_block_addr, new_block));
188 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
189 :
190 : // Block in virtual portion of section with explicit data.
191 : section_headers[last_section_id].SizeOfRawData =
192 E : new_block_addr - last_section_start;
193 : section_headers[last_section_id].Misc.VirtualSize =
194 E : new_block_addr + new_block->size() - last_section_start;
195 : image_layout.sections.back().data_size =
196 E : section_headers[last_section_id].SizeOfRawData;
197 : image_layout.sections.back().size =
198 E : section_headers[last_section_id].Misc.VirtualSize;
199 E : ASSERT_TRUE(new_block->AllocateData(file_alignment) != NULL);
200 E : ::memset(new_block->GetMutableData(), 0xCC, new_block->data_size());
201 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
202 :
203 : // Extending the initialized data size of the section to cover the initialized
204 : // portion of the block should allow it to be written, even though half of it
205 : // is implicit.
206 E : section_headers[last_section_id].SizeOfRawData += file_alignment;
207 E : image_layout.sections.back().data_size += file_alignment;
208 E : EXPECT_TRUE(WriteImageLayout(image_layout, temp_file));
209 :
210 : // Finally, having a FileOffsetAddress reference to a location in implicit
211 : // data should fail.
212 : ASSERT_TRUE(new_block->SetReference(
213 : 0,
214 : BlockGraph::Reference(BlockGraph::FILE_OFFSET_REF,
215 : 4,
216 : new_block,
217 : file_alignment,
218 E : file_alignment)));
219 E : EXPECT_FALSE(WriteImageLayout(image_layout, temp_file));
220 E : }
221 :
222 : } // namespace pe
|