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_utils.h"
16 :
17 : #include "gmock/gmock.h"
18 : #include "gtest/gtest.h"
19 : #include "syzygy/pe/transforms/add_imports_transform.h"
20 :
21 : namespace pe {
22 :
23 : using block_graph::BlockGraph;
24 : using core::RelativeAddress;
25 : using pe::transforms::AddImportsTransform;
26 :
27 : namespace {
28 :
29 : class PEUtilsTest : public testing::Test {
30 : public:
31 : PEUtilsTest()
32 : : nt_headers_block_(NULL),
33 : dos_header_block_(NULL),
34 : main_entry_point_block_(NULL),
35 : tls_initializer_block_(NULL),
36 : nt_headers_(NULL),
37 E : dos_header_(NULL) {
38 E : }
39 :
40 E : virtual void SetUp() {
41 : // Create the NT headers block.
42 E : ASSERT_NO_FATAL_FAILURE(CreateNtHeadersBlock());
43 : // And the DOS header block.
44 E : ASSERT_NO_FATAL_FAILURE(CreateDosHeaderBlock());
45 : // And set-up some entry points.
46 E : ASSERT_NO_FATAL_FAILURE(CreateEntryPoints());
47 E : }
48 :
49 : protected:
50 : void CreateDosHeaderBlock();
51 : void CreateNtHeadersBlock();
52 : void CreateEntryPoints();
53 :
54 : BlockGraph block_graph_;
55 :
56 : BlockGraph::Block* nt_headers_block_;
57 : BlockGraph::Block* dos_header_block_;
58 : BlockGraph::Block* main_entry_point_block_;
59 : BlockGraph::Block* tls_initializer_block_;
60 : IMAGE_DOS_HEADER* dos_header_;
61 : IMAGE_NT_HEADERS* nt_headers_;
62 : };
63 :
64 E : void PEUtilsTest::CreateNtHeadersBlock() {
65 : nt_headers_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
66 : sizeof(IMAGE_NT_HEADERS),
67 E : "NT Headers");
68 E : ASSERT_TRUE(nt_headers_block_ != NULL);
69 :
70 : nt_headers_ = reinterpret_cast<IMAGE_NT_HEADERS*>(
71 E : nt_headers_block_->AllocateData(sizeof(IMAGE_NT_HEADERS)));
72 E : ASSERT_TRUE(nt_headers_ != NULL);
73 :
74 E : nt_headers_->Signature = IMAGE_NT_SIGNATURE;
75 : nt_headers_->FileHeader.SizeOfOptionalHeader =
76 E : sizeof(nt_headers_->OptionalHeader);
77 E : nt_headers_->OptionalHeader.Magic = IMAGE_NT_OPTIONAL_HDR_MAGIC;
78 E : }
79 :
80 E : void PEUtilsTest::CreateDosHeaderBlock() {
81 : dos_header_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
82 : sizeof(IMAGE_DOS_HEADER),
83 E : "DOS Header");
84 E : ASSERT_TRUE(dos_header_block_ != NULL);
85 :
86 : dos_header_ = reinterpret_cast<IMAGE_DOS_HEADER*>(
87 E : dos_header_block_->AllocateData(dos_header_block_->size()));
88 E : ASSERT_TRUE(dos_header_ != NULL);
89 :
90 : // Set the correct magic constants in the manufactured DOS header.
91 E : dos_header_->e_magic = IMAGE_DOS_SIGNATURE;
92 : // Set the "DOS File Size" headers.
93 E : dos_header_->e_cblp = dos_header_block_->size() % 512;
94 E : dos_header_->e_cp = dos_header_block_->size() / 512;
95 E : if (dos_header_->e_cblp != 0)
96 E : dos_header_->e_cp++;
97 : // Set the header paragraph size.
98 E : dos_header_->e_cparhdr = dos_header_block_->size() / 16;
99 :
100 E : if (nt_headers_block_ != NULL) {
101 : // Set the NT headers reference.
102 : dos_header_block_->SetReference(
103 : offsetof(IMAGE_DOS_HEADER, e_lfanew),
104 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
105 : sizeof(RelativeAddress),
106 : nt_headers_block_,
107 E : 0, 0));
108 : }
109 E : }
110 :
111 E : void PEUtilsTest::CreateEntryPoints() {
112 : // Setup the main entry-point.
113 : main_entry_point_block_ = block_graph_.AddBlock(
114 E : BlockGraph::CODE_BLOCK, 1, "main_entry_point");
115 E : ASSERT_TRUE(main_entry_point_block_ != NULL);
116 : ASSERT_TRUE(nt_headers_block_->SetReference(
117 : offsetof(IMAGE_NT_HEADERS, OptionalHeader.AddressOfEntryPoint),
118 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
119 : BlockGraph::Reference::kMaximumSize,
120 E : main_entry_point_block_, 0, 0)));
121 :
122 : // Setup the TLS directory.
123 : static const size_t kTlsDirectoryOffset = offsetof(
124 : IMAGE_NT_HEADERS,
125 : OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].VirtualAddress);
126 : BlockGraph::Block* tls_directory_block = block_graph_.AddBlock(
127 : BlockGraph::DATA_BLOCK,
128 : sizeof(IMAGE_TLS_DIRECTORY),
129 E : "tls_directory");
130 E : ASSERT_TRUE(tls_directory_block != NULL);
131 E : ASSERT_TRUE(tls_directory_block->AllocateData(tls_directory_block->size()));
132 : nt_headers_->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].Size =
133 E : tls_directory_block->size();
134 : ASSERT_TRUE(nt_headers_block_->SetReference(
135 : kTlsDirectoryOffset,
136 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
137 : BlockGraph::Reference::kMaximumSize,
138 E : tls_directory_block, 0, 0)));
139 :
140 : // Setup the TLS callbacks table. Reserving enough space for one callback
141 : // and the trailing NULL sentinel.
142 : BlockGraph::Block* tls_callbacks_block = block_graph_.AddBlock(
143 : BlockGraph::DATA_BLOCK,
144 : 2 * BlockGraph::Reference::kMaximumSize,
145 E : "tls_callbacks");
146 E : ASSERT_TRUE(tls_callbacks_block != NULL);
147 E : ASSERT_TRUE(tls_callbacks_block->AllocateData(tls_callbacks_block->size()));
148 : nt_headers_->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].Size =
149 E : tls_directory_block->size();
150 : ASSERT_TRUE(tls_directory_block->SetReference(
151 : offsetof(IMAGE_TLS_DIRECTORY, AddressOfCallBacks),
152 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
153 : BlockGraph::Reference::kMaximumSize,
154 E : tls_callbacks_block, 0, 0)));
155 :
156 : // Add a TLS initializer.
157 : tls_initializer_block_ = block_graph_.AddBlock(
158 E : BlockGraph::CODE_BLOCK, 1, "tls_initializer");
159 E : ASSERT_TRUE(tls_initializer_block_ != NULL);
160 : ASSERT_TRUE(tls_callbacks_block->SetReference(
161 : 0,
162 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
163 : BlockGraph::Reference::kMaximumSize,
164 E : tls_initializer_block_, 0, 0)));
165 E : }
166 :
167 : } // namespace
168 :
169 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockSuccess) {
170 : // This DOS header should test valid.
171 E : EXPECT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
172 E : }
173 :
174 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockNoDataFails) {
175 E : dos_header_block_->SetData(NULL, 0);
176 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
177 E : }
178 :
179 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockTooShortFails) {
180 E : dos_header_block_->ResizeData(sizeof(IMAGE_DOS_HEADER) - 1);
181 E : dos_header_block_->set_size(sizeof(IMAGE_DOS_HEADER) - 1);
182 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
183 E : }
184 :
185 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockInvalidMagicFails) {
186 E : ++dos_header_->e_magic;
187 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
188 E : }
189 :
190 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockInvalidDosFileSizeFails) {
191 E : dos_header_->e_cp = 0;
192 E : dos_header_->e_cblp = 0;
193 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
194 :
195 : // This is invalid, as there are zero pages, and thus no last page.
196 E : dos_header_->e_cblp = 10;
197 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
198 E : }
199 :
200 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockInvalidHeaderSizeFails) {
201 E : --dos_header_->e_cparhdr;
202 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
203 E : }
204 :
205 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockInvalidNTHeaderRefFails) {
206 : // Set the NT headers reference to a non-zero offset.
207 : dos_header_block_->SetReference(
208 : offsetof(IMAGE_DOS_HEADER, e_lfanew),
209 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
210 : sizeof(RelativeAddress),
211 : nt_headers_block_,
212 E : 10, 10));
213 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
214 E : }
215 :
216 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockNoNTHeaderRefFails) {
217 : // Clear the NT headers reference.
218 E : dos_header_block_->RemoveReference(offsetof(IMAGE_DOS_HEADER, e_lfanew));
219 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
220 E : }
221 :
222 E : TEST_F(PEUtilsTest, IsValidNtHeaderBlockSuccess) {
223 : // The NT headers are valid.
224 E : EXPECT_TRUE(IsValidNtHeadersBlock(nt_headers_block_));
225 E : }
226 :
227 E : TEST_F(PEUtilsTest, IsValidNtHeaderBlockInvalidSigFails) {
228 E : ++nt_headers_->Signature;
229 : // Invalid signature.
230 E : EXPECT_FALSE(IsValidNtHeadersBlock(nt_headers_block_));
231 E : }
232 :
233 E : TEST_F(PEUtilsTest, IsValidNtHeaderBlockInvalidOptionalSigFails) {
234 E : ++nt_headers_->OptionalHeader.Magic;
235 : // Invalid signature.
236 E : EXPECT_FALSE(IsValidNtHeadersBlock(nt_headers_block_));
237 E : }
238 :
239 E : TEST_F(PEUtilsTest, IsValidNtHeaderBlockInvalidOptionalSizeFails) {
240 E : ++nt_headers_->FileHeader.SizeOfOptionalHeader;
241 : // Invalid signature.
242 E : EXPECT_FALSE(IsValidNtHeadersBlock(nt_headers_block_));
243 E : }
244 :
245 E : TEST_F(PEUtilsTest, GetNtHeadersBlockFromDosHeaderBlock) {
246 : ASSERT_EQ(nt_headers_block_,
247 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block_));
248 E : }
249 :
250 E : TEST_F(PEUtilsTest, GetNtHeadersBlockFromDosHeaderBlockConst) {
251 : ASSERT_EQ(nt_headers_block_,
252 : GetNtHeadersBlockFromDosHeaderBlock(
253 E : const_cast<const BlockGraph::Block*>(dos_header_block_)));
254 E : }
255 :
256 E : TEST_F(PEUtilsTest, GetExeEntryPoint) {
257 E : EntryPoint entry_point;
258 :
259 : // Get the entry point for an EXE.
260 E : nt_headers_->FileHeader.Characteristics &= ~IMAGE_FILE_DLL;
261 E : EXPECT_TRUE(GetExeEntryPoint(dos_header_block_, &entry_point));
262 E : EXPECT_EQ(main_entry_point_block_, entry_point.first);
263 E : EXPECT_EQ(0, entry_point.second);
264 :
265 : // Should return no entry points if the image is a DLL.
266 E : nt_headers_->FileHeader.Characteristics |= IMAGE_FILE_DLL;
267 E : EXPECT_TRUE(GetExeEntryPoint(dos_header_block_, &entry_point));
268 E : EXPECT_EQ(NULL, entry_point.first);
269 :
270 : // Should fail if the image is an EXE with no entry-point.
271 E : nt_headers_->FileHeader.Characteristics &= ~IMAGE_FILE_DLL;
272 : ASSERT_TRUE(nt_headers_block_->RemoveReference(
273 E : offsetof(IMAGE_NT_HEADERS, OptionalHeader.AddressOfEntryPoint)));
274 E : EXPECT_FALSE(GetExeEntryPoint(dos_header_block_, &entry_point));
275 E : }
276 :
277 E : TEST_F(PEUtilsTest, GetDllEntryPoint) {
278 E : EntryPoint entry_point;
279 :
280 : // Get the DLL entry point.
281 E : nt_headers_->FileHeader.Characteristics |= IMAGE_FILE_DLL;
282 E : EXPECT_TRUE(GetDllEntryPoint(dos_header_block_, &entry_point));
283 E : EXPECT_EQ(main_entry_point_block_, entry_point.first);
284 E : EXPECT_EQ(0, entry_point.second);
285 :
286 : // Should return no entry points if the image is an EXE.
287 E : nt_headers_->FileHeader.Characteristics &= ~IMAGE_FILE_DLL;
288 E : EXPECT_TRUE(GetDllEntryPoint(dos_header_block_, &entry_point));
289 E : EXPECT_EQ(NULL, entry_point.first);
290 :
291 : // Should return no entry points if the image is a DLL without an entry-point.
292 E : nt_headers_->FileHeader.Characteristics |= IMAGE_FILE_DLL;
293 : ASSERT_TRUE(nt_headers_block_->RemoveReference(
294 E : offsetof(IMAGE_NT_HEADERS, OptionalHeader.AddressOfEntryPoint)));
295 E : EXPECT_TRUE(GetDllEntryPoint(dos_header_block_, &entry_point));
296 E : EXPECT_EQ(NULL, entry_point.first);
297 E : }
298 :
299 E : TEST_F(PEUtilsTest, GetTlsInitializers) {
300 : // A container to store the entry-points.
301 E : EntryPointSet entry_points;
302 :
303 : // Get the entry points.
304 E : EXPECT_TRUE(GetTlsInitializers(dos_header_block_, &entry_points));
305 E : EXPECT_EQ(1U, entry_points.size());
306 E : EXPECT_EQ(tls_initializer_block_, entry_points.begin()->first);
307 E : }
308 :
309 E : TEST_F(PEUtilsTest, HasImportEntry) {
310 : // Creates an imported module.
311 E : AddImportsTransform::ImportedModule module("foo.dll");
312 E : const char* kFooFunc = "foo_func";
313 : size_t function_foo = module.AddSymbol(
314 E : kFooFunc, AddImportsTransform::ImportedModule::kAlwaysImport);
315 E : ASSERT_EQ(kFooFunc, module.GetSymbolName(function_foo));
316 :
317 : // Apply the transform to add this module import to the block-graph.
318 E : AddImportsTransform transform;
319 E : transform.AddModule(&module);
320 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
321 E : &transform, &block_graph_, dos_header_block_));
322 :
323 : // Ensure that we can find this module, and that we can't find a
324 : // non-imported module.
325 E : bool has_import = false;
326 E : EXPECT_TRUE(HasImportEntry(dos_header_block_, "foo.dll", &has_import));
327 E : EXPECT_TRUE(has_import);
328 E : EXPECT_TRUE(HasImportEntry(dos_header_block_, "bar.dll", &has_import));
329 E : EXPECT_FALSE(has_import);
330 E : }
331 :
332 : } // namespace pe
|