| OLD | NEW | 
|---|
|  | (Empty) | 
| 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 |  | 
| 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 |  | 
| 59 PeImageReader::PeImageReader() |  | 
| 60     : image_data_(), |  | 
| 61       image_size_(), |  | 
| 62       validation_state_() {} |  | 
| 63 |  | 
| 64 PeImageReader::~PeImageReader() { |  | 
| 65   Clear(); |  | 
| 66 } |  | 
| 67 |  | 
| 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 |  | 
| 84 PeImageReader::WordSize PeImageReader::GetWordSize() { |  | 
| 85   return optional_header_->GetWordSize(); |  | 
| 86 } |  | 
| 87 |  | 
| 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 |  | 
| 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 |  | 
| 99 const uint8_t* PeImageReader::GetOptionalHeaderData( |  | 
| 100     size_t* optional_header_size) { |  | 
| 101   *optional_header_size = GetOptionalHeaderSize(); |  | 
| 102   return GetOptionalHeaderStart(); |  | 
| 103 } |  | 
| 104 |  | 
| 105 size_t PeImageReader::GetNumberOfSections() { |  | 
| 106   return GetCoffFileHeader()->NumberOfSections; |  | 
| 107 } |  | 
| 108 |  | 
| 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 |  | 
| 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 |  | 
| 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 |  | 
| 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 |  | 
| 159 void PeImageReader::Clear() { |  | 
| 160   image_data_ = NULL; |  | 
| 161   image_size_ = 0; |  | 
| 162   validation_state_ = 0; |  | 
| 163   optional_header_.reset(); |  | 
| 164 } |  | 
| 165 |  | 
| 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 |  | 
| 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 |  | 
| 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 |  | 
| 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 |  | 
| 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 |  | 
| 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 |  | 
| 268 size_t PeImageReader::GetOptionalHeaderSize() { |  | 
| 269   return GetCoffFileHeader()->SizeOfOptionalHeader; |  | 
| 270 } |  | 
| 271 |  | 
| 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 |  | 
| 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                         §ion_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 |  | 
| 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 |  | 
| OLD | NEW | 
|---|