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 "syzygy/block_graph/typed_block.h"
18 : #include "syzygy/pe/dos_stub.h"
19 :
20 : namespace pe {
21 :
22 : using block_graph::BlockGraph;
23 : using block_graph::ConstTypedBlock;
24 : using block_graph::TypedBlock;
25 : using core::RelativeAddress;
26 :
27 : namespace {
28 :
29 : template <typename BlockPtr>
30 : BlockPtr UncheckedGetNtHeadersBlockFromDosHeaderBlock(
31 E : BlockPtr dos_header_block) {
32 E : BlockGraph::Reference ref;
33 : if (!dos_header_block->GetReference(offsetof(IMAGE_DOS_HEADER, e_lfanew),
34 E : &ref)) {
35 : // No NT headers reference.
36 E : return NULL;
37 : }
38 :
39 : if (ref.offset() != 0 ||
40 : ref.type() != BlockGraph::RELATIVE_REF ||
41 E : ref.size() != sizeof(RelativeAddress)) {
42 : // The reference is of incorrect type.
43 E : return NULL;
44 : }
45 :
46 E : return ref.referenced();
47 E : }
48 :
49 : template <typename BlockPtr>
50 : BlockPtr CheckedGetNtHeadersBlockFromDosHeaderBlock(
51 E : BlockPtr dos_header_block) {
52 E : DCHECK(IsValidDosHeaderBlock(dos_header_block));
53 :
54 : BlockPtr nt_headers_block =
55 E : UncheckedGetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
56 : if (nt_headers_block == NULL ||
57 E : !IsValidNtHeadersBlock(nt_headers_block)) {
58 i : return NULL;
59 : }
60 :
61 E : return nt_headers_block;
62 E : }
63 :
64 : } // namespace
65 :
66 : const char kCodeSectionName[] = ".text";
67 : const char kReadOnlyDataSectionName[] = ".rdata";
68 : const char kReadWriteDataSectionName[] = ".data";
69 : const char kRelocSectionName[] = ".reloc";
70 : const char kResourceSectionName[] = ".rsrc";
71 : const char kTlsSectionName[] = ".tls";
72 :
73 : const DWORD kCodeCharacteristics =
74 : IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE;
75 : const DWORD kReadOnlyDataCharacteristics =
76 : IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ;
77 : const DWORD kReadWriteDataCharacteristics =
78 : IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
79 : const DWORD kRelocCharacteristics =
80 : IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_DISCARDABLE |
81 : IMAGE_SCN_MEM_READ;
82 :
83 E : bool IsValidDosHeaderBlock(const BlockGraph::Block* dos_header_block) {
84 E : ConstTypedBlock<IMAGE_DOS_HEADER> dos_header;
85 :
86 E : if (!dos_header.Init(0, dos_header_block)) {
87 : // Too small or no data.
88 E : return false;
89 : }
90 :
91 E : if (dos_header->e_magic != IMAGE_DOS_SIGNATURE) {
92 : // Wrong signature.
93 E : return false;
94 : }
95 :
96 : // The "DOS file size" is encoded in a rather wonky manner.
97 : // - e_cp is the number of "pages" in the file, but
98 : // - e_cblp is the number of bytes used on the last page.
99 E : size_t dos_file_size = 512 * dos_header->e_cp;
100 E : if (dos_header->e_cblp != 0) {
101 : // Let's not go below zero size.
102 E : if (dos_file_size < 512)
103 E : return false;
104 :
105 E : dos_file_size -= 512;
106 E : dos_file_size += dos_header->e_cblp;
107 : }
108 : // The VC linker yields a DOS header with a size that's larger than
109 : // the DOS header and the NT headers combined. I wonder if anyone cares
110 : // about these sizes anymore.
111 E : if (dos_file_size < dos_header_block->size())
112 E : return false;
113 :
114 : // Check the paragraph size of the header.
115 E : if (dos_header->e_cparhdr * 16 < sizeof(IMAGE_DOS_HEADER))
116 E : return false;
117 :
118 : // Retrieve the NT headers.
119 : const BlockGraph::Block* nt_headers =
120 E : UncheckedGetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
121 E : if (nt_headers == NULL) {
122 : // No DOS header reference.
123 E : return false;
124 : }
125 :
126 E : return true;
127 E : }
128 :
129 E : bool IsValidNtHeadersBlock(const BlockGraph::Block* nt_headers_block) {
130 : // Check the signatures.
131 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
132 :
133 E : if (!nt_headers.Init(0, nt_headers_block)) {
134 : // Short or no data.
135 i : return false;
136 : }
137 :
138 E : if (nt_headers->Signature!= IMAGE_NT_SIGNATURE) {
139 : // Wrong signature.
140 E : return false;
141 : }
142 : if (nt_headers->FileHeader.SizeOfOptionalHeader !=
143 E : sizeof(IMAGE_OPTIONAL_HEADER)) {
144 : // Wrong optional header size.
145 E : return false;
146 : }
147 E : if (nt_headers->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR_MAGIC) {
148 : // Wrong magic for optional header.
149 E : return false;
150 : }
151 :
152 : // Calculate the minimum size for the NT headers and the section header.
153 : size_t header_size = sizeof(IMAGE_NT_HEADERS) +
154 E : sizeof(IMAGE_SECTION_HEADER) * nt_headers->FileHeader.NumberOfSections;
155 :
156 : if (nt_headers_block->size() < header_size ||
157 E : nt_headers_block->data_size() < header_size) {
158 : // The block's size isn't large enough for the section headers.
159 i : return false;
160 : }
161 :
162 E : return true;
163 E : }
164 :
165 : const BlockGraph::Block* GetNtHeadersBlockFromDosHeaderBlock(
166 E : const BlockGraph::Block* dos_header_block) {
167 E : return CheckedGetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
168 E : }
169 :
170 : BlockGraph::Block* GetNtHeadersBlockFromDosHeaderBlock(
171 E : BlockGraph::Block* dos_header_block) {
172 E : return CheckedGetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
173 E : }
174 :
175 E : bool UpdateDosHeader(BlockGraph::Block* dos_header_block) {
176 E : DCHECK(dos_header_block != NULL);
177 :
178 : // The DOS header has to be a multiple of 16 bytes for historic reasons.
179 : size_t dos_header_size = common::AlignUp(
180 E : sizeof(IMAGE_DOS_HEADER) + pe::kDosStubSize, 16);
181 :
182 : // If the new header block is shorter than it was, go ahead and
183 : // trim the source ranges to match the new, shorter size.
184 E : if (dos_header_block->size() > dos_header_size) {
185 : BlockGraph::Block::DataRange range(
186 E : dos_header_size, dos_header_block->size() - dos_header_size);
187 E : dos_header_block->source_ranges().RemoveMappedRange(range);
188 : }
189 :
190 E : dos_header_block->ResizeData(dos_header_size);
191 E : dos_header_block->set_size(dos_header_size);
192 E : DCHECK_EQ(dos_header_size, dos_header_block->size());
193 E : DCHECK_EQ(dos_header_size, dos_header_block->data_size());
194 :
195 E : TypedBlock<IMAGE_DOS_HEADER> dos_header;
196 E : if (!dos_header.InitWithSize(0, dos_header_size, dos_header_block)) {
197 i : LOG(ERROR) << "Unable to cast IMAGE_DOS_HEADER.";
198 i : return false;
199 : }
200 :
201 : // Wipe the DOS header and fill in the stub.
202 E : memset(dos_header.Get(), 0, sizeof(IMAGE_DOS_HEADER));
203 E : memcpy(dos_header.Get() + 1, pe::kDosStub, pe::kDosStubSize);
204 :
205 E : dos_header->e_magic = IMAGE_DOS_SIGNATURE;
206 : // Calculate the number of bytes used on the last DOS executable "page".
207 E : dos_header->e_cblp = dos_header_size % 512;
208 : // Calculate the number of pages used by the DOS executable.
209 E : dos_header->e_cp = dos_header_size / 512;
210 : // Count the last page if we didn't have an even multiple
211 E : if (dos_header->e_cblp != 0)
212 E : dos_header->e_cp++;
213 :
214 : // Header length in "paragraphs".
215 E : dos_header->e_cparhdr = sizeof(*dos_header) / 16;
216 :
217 : // Set this to max allowed, just because.
218 E : dos_header->e_maxalloc = 0xFFFF;
219 :
220 : // Location of relocs - our header has zero relocs, but we set this anyway.
221 E : dos_header->e_lfarlc = sizeof(*dos_header);
222 :
223 E : DCHECK(IsValidDosHeaderBlock(dos_header_block));
224 :
225 E : return true;
226 E : }
227 :
228 : namespace {
229 :
230 : enum EntryPointTypeEnum { kExeEntryPoint, kDllEntryPoint };
231 :
232 : bool GetImageEntryPoint(BlockGraph::Block* dos_header_block,
233 : EntryPointTypeEnum desired_entry_point_type,
234 E : EntryPoint* entry_point) {
235 E : DCHECK(dos_header_block != NULL);
236 E : DCHECK(entry_point != NULL);
237 :
238 E : *entry_point = EntryPoint(static_cast<BlockGraph::Block*>(NULL), 0);
239 :
240 : BlockGraph::Block* nt_headers_block =
241 E : pe::GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
242 :
243 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
244 E : if (nt_headers_block == NULL || !nt_headers.Init(0, nt_headers_block)) {
245 i : LOG(ERROR) << "Unable to retrieve NT Headers.";
246 i : return false;
247 : }
248 :
249 E : EntryPointTypeEnum entry_point_type = kExeEntryPoint;
250 E : if ((nt_headers->FileHeader.Characteristics & IMAGE_FILE_DLL) != 0)
251 E : entry_point_type = kDllEntryPoint;
252 :
253 E : if (entry_point_type != desired_entry_point_type)
254 E : return true;
255 :
256 E : BlockGraph::Reference entry_point_ref;
257 : bool found = nt_headers.block()->GetReference(
258 : offsetof(IMAGE_NT_HEADERS, OptionalHeader.AddressOfEntryPoint),
259 E : &entry_point_ref);
260 :
261 E : if (!found && entry_point_type == kExeEntryPoint) {
262 E : LOG(ERROR) << "Malformed PE Headers: No entry point found for executable.";
263 E : return false;
264 : }
265 :
266 E : if (found) {
267 : *entry_point = EntryPoint(entry_point_ref.referenced(),
268 E : entry_point_ref.offset());
269 : }
270 :
271 E : return true;
272 E : }
273 :
274 : } // namespace
275 :
276 : bool GetExeEntryPoint(BlockGraph::Block* dos_header_block,
277 E : EntryPoint* entry_point) {
278 E : return GetImageEntryPoint(dos_header_block, kExeEntryPoint, entry_point);
279 E : }
280 :
281 : bool GetDllEntryPoint(BlockGraph::Block* dos_header_block,
282 E : EntryPoint* entry_point) {
283 E : return GetImageEntryPoint(dos_header_block, kDllEntryPoint, entry_point);
284 E : }
285 :
286 : bool GetTlsInitializers(BlockGraph::Block* dos_header_block,
287 E : EntryPointSet* entry_points) {
288 E : DCHECK(dos_header_block != NULL);
289 E : DCHECK(entry_points != NULL);
290 :
291 : BlockGraph::Block* nt_headers_block =
292 E : pe::GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
293 :
294 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
295 E : if (nt_headers_block == NULL || !nt_headers.Init(0, nt_headers_block)) {
296 i : LOG(ERROR) << "Unable to retrieve NT Headers.";
297 i : return false;
298 : }
299 :
300 : // If the module has no TLS directory then there are no TLS initializers
301 : // and hence nothing to do.
302 : const IMAGE_DATA_DIRECTORY& data_dir =
303 E : nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS];
304 E : if (data_dir.Size == 0 || !nt_headers.HasReference(data_dir.VirtualAddress)) {
305 E : return true;
306 : }
307 :
308 : // Find the TLS directory.
309 E : TypedBlock<IMAGE_TLS_DIRECTORY> tls_dir;
310 E : if (!nt_headers.Dereference(data_dir.VirtualAddress, &tls_dir)) {
311 i : LOG(ERROR) << "Failed to cast TLS directory.";
312 i : return false;
313 : }
314 :
315 : // Get the TLS initializer callbacks. We manually lookup the reference
316 : // because it is an indirect reference, which can't be dereferenced by
317 : // TypedBlock.
318 : typedef BlockGraph::Block::ReferenceMap ReferenceMap;
319 : ReferenceMap::const_iterator callback_ref =
320 : tls_dir.block()->references().find(
321 E : tls_dir.OffsetOf(tls_dir->AddressOfCallBacks));
322 E : if (callback_ref == tls_dir.block()->references().end()) {
323 i : LOG(ERROR) << "Failed to locate TLS initializers.";
324 i : return false;
325 : }
326 :
327 : // Note each of the TLS entry points.
328 E : const BlockGraph::Block* callbacks_block = callback_ref->second.referenced();
329 E : const ReferenceMap& ref_map = callbacks_block->references();
330 E : ReferenceMap::const_iterator iter = ref_map.begin();
331 E : for (; iter != ref_map.end(); ++iter) {
332 E : const BlockGraph::Reference& ref = iter->second;
333 E : DCHECK(ref.size() == sizeof(core::AbsoluteAddress));
334 : entry_points->insert(
335 E : std::make_pair(ref.referenced(), ref.offset()));
336 E : }
337 :
338 E : return true;
339 E : }
340 :
341 : } // namespace pe
|