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: ~OptionalHeader()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. GetStructureAt(size_t offset,const T ** structure)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. GetStructureAt(size_t offset,size_t structure_size,const T ** structure)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