chrome/common/safe_browsing/pe_image_reader_win.h

   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_COMMON_SAFE_BROWSING_PE_IMAGE_READER_WIN_H_
   6 #define CHROME_COMMON_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   // A callback invoked by EnumCertificates once for each attribute certificate
  25   // entry in the image's attribute certificate table. |revision| and
  26   // |certificate_type| identify the contents of |certificate_data| (which is of
  27   // |certificate_data_size| bytes). |context| is the value provided by the
  28   // caller to EnumCertificates(). Implementations must return true to continue
  29   // the enumeration, or false to abort.
  30   typedef bool (*EnumCertificatesCallback)(uint16_t revision,
  31                                            uint16_t certificate_type,
  32                                            const uint8_t* certificate_data,
  33                                            size_t certificate_data_size,
  34                                            void* context);
  35
  36   PeImageReader();
  37   ~PeImageReader();
  38
  39   // Returns false if the given data does not appear to be a valid PE image.
  40   bool Initialize(const uint8_t* image_data, size_t image_size);
  41
  42   // Returns the machine word size for the image.
  43   WordSize GetWordSize();
  44
  45   const IMAGE_DOS_HEADER* GetDosHeader();
  46   const IMAGE_FILE_HEADER* GetCoffFileHeader();
  47
  48   // Returns a pointer to the optional header and its size.
  49   const uint8_t* GetOptionalHeaderData(size_t* optional_data_size);
  50   size_t GetNumberOfSections();
  51   const IMAGE_SECTION_HEADER* GetSectionHeaderAt(size_t index);
  52
  53   // Returns a pointer to the image's export data (.edata) section and its size,
  54   // or NULL if the section is not present.
  55   const uint8_t* GetExportSection(size_t* section_size);
  56
  57   size_t GetNumberOfDebugEntries();
  58   const IMAGE_DEBUG_DIRECTORY* GetDebugEntry(size_t index,
  59                                              const uint8_t** raw_data,
  60                                              size_t* raw_data_size);
  61
  62   // Invokes |callback| once per attribute certificate entry. |context| is a
  63   // caller-specific value that is passed to |callback|. Returns true if all
  64   // certificate entries are visited (even if there are no such entries) and
  65   // |callback| returns true for each. Conversely, returns |false| if |callback|
  66   // returns false or if the image is malformed in any way.
  67   bool EnumCertificates(EnumCertificatesCallback callback,
  68                         void* context);
  69
  70  private:
  71   // Bits indicating what portions of the image have been validated.
  72   enum ValidationStages {
  73     VALID_DOS_HEADER = 1 << 0,
  74     VALID_PE_SIGNATURE = 1 << 1,
  75     VALID_COFF_FILE_HEADER = 1 << 2,
  76     VALID_OPTIONAL_HEADER = 1 << 3,
  77     VALID_SECTION_HEADERS = 1 << 4,
  78   };
  79
  80   // An interface to an image's optional header.
  81   class OptionalHeader {
  82    public:
  83     virtual ~OptionalHeader() {}
  84
  85     virtual WordSize GetWordSize() = 0;
  86
  87     // Returns the offset of the DataDirectory member relative to the start of
  88     // the optional header.
  89     virtual size_t GetDataDirectoryOffset() = 0;
  90
  91     // Returns the number of entries in the data directory.
  92     virtual DWORD GetDataDirectorySize() = 0;
  93
  94     // Returns a pointer to the first data directory entry.
  95     virtual const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntries() = 0;
  96   };
  97
  98   template<class OPTIONAL_HEADER_TYPE>
  99   class OptionalHeaderImpl;
 100
 101   void Clear();
 102   bool ValidateDosHeader();
 103   bool ValidatePeSignature();
 104   bool ValidateCoffFileHeader();
 105   bool ValidateOptionalHeader();
 106   bool ValidateSectionHeaders();
 107
 108   // Return a pointer to the first byte of the image's optional header.
 109   const uint8_t* GetOptionalHeaderStart();
 110   size_t GetOptionalHeaderSize();
 111
 112   // Returns the desired directory entry, or NULL if |index| is out of bounds.
 113   const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntryAt(size_t index);
 114
 115   // Returns the header for the section that contains the given address, or NULL
 116   // if the address is out of bounds or the image does not contain the section.
 117   const IMAGE_SECTION_HEADER* FindSectionFromRva(uint32_t relative_address);
 118
 119   // Returns a pointer to the |data_length| bytes referenced by the |index|'th
 120   // data directory entry.
 121   const uint8_t* GetImageData(size_t index, size_t* data_length);
 122
 123   // Populates |structure| with a pointer to a desired structure of type T at
 124   // the given offset if the image is sufficiently large to contain it. Returns
 125   // false if the structure does not fully fit within the image at the given
 126   // offset.
 127   template<typename T> bool GetStructureAt(size_t offset, const T** structure) {
 128     return GetStructureAt(offset, sizeof(**structure), structure);
 129   }
 130
 131   // Populates |structure| with a pointer to a desired structure of type T at
 132   // the given offset if the image is sufficiently large to contain
 133   // |structure_size| bytes. Returns false if the structure does not fully fit
 134   // within the image at the given offset.
 135   template<typename T> bool GetStructureAt(size_t offset,
 136                                            size_t structure_size,
 137                                            const T** structure) {
 138     if (offset > image_size_)
 139       return false;
 140     if (structure_size > image_size_ - offset)
 141       return false;
 142     *structure = reinterpret_cast<const T*>(image_data_ + offset);
 143     return true;
 144   }
 145
 146   const uint8_t* image_data_;
 147   size_t image_size_;
 148   uint32_t validation_state_;
 149   scoped_ptr<OptionalHeader> optional_header_;
 150   DISALLOW_COPY_AND_ASSIGN(PeImageReader);
 151 };
 152
 153 }  // namespace safe_browsing
 154
 155 #endif  // CHROME_COMMON_SAFE_BROWSING_PE_IMAGE_READER_WIN_H_