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