chrome/browser/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_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_