1 // Copyright 2014 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "base/win/pe_image_reader.h"
6
7 #include <wintrust.h>
8
9 #include <memory>
10
11 #include "base/check_op.h"
12 #include "base/macros.h"
13 #include "base/numerics/safe_math.h"
14
15 namespace base {
16 namespace win {
17
18 // A class template of traits pertaining to IMAGE_OPTIONAL_HEADER{32,64}.
19 template <class HEADER_TYPE>
20 struct OptionalHeaderTraits {};
21
22 template <>
23 struct OptionalHeaderTraits<IMAGE_OPTIONAL_HEADER32> {
24 static const PeImageReader::WordSize word_size = PeImageReader::WORD_SIZE_32;
25 };
26
27 template <>
28 struct OptionalHeaderTraits<IMAGE_OPTIONAL_HEADER64> {
29 static const PeImageReader::WordSize word_size = PeImageReader::WORD_SIZE_64;
30 };
31
32 // A template for type-specific optional header implementations. This, in
33 // conjunction with the OptionalHeader interface, effectively erases the
34 // underlying structure type from the point of view of the PeImageReader.
35 template <class OPTIONAL_HEADER_TYPE>
36 class PeImageReader::OptionalHeaderImpl : public PeImageReader::OptionalHeader {
37 public:
38 using TraitsType = OptionalHeaderTraits<OPTIONAL_HEADER_TYPE>;
39
OptionalHeaderImpl(const uint8_t * optional_header_start)40 explicit OptionalHeaderImpl(const uint8_t* optional_header_start)
41 : optional_header_(reinterpret_cast<const OPTIONAL_HEADER_TYPE*>(
42 optional_header_start)) {}
43
GetWordSize()44 WordSize GetWordSize() override { return TraitsType::word_size; }
45
GetDataDirectoryOffset()46 size_t GetDataDirectoryOffset() override {
47 return offsetof(OPTIONAL_HEADER_TYPE, DataDirectory);
48 }
49
GetDataDirectorySize()50 DWORD GetDataDirectorySize() override {
51 return optional_header_->NumberOfRvaAndSizes;
52 }
53
GetDataDirectoryEntries()54 const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntries() override {
55 return &optional_header_->DataDirectory[0];
56 }
57
GetSizeOfImage()58 DWORD GetSizeOfImage() override { return optional_header_->SizeOfImage; }
59
60 private:
61 const OPTIONAL_HEADER_TYPE* optional_header_;
62 DISALLOW_COPY_AND_ASSIGN(OptionalHeaderImpl);
63 };
64
PeImageReader()65 PeImageReader::PeImageReader() {}
66
~PeImageReader()67 PeImageReader::~PeImageReader() {
68 Clear();
69 }
70
Initialize(const uint8_t * image_data,size_t image_size)71 bool PeImageReader::Initialize(const uint8_t* image_data, size_t image_size) {
72 image_data_ = image_data;
73 image_size_ = image_size;
74
75 if (!ValidateDosHeader() || !ValidatePeSignature() ||
76 !ValidateCoffFileHeader() || !ValidateOptionalHeader() ||
77 !ValidateSectionHeaders()) {
78 Clear();
79 return false;
80 }
81
82 return true;
83 }
84
GetWordSize()85 PeImageReader::WordSize PeImageReader::GetWordSize() {
86 return optional_header_->GetWordSize();
87 }
88
GetDosHeader()89 const IMAGE_DOS_HEADER* PeImageReader::GetDosHeader() {
90 DCHECK_NE((validation_state_ & VALID_DOS_HEADER), 0U);
91 return reinterpret_cast<const IMAGE_DOS_HEADER*>(image_data_);
92 }
93
GetCoffFileHeader()94 const IMAGE_FILE_HEADER* PeImageReader::GetCoffFileHeader() {
95 DCHECK_NE((validation_state_ & VALID_COFF_FILE_HEADER), 0U);
96 return reinterpret_cast<const IMAGE_FILE_HEADER*>(
97 image_data_ + GetDosHeader()->e_lfanew + sizeof(DWORD));
98 }
99
GetOptionalHeaderData(size_t * optional_header_size)100 const uint8_t* PeImageReader::GetOptionalHeaderData(
101 size_t* optional_header_size) {
102 *optional_header_size = GetOptionalHeaderSize();
103 return GetOptionalHeaderStart();
104 }
105
GetNumberOfSections()106 size_t PeImageReader::GetNumberOfSections() {
107 return GetCoffFileHeader()->NumberOfSections;
108 }
109
GetSectionHeaderAt(size_t index)110 const IMAGE_SECTION_HEADER* PeImageReader::GetSectionHeaderAt(size_t index) {
111 DCHECK_NE((validation_state_ & VALID_SECTION_HEADERS), 0U);
112 DCHECK_LT(index, GetNumberOfSections());
113 return reinterpret_cast<const IMAGE_SECTION_HEADER*>(
114 GetOptionalHeaderStart() + GetOptionalHeaderSize() +
115 (sizeof(IMAGE_SECTION_HEADER) * index));
116 }
117
GetExportSection(size_t * section_size)118 const uint8_t* PeImageReader::GetExportSection(size_t* section_size) {
119 size_t data_size = 0;
120 const uint8_t* data = GetImageData(IMAGE_DIRECTORY_ENTRY_EXPORT, &data_size);
121
122 // The export section data must be big enough for the export directory.
123 if (!data || data_size < sizeof(IMAGE_EXPORT_DIRECTORY))
124 return nullptr;
125
126 *section_size = data_size;
127 return data;
128 }
129
GetNumberOfDebugEntries()130 size_t PeImageReader::GetNumberOfDebugEntries() {
131 size_t data_size = 0;
132 const uint8_t* data = GetImageData(IMAGE_DIRECTORY_ENTRY_DEBUG, &data_size);
133 return data ? (data_size / sizeof(IMAGE_DEBUG_DIRECTORY)) : 0;
134 }
135
GetDebugEntry(size_t index,const uint8_t ** raw_data,size_t * raw_data_size)136 const IMAGE_DEBUG_DIRECTORY* PeImageReader::GetDebugEntry(
137 size_t index,
138 const uint8_t** raw_data,
139 size_t* raw_data_size) {
140 DCHECK_LT(index, GetNumberOfDebugEntries());
141
142 // Get the debug directory.
143 size_t debug_directory_size = 0;
144 const IMAGE_DEBUG_DIRECTORY* entries =
145 reinterpret_cast<const IMAGE_DEBUG_DIRECTORY*>(
146 GetImageData(IMAGE_DIRECTORY_ENTRY_DEBUG, &debug_directory_size));
147 if (!entries)
148 return nullptr;
149
150 const IMAGE_DEBUG_DIRECTORY& entry = entries[index];
151 const uint8_t* debug_data = nullptr;
152 if (GetStructureAt(entry.PointerToRawData, entry.SizeOfData, &debug_data)) {
153 *raw_data = debug_data;
154 *raw_data_size = entry.SizeOfData;
155 }
156 return &entry;
157 }
158
EnumCertificates(EnumCertificatesCallback callback,void * context)159 bool PeImageReader::EnumCertificates(EnumCertificatesCallback callback,
160 void* context) {
161 size_t data_size = 0;
162 const uint8_t* data =
163 GetImageData(IMAGE_DIRECTORY_ENTRY_SECURITY, &data_size);
164 if (!data)
165 return false; // Certificate table is out of bounds.
166 const size_t kWinCertificateSize = offsetof(WIN_CERTIFICATE, bCertificate);
167 while (data_size) {
168 const WIN_CERTIFICATE* win_certificate =
169 reinterpret_cast<const WIN_CERTIFICATE*>(data);
170 if (kWinCertificateSize > data_size ||
171 kWinCertificateSize > win_certificate->dwLength ||
172 win_certificate->dwLength > data_size) {
173 return false;
174 }
175 if (!(*callback)(
176 win_certificate->wRevision, win_certificate->wCertificateType,
177 &win_certificate->bCertificate[0],
178 win_certificate->dwLength - kWinCertificateSize, context)) {
179 return false;
180 }
181 size_t padded_length = (win_certificate->dwLength + 7) & ~0x7;
182 // Don't overflow when recalculating data_size, since padded_length can be
183 // attacker controlled.
184 if (!CheckSub(data_size, padded_length).AssignIfValid(&data_size))
185 return false;
186 data += padded_length;
187 }
188 return true;
189 }
190
GetSizeOfImage()191 DWORD PeImageReader::GetSizeOfImage() {
192 return optional_header_->GetSizeOfImage();
193 }
194
Clear()195 void PeImageReader::Clear() {
196 image_data_ = nullptr;
197 image_size_ = 0;
198 validation_state_ = 0;
199 optional_header_.reset();
200 }
201
ValidateDosHeader()202 bool PeImageReader::ValidateDosHeader() {
203 const IMAGE_DOS_HEADER* dos_header = nullptr;
204 if (!GetStructureAt(0, &dos_header) ||
205 dos_header->e_magic != IMAGE_DOS_SIGNATURE || dos_header->e_lfanew < 0) {
206 return false;
207 }
208
209 validation_state_ |= VALID_DOS_HEADER;
210 return true;
211 }
212
ValidatePeSignature()213 bool PeImageReader::ValidatePeSignature() {
214 const DWORD* signature = nullptr;
215 if (!GetStructureAt(GetDosHeader()->e_lfanew, &signature) ||
216 *signature != IMAGE_NT_SIGNATURE) {
217 return false;
218 }
219
220 validation_state_ |= VALID_PE_SIGNATURE;
221 return true;
222 }
223
ValidateCoffFileHeader()224 bool PeImageReader::ValidateCoffFileHeader() {
225 DCHECK_NE((validation_state_ & VALID_PE_SIGNATURE), 0U);
226 const IMAGE_FILE_HEADER* file_header = nullptr;
227 if (!GetStructureAt(
228 GetDosHeader()->e_lfanew + offsetof(IMAGE_NT_HEADERS32, FileHeader),
229 &file_header)) {
230 return false;
231 }
232
233 validation_state_ |= VALID_COFF_FILE_HEADER;
234 return true;
235 }
236
ValidateOptionalHeader()237 bool PeImageReader::ValidateOptionalHeader() {
238 const IMAGE_FILE_HEADER* file_header = GetCoffFileHeader();
239 const size_t optional_header_offset =
240 GetDosHeader()->e_lfanew + offsetof(IMAGE_NT_HEADERS32, OptionalHeader);
241 const size_t optional_header_size = file_header->SizeOfOptionalHeader;
242 const WORD* optional_header_magic = nullptr;
243
244 if (optional_header_size < sizeof(*optional_header_magic) ||
245 !GetStructureAt(optional_header_offset, &optional_header_magic)) {
246 return false;
247 }
248
249 std::unique_ptr<OptionalHeader> optional_header;
250 if (*optional_header_magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
251 optional_header =
252 std::make_unique<OptionalHeaderImpl<IMAGE_OPTIONAL_HEADER32>>(
253 image_data_ + optional_header_offset);
254 } else if (*optional_header_magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
255 optional_header =
256 std::make_unique<OptionalHeaderImpl<IMAGE_OPTIONAL_HEADER64>>(
257 image_data_ + optional_header_offset);
258 } else {
259 return false;
260 }
261
262 // Does all of the claimed optional header fit in the image?
263 if (optional_header_size > image_size_ - optional_header_offset)
264 return false;
265
266 // Is the claimed optional header big enough for everything but the dir?
267 if (optional_header->GetDataDirectoryOffset() > optional_header_size)
268 return false;
269
270 // Is there enough room for all of the claimed directory entries?
271 if (optional_header->GetDataDirectorySize() >
272 ((optional_header_size - optional_header->GetDataDirectoryOffset()) /
273 sizeof(IMAGE_DATA_DIRECTORY))) {
274 return false;
275 }
276
277 optional_header_.swap(optional_header);
278 validation_state_ |= VALID_OPTIONAL_HEADER;
279 return true;
280 }
281
ValidateSectionHeaders()282 bool PeImageReader::ValidateSectionHeaders() {
283 const uint8_t* first_section_header =
284 GetOptionalHeaderStart() + GetOptionalHeaderSize();
285 const size_t number_of_sections = GetNumberOfSections();
286
287 // Do all section headers fit in the image?
288 if (!GetStructureAt(first_section_header - image_data_,
289 number_of_sections * sizeof(IMAGE_SECTION_HEADER),
290 &first_section_header)) {
291 return false;
292 }
293
294 validation_state_ |= VALID_SECTION_HEADERS;
295 return true;
296 }
297
GetOptionalHeaderStart()298 const uint8_t* PeImageReader::GetOptionalHeaderStart() {
299 DCHECK_NE((validation_state_ & VALID_OPTIONAL_HEADER), 0U);
300 return (image_data_ + GetDosHeader()->e_lfanew +
301 offsetof(IMAGE_NT_HEADERS32, OptionalHeader));
302 }
303
GetOptionalHeaderSize()304 size_t PeImageReader::GetOptionalHeaderSize() {
305 return GetCoffFileHeader()->SizeOfOptionalHeader;
306 }
307
GetDataDirectoryEntryAt(size_t index)308 const IMAGE_DATA_DIRECTORY* PeImageReader::GetDataDirectoryEntryAt(
309 size_t index) {
310 DCHECK_NE((validation_state_ & VALID_OPTIONAL_HEADER), 0U);
311 if (index >= optional_header_->GetDataDirectorySize())
312 return nullptr;
313 return &optional_header_->GetDataDirectoryEntries()[index];
314 }
315
FindSectionFromRva(uint32_t relative_address)316 const IMAGE_SECTION_HEADER* PeImageReader::FindSectionFromRva(
317 uint32_t relative_address) {
318 const size_t number_of_sections = GetNumberOfSections();
319 for (size_t i = 0; i < number_of_sections; ++i) {
320 const IMAGE_SECTION_HEADER* section_header = GetSectionHeaderAt(i);
321 // Is the raw data present in the image? If no, optimistically keep looking.
322 const uint8_t* section_data = nullptr;
323 if (!GetStructureAt(section_header->PointerToRawData,
324 section_header->SizeOfRawData, §ion_data)) {
325 continue;
326 }
327 // Does the RVA lie on or after this section's start when mapped? If no,
328 // bail.
329 if (section_header->VirtualAddress > relative_address)
330 break;
331 // Does the RVA lie within the section when mapped? If no, keep looking.
332 size_t address_offset = relative_address - section_header->VirtualAddress;
333 if (address_offset > section_header->Misc.VirtualSize)
334 continue;
335 // We have a winner.
336 return section_header;
337 }
338 return nullptr;
339 }
340
GetImageData(size_t index,size_t * data_length)341 const uint8_t* PeImageReader::GetImageData(size_t index, size_t* data_length) {
342 // Get the requested directory entry.
343 const IMAGE_DATA_DIRECTORY* entry = GetDataDirectoryEntryAt(index);
344 if (!entry)
345 return nullptr;
346
347 // The entry for the certificate table is special in that its address is a
348 // file pointer rather than an RVA.
349 if (index == IMAGE_DIRECTORY_ENTRY_SECURITY) {
350 // Does the data fit within the file.
351 if (entry->VirtualAddress > image_size_ ||
352 image_size_ - entry->VirtualAddress < entry->Size) {
353 return nullptr;
354 }
355 *data_length = entry->Size;
356 return image_data_ + entry->VirtualAddress;
357 }
358
359 // Find the section containing the data.
360 const IMAGE_SECTION_HEADER* header =
361 FindSectionFromRva(entry->VirtualAddress);
362 if (!header)
363 return nullptr;
364
365 // Does the data fit within the section when mapped?
366 size_t data_offset = entry->VirtualAddress - header->VirtualAddress;
367 if (entry->Size > (header->Misc.VirtualSize - data_offset))
368 return nullptr;
369
370 // Is the data entirely present on disk (if not it's zeroed out when loaded)?
371 if (data_offset >= header->SizeOfRawData ||
372 header->SizeOfRawData - data_offset < entry->Size) {
373 return nullptr;
374 }
375
376 *data_length = entry->Size;
377 return image_data_ + header->PointerToRawData + data_offset;
378 }
379
380 } // namespace win
381 } // namespace base
382