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