Home | History | Annotate | Download | only in safe_browsing 
      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                         &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 
    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