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 "chrome/browser/safe_browsing/pe_image_reader_win.h"
6
7#include "base/logging.h"
8
9namespace safe_browsing {
10
11// A class template of traits pertaining to IMAGE_OPTIONAL_HEADER{32,64}.
12template<class HEADER_TYPE>
13struct OptionalHeaderTraits {
14};
15
16template<>
17struct OptionalHeaderTraits<IMAGE_OPTIONAL_HEADER32> {
18  static const PeImageReader::WordSize word_size = PeImageReader::WORD_SIZE_32;
19};
20
21template<>
22struct OptionalHeaderTraits<IMAGE_OPTIONAL_HEADER64> {
23  static const PeImageReader::WordSize word_size = PeImageReader::WORD_SIZE_64;
24};
25
26// A template for type-specific optional header implementations. This, in
27// conjunction with the OptionalHeader interface, effectively erases the
28// underlying structure type from the point of view of the PeImageReader.
29template<class OPTIONAL_HEADER_TYPE>
30class PeImageReader::OptionalHeaderImpl : public PeImageReader::OptionalHeader {
31 public:
32  typedef OptionalHeaderTraits<OPTIONAL_HEADER_TYPE> TraitsType;
33
34  explicit OptionalHeaderImpl(const uint8_t* optional_header_start)
35      : optional_header_(reinterpret_cast<const OPTIONAL_HEADER_TYPE*>(
36                             optional_header_start)) {}
37
38  virtual WordSize GetWordSize() OVERRIDE {
39    return TraitsType::word_size;
40  }
41
42  virtual size_t GetDataDirectoryOffset() OVERRIDE {
43    return offsetof(OPTIONAL_HEADER_TYPE, DataDirectory);
44  }
45
46  virtual DWORD GetDataDirectorySize() OVERRIDE {
47    return optional_header_->NumberOfRvaAndSizes;
48  }
49
50  virtual const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntries() OVERRIDE {
51    return &optional_header_->DataDirectory[0];
52  }
53
54 private:
55  const OPTIONAL_HEADER_TYPE* optional_header_;
56  DISALLOW_COPY_AND_ASSIGN(OptionalHeaderImpl);
57};
58
59PeImageReader::PeImageReader()
60    : image_data_(),
61      image_size_(),
62      validation_state_() {}
63
64PeImageReader::~PeImageReader() {
65  Clear();
66}
67
68bool PeImageReader::Initialize(const uint8_t* image_data, size_t image_size) {
69  image_data_ = image_data;
70  image_size_ = image_size;
71
72  if (!ValidateDosHeader() ||
73      !ValidatePeSignature() ||
74      !ValidateCoffFileHeader() ||
75      !ValidateOptionalHeader() ||
76      !ValidateSectionHeaders()) {
77    Clear();
78    return false;
79  }
80
81  return true;
82}
83
84PeImageReader::WordSize PeImageReader::GetWordSize() {
85  return optional_header_->GetWordSize();
86}
87
88const IMAGE_DOS_HEADER* PeImageReader::GetDosHeader() {
89  DCHECK_NE((validation_state_ & VALID_DOS_HEADER), 0U);
90  return reinterpret_cast<const IMAGE_DOS_HEADER*>(image_data_);
91}
92
93const IMAGE_FILE_HEADER* PeImageReader::GetCoffFileHeader() {
94  DCHECK_NE((validation_state_ & VALID_COFF_FILE_HEADER), 0U);
95  return reinterpret_cast<const IMAGE_FILE_HEADER*>(
96      image_data_ + GetDosHeader()->e_lfanew + sizeof(DWORD));
97}
98
99const uint8_t* PeImageReader::GetOptionalHeaderData(
100    size_t* optional_header_size) {
101  *optional_header_size = GetOptionalHeaderSize();
102  return GetOptionalHeaderStart();
103}
104
105size_t PeImageReader::GetNumberOfSections() {
106  return GetCoffFileHeader()->NumberOfSections;
107}
108
109const IMAGE_SECTION_HEADER* PeImageReader::GetSectionHeaderAt(size_t index) {
110  DCHECK_NE((validation_state_ & VALID_SECTION_HEADERS), 0U);
111  DCHECK_LT(index, GetNumberOfSections());
112  return reinterpret_cast<const IMAGE_SECTION_HEADER*>(
113      GetOptionalHeaderStart() +
114      GetOptionalHeaderSize() +
115      (sizeof(IMAGE_SECTION_HEADER) * index));
116}
117
118const 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 NULL;
125
126  *section_size = data_size;
127  return data;
128}
129
130size_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
136const 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 NULL;
149
150  const IMAGE_DEBUG_DIRECTORY& entry = entries[index];
151  const uint8_t* debug_data = NULL;
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
159void PeImageReader::Clear() {
160  image_data_ = NULL;
161  image_size_ = 0;
162  validation_state_ = 0;
163  optional_header_.reset();
164}
165
166bool PeImageReader::ValidateDosHeader() {
167  const IMAGE_DOS_HEADER* dos_header = NULL;
168  if (!GetStructureAt(0, &dos_header) ||
169      dos_header->e_magic != IMAGE_DOS_SIGNATURE ||
170      dos_header->e_lfanew < 0) {
171    return false;
172  }
173
174  validation_state_ |= VALID_DOS_HEADER;
175  return true;
176}
177
178bool PeImageReader::ValidatePeSignature() {
179  const DWORD* signature = NULL;
180  if (!GetStructureAt(GetDosHeader()->e_lfanew, &signature) ||
181      *signature != IMAGE_NT_SIGNATURE) {
182    return false;
183  }
184
185  validation_state_ |= VALID_PE_SIGNATURE;
186  return true;
187}
188
189bool PeImageReader::ValidateCoffFileHeader() {
190  DCHECK_NE((validation_state_ & VALID_PE_SIGNATURE), 0U);
191  const IMAGE_FILE_HEADER* file_header = NULL;
192  if (!GetStructureAt(GetDosHeader()->e_lfanew +
193                          offsetof(IMAGE_NT_HEADERS32, FileHeader),
194                      &file_header)) {
195    return false;
196  }
197
198  validation_state_ |= VALID_COFF_FILE_HEADER;
199  return true;
200}
201
202bool PeImageReader::ValidateOptionalHeader() {
203  const IMAGE_FILE_HEADER* file_header = GetCoffFileHeader();
204  const size_t optional_header_offset =
205      GetDosHeader()->e_lfanew + offsetof(IMAGE_NT_HEADERS32, OptionalHeader);
206  const size_t optional_header_size = file_header->SizeOfOptionalHeader;
207  const WORD* optional_header_magic = NULL;
208
209  if (optional_header_size < sizeof(*optional_header_magic) ||
210      !GetStructureAt(optional_header_offset, &optional_header_magic)) {
211    return false;
212  }
213
214  scoped_ptr<OptionalHeader> optional_header;
215  if (*optional_header_magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
216    optional_header.reset(new OptionalHeaderImpl<IMAGE_OPTIONAL_HEADER32>(
217        image_data_ + optional_header_offset));
218  } else if (*optional_header_magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
219    optional_header.reset(new OptionalHeaderImpl<IMAGE_OPTIONAL_HEADER64>(
220        image_data_ + optional_header_offset));
221  } else {
222    return false;
223  }
224
225  // Does all of the claimed optional header fit in the image?
226  if (optional_header_size > image_size_ - optional_header_offset)
227    return false;
228
229  // Is the claimed optional header big enough for everything but the dir?
230  if (optional_header->GetDataDirectoryOffset() > optional_header_size)
231    return false;
232
233  // Is there enough room for all of the claimed directory entries?
234  if (optional_header->GetDataDirectorySize() >
235      ((optional_header_size - optional_header->GetDataDirectoryOffset()) /
236       sizeof(IMAGE_DATA_DIRECTORY))) {
237    return false;
238  }
239
240  optional_header_.swap(optional_header);
241  validation_state_ |= VALID_OPTIONAL_HEADER;
242  return true;
243}
244
245bool PeImageReader::ValidateSectionHeaders() {
246  const uint8_t* first_section_header =
247      GetOptionalHeaderStart() + GetOptionalHeaderSize();
248  const size_t number_of_sections = GetNumberOfSections();
249
250  // Do all section headers fit in the image?
251  if (!GetStructureAt(first_section_header - image_data_,
252                      number_of_sections * sizeof(IMAGE_SECTION_HEADER),
253                      &first_section_header)) {
254    return false;
255  }
256
257  validation_state_ |= VALID_SECTION_HEADERS;
258  return true;
259}
260
261const uint8_t* PeImageReader::GetOptionalHeaderStart() {
262  DCHECK_NE((validation_state_ & VALID_OPTIONAL_HEADER), 0U);
263  return (image_data_ +
264          GetDosHeader()->e_lfanew +
265          offsetof(IMAGE_NT_HEADERS32, OptionalHeader));
266}
267
268size_t PeImageReader::GetOptionalHeaderSize() {
269  return GetCoffFileHeader()->SizeOfOptionalHeader;
270}
271
272const IMAGE_DATA_DIRECTORY* PeImageReader::GetDataDirectoryEntryAt(
273    size_t index) {
274  DCHECK_NE((validation_state_ & VALID_OPTIONAL_HEADER), 0U);
275  if (index >= optional_header_->GetDataDirectorySize())
276    return NULL;
277  return &optional_header_->GetDataDirectoryEntries()[index];
278}
279
280const IMAGE_SECTION_HEADER* PeImageReader::FindSectionFromRva(
281    uint32_t relative_address) {
282  const size_t number_of_sections = GetNumberOfSections();
283  for (size_t i = 0; i < number_of_sections; ++i) {
284    const IMAGE_SECTION_HEADER* section_header = GetSectionHeaderAt(i);
285    // Is the raw data present in the image? If no, optimistically keep looking.
286    const uint8_t* section_data = NULL;
287    if (!GetStructureAt(section_header->PointerToRawData,
288                        section_header->SizeOfRawData,
289                        &section_data)) {
290      continue;
291    }
292    // Does the RVA lie on or after this section's start when mapped? If no,
293    // bail.
294    if (section_header->VirtualAddress > relative_address)
295      break;
296    // Does the RVA lie within the section when mapped? If no, keep looking.
297    size_t address_offset = relative_address - section_header->VirtualAddress;
298    if (address_offset > section_header->Misc.VirtualSize)
299      continue;
300    // We have a winner.
301    return section_header;
302  }
303  return NULL;
304}
305
306const uint8_t* PeImageReader::GetImageData(size_t index, size_t* data_length) {
307  // Get the requested directory entry.
308  const IMAGE_DATA_DIRECTORY* entry = GetDataDirectoryEntryAt(index);
309  if (!entry)
310    return NULL;
311
312  // Find the section containing the data.
313  const IMAGE_SECTION_HEADER* header =
314      FindSectionFromRva(entry->VirtualAddress);
315  if (!header)
316    return NULL;
317
318  // Does the data fit within the section when mapped?
319  size_t data_offset = entry->VirtualAddress - header->VirtualAddress;
320  if (entry->Size > (header->Misc.VirtualSize - data_offset))
321    return NULL;
322
323  // Is the data entirely present on disk (if not it's zeroed out when loaded)?
324  if (data_offset >= header->SizeOfRawData ||
325      header->SizeOfRawData - data_offset < entry->Size) {
326    return NULL;
327  }
328
329  *data_length = entry->Size;
330  return image_data_ + header->PointerToRawData + data_offset;
331}
332
333}  // namespace safe_browsing
334