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 #ifndef CHROME_BROWSER_SAFE_BROWSING_PE_IMAGE_READER_WIN_H_
      6 #define CHROME_BROWSER_SAFE_BROWSING_PE_IMAGE_READER_WIN_H_
      7 
      8 #include <windows.h>
      9 
     10 #include "base/basictypes.h"
     11 #include "base/memory/scoped_ptr.h"
     12 
     13 namespace safe_browsing {
     14 
     15 // Parses headers and various data from a PE image. This parser is safe for use
     16 // on untrusted data.
     17 class PeImageReader {
     18  public:
     19   enum WordSize {
     20     WORD_SIZE_32,
     21     WORD_SIZE_64,
     22   };
     23 
     24   PeImageReader();
     25   ~PeImageReader();
     26 
     27   // Returns false if the given data does not appear to be a valid PE image.
     28   bool Initialize(const uint8_t* image_data, size_t image_size);
     29 
     30   // Returns the machine word size for the image.
     31   WordSize GetWordSize();
     32 
     33   const IMAGE_DOS_HEADER* GetDosHeader();
     34   const IMAGE_FILE_HEADER* GetCoffFileHeader();
     35 
     36   // Returns a pointer to the optional header and its size.
     37   const uint8_t* GetOptionalHeaderData(size_t* optional_data_size);
     38   size_t GetNumberOfSections();
     39   const IMAGE_SECTION_HEADER* GetSectionHeaderAt(size_t index);
     40 
     41   // Returns a pointer to the image's export data (.edata) section and its size,
     42   // or NULL if the section is not present.
     43   const uint8_t* GetExportSection(size_t* section_size);
     44 
     45   size_t GetNumberOfDebugEntries();
     46   const IMAGE_DEBUG_DIRECTORY* GetDebugEntry(size_t index,
     47                                              const uint8_t** raw_data,
     48                                              size_t* raw_data_size);
     49 
     50  private:
     51   // Bits indicating what portions of the image have been validated.
     52   enum ValidationStages {
     53     VALID_DOS_HEADER = 1 << 0,
     54     VALID_PE_SIGNATURE = 1 << 1,
     55     VALID_COFF_FILE_HEADER = 1 << 2,
     56     VALID_OPTIONAL_HEADER = 1 << 3,
     57     VALID_SECTION_HEADERS = 1 << 4,
     58   };
     59 
     60   // An interface to an image's optional header.
     61   class OptionalHeader {
     62    public:
     63     virtual ~OptionalHeader() {}
     64 
     65     virtual WordSize GetWordSize() = 0;
     66 
     67     // Returns the offset of the DataDirectory member relative to the start of
     68     // the optional header.
     69     virtual size_t GetDataDirectoryOffset() = 0;
     70 
     71     // Returns the number of entries in the data directory.
     72     virtual DWORD GetDataDirectorySize() = 0;
     73 
     74     // Returns a pointer to the first data directory entry.
     75     virtual const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntries() = 0;
     76   };
     77 
     78   template<class OPTIONAL_HEADER_TYPE>
     79   class OptionalHeaderImpl;
     80 
     81   void Clear();
     82   bool ValidateDosHeader();
     83   bool ValidatePeSignature();
     84   bool ValidateCoffFileHeader();
     85   bool ValidateOptionalHeader();
     86   bool ValidateSectionHeaders();
     87 
     88   // Return a pointer to the first byte of the image's optional header.
     89   const uint8_t* GetOptionalHeaderStart();
     90   size_t GetOptionalHeaderSize();
     91 
     92   // Returns the desired directory entry, or NULL if |index| is out of bounds.
     93   const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntryAt(size_t index);
     94 
     95   // Returns the header for the section that contains the given address, or NULL
     96   // if the address is out of bounds or the image does not contain the section.
     97   const IMAGE_SECTION_HEADER* FindSectionFromRva(uint32_t relative_address);
     98 
     99   // Returns a pointer to the |data_length| bytes referenced by the |index|'th
    100   // data directory entry.
    101   const uint8_t* GetImageData(size_t index, size_t* data_length);
    102 
    103   // Populates |structure| with a pointer to a desired structure of type T at
    104   // the given offset if the image is sufficiently large to contain it. Returns
    105   // false if the structure does not fully fit within the image at the given
    106   // offset.
    107   template<typename T> bool GetStructureAt(size_t offset, const T** structure) {
    108     return GetStructureAt(offset, sizeof(**structure), structure);
    109   }
    110 
    111   // Populates |structure| with a pointer to a desired structure of type T at
    112   // the given offset if the image is sufficiently large to contain
    113   // |structure_size| bytes. Returns false if the structure does not fully fit
    114   // within the image at the given offset.
    115   template<typename T> bool GetStructureAt(size_t offset,
    116                                            size_t structure_size,
    117                                            const T** structure) {
    118     if (offset > image_size_)
    119       return false;
    120     if (structure_size > image_size_ - offset)
    121       return false;
    122     *structure = reinterpret_cast<const T*>(image_data_ + offset);
    123     return true;
    124   }
    125 
    126   const uint8_t* image_data_;
    127   size_t image_size_;
    128   uint32_t validation_state_;
    129   scoped_ptr<OptionalHeader> optional_header_;
    130   DISALLOW_COPY_AND_ASSIGN(PeImageReader);
    131 };
    132 
    133 }  // namespace safe_browsing
    134 
    135 #endif  // CHROME_BROWSER_SAFE_BROWSING_PE_IMAGE_READER_WIN_H_
    136