• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 
9 namespace safe_browsing {
10 
11 // A class template of traits pertaining to IMAGE_OPTIONAL_HEADER{32,64}.
12 template<class HEADER_TYPE>
13 struct OptionalHeaderTraits {
14 };
15 
16 template<>
17 struct OptionalHeaderTraits<IMAGE_OPTIONAL_HEADER32> {
18   static const PeImageReader::WordSize word_size = PeImageReader::WORD_SIZE_32;
19 };
20 
21 template<>
22 struct 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.
29 template<class OPTIONAL_HEADER_TYPE>
30 class PeImageReader::OptionalHeaderImpl : public PeImageReader::OptionalHeader {
31  public:
32   typedef OptionalHeaderTraits<OPTIONAL_HEADER_TYPE> TraitsType;
33 
OptionalHeaderImpl(const uint8_t * optional_header_start)34   explicit OptionalHeaderImpl(const uint8_t* optional_header_start)
35       : optional_header_(reinterpret_cast<const OPTIONAL_HEADER_TYPE*>(
36                              optional_header_start)) {}
37 
GetWordSize()38   virtual WordSize GetWordSize() OVERRIDE {
39     return TraitsType::word_size;
40   }
41 
GetDataDirectoryOffset()42   virtual size_t GetDataDirectoryOffset() OVERRIDE {
43     return offsetof(OPTIONAL_HEADER_TYPE, DataDirectory);
44   }
45 
GetDataDirectorySize()46   virtual DWORD GetDataDirectorySize() OVERRIDE {
47     return optional_header_->NumberOfRvaAndSizes;
48   }
49 
GetDataDirectoryEntries()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 
PeImageReader()59 PeImageReader::PeImageReader()
60     : image_data_(),
61       image_size_(),
62       validation_state_() {}
63 
~PeImageReader()64 PeImageReader::~PeImageReader() {
65   Clear();
66 }
67 
Initialize(const uint8_t * image_data,size_t image_size)68 bool 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 
GetWordSize()84 PeImageReader::WordSize PeImageReader::GetWordSize() {
85   return optional_header_->GetWordSize();
86 }
87 
GetDosHeader()88 const 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 
GetCoffFileHeader()93 const 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 
GetOptionalHeaderData(size_t * optional_header_size)99 const uint8_t* PeImageReader::GetOptionalHeaderData(
100     size_t* optional_header_size) {
101   *optional_header_size = GetOptionalHeaderSize();
102   return GetOptionalHeaderStart();
103 }
104 
GetNumberOfSections()105 size_t PeImageReader::GetNumberOfSections() {
106   return GetCoffFileHeader()->NumberOfSections;
107 }
108 
GetSectionHeaderAt(size_t index)109 const 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 
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 NULL;
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 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 
Clear()159 void PeImageReader::Clear() {
160   image_data_ = NULL;
161   image_size_ = 0;
162   validation_state_ = 0;
163   optional_header_.reset();
164 }
165 
ValidateDosHeader()166 bool 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 
ValidatePeSignature()178 bool 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 
ValidateCoffFileHeader()189 bool 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 
ValidateOptionalHeader()202 bool 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 
ValidateSectionHeaders()245 bool 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 
GetOptionalHeaderStart()261 const 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 
GetOptionalHeaderSize()268 size_t PeImageReader::GetOptionalHeaderSize() {
269   return GetCoffFileHeader()->SizeOfOptionalHeader;
270 }
271 
GetDataDirectoryEntryAt(size_t index)272 const 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 
FindSectionFromRva(uint32_t relative_address)280 const 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 
GetImageData(size_t index,size_t * data_length)306 const 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