1 // Copyright 2014 The Chromium Authors 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 BASE_WIN_PE_IMAGE_READER_H_ 6 #define BASE_WIN_PE_IMAGE_READER_H_ 7 8 #include <windows.h> 9 10 #include <stddef.h> 11 #include <stdint.h> 12 13 #include <memory> 14 15 #include "base/base_export.h" 16 #include "base/containers/span.h" 17 #include "base/memory/raw_span.h" 18 #include "base/numerics/safe_math.h" 19 20 namespace base { 21 namespace win { 22 23 // Parses headers and various data from a PE image. This parser is safe for use 24 // on untrusted data and works on PE files with different bitness from the 25 // current process. The PeImageReader is initialized after construction by 26 // passing the address and size of a PE file that has been read into memory - 27 // not loaded by the OS as an image. Parsing of a PE file that has been loaded 28 // as an image can be done with PEImage. 29 class BASE_EXPORT PeImageReader { 30 public: 31 enum WordSize { 32 WORD_SIZE_32, 33 WORD_SIZE_64, 34 }; 35 36 // A callback invoked by EnumCertificates once for each attribute certificate 37 // entry in the image's attribute certificate table. |revision| and 38 // |certificate_type| identify the contents of |certificate_data| (which is of 39 // |certificate_data_size| bytes). |context| is the value provided by the 40 // caller to EnumCertificates(). Implementations must return true to continue 41 // the enumeration, or false to abort. 42 using EnumCertificatesCallback = bool (*)(uint16_t revision, 43 uint16_t certificate_type, 44 const uint8_t* certificate_data, 45 size_t certificate_data_size, 46 void* context); 47 48 PeImageReader(); 49 50 PeImageReader(const PeImageReader&) = delete; 51 PeImageReader& operator=(const PeImageReader&) = delete; 52 53 ~PeImageReader(); 54 55 // Returns false if the given data does not appear to be a valid PE image. 56 bool Initialize(span<const uint8_t> image_data); 57 58 // Returns the machine word size for the image. 59 WordSize GetWordSize(); 60 61 const IMAGE_DOS_HEADER* GetDosHeader(); 62 const IMAGE_FILE_HEADER* GetCoffFileHeader(); 63 64 // Returns the optional header data. 65 span<const uint8_t> GetOptionalHeaderData(); 66 size_t GetNumberOfSections(); 67 const IMAGE_SECTION_HEADER* GetSectionHeaderAt(size_t index); 68 69 // Returns the image's export data (.edata) section, or an empty span if the 70 // section is not present. 71 span<const uint8_t> GetExportSection(); 72 73 size_t GetNumberOfDebugEntries(); 74 // Returns a pointer to the |index|'th debug directory entry, or nullptr if 75 // |index| is out of bounds. |raw_data| is an out-param which will be filled 76 // with the corresponding raw data. 77 const IMAGE_DEBUG_DIRECTORY* GetDebugEntry(size_t index, 78 span<const uint8_t>& raw_data); 79 80 // Invokes |callback| once per attribute certificate entry. |context| is a 81 // caller-specific value that is passed to |callback|. Returns true if all 82 // certificate entries are visited (even if there are no such entries) and 83 // |callback| returns true for each. Conversely, returns |false| if |callback| 84 // returns false or if the image is malformed in any way. 85 bool EnumCertificates(EnumCertificatesCallback callback, void* context); 86 87 // Returns the size of the image file. 88 DWORD GetSizeOfImage(); 89 90 private: 91 // Bits indicating what portions of the image have been validated. 92 enum ValidationStages { 93 VALID_DOS_HEADER = 1 << 0, 94 VALID_PE_SIGNATURE = 1 << 1, 95 VALID_COFF_FILE_HEADER = 1 << 2, 96 VALID_OPTIONAL_HEADER = 1 << 3, 97 VALID_SECTION_HEADERS = 1 << 4, 98 }; 99 100 // An interface to an image's optional header. 101 class OptionalHeader { 102 public: 103 virtual ~OptionalHeader() = default; 104 105 virtual WordSize GetWordSize() = 0; 106 107 // Returns the offset of the DataDirectory member relative to the start of 108 // the optional header. 109 virtual size_t GetDataDirectoryOffset() = 0; 110 111 // Returns the number of entries in the data directory. 112 virtual DWORD GetDataDirectorySize() = 0; 113 114 // Returns a pointer to the first data directory entry. 115 virtual const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntries() = 0; 116 117 // Returns the size of the image file. 118 virtual DWORD GetSizeOfImage() = 0; 119 }; 120 121 template <class OPTIONAL_HEADER_TYPE> 122 class OptionalHeaderImpl; 123 124 void Clear(); 125 bool ValidateDosHeader(); 126 bool ValidatePeSignature(); 127 bool ValidateCoffFileHeader(); 128 bool ValidateOptionalHeader(); 129 bool ValidateSectionHeaders(); 130 131 // Return a pointer to the first byte of the image's optional header. 132 const uint8_t* GetOptionalHeaderStart(); 133 size_t GetOptionalHeaderSize(); 134 135 // Returns the desired directory entry, or nullptr if |index| is out of 136 // bounds. 137 const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntryAt(size_t index); 138 139 // Returns the header for the section that contains the given address, or 140 // nullptr if the address is out of bounds or the image does not contain the 141 // section. 142 const IMAGE_SECTION_HEADER* FindSectionFromRva(uint32_t relative_address); 143 144 // Returns the bytes referenced by the |index|'th data directory entry. 145 span<const uint8_t> GetImageData(size_t index); 146 147 // Populates |structure| with a pointer to a desired structure of type T at 148 // the given offset if the image is sufficiently large to contain it. Returns 149 // false if the structure does not fully fit within the image at the given 150 // offset. 151 template <typename T> GetStructureAt(size_t offset,const T ** structure)152 bool GetStructureAt(size_t offset, const T** structure) { 153 return GetStructureAt(offset, sizeof(**structure), structure); 154 } 155 156 // Populates |structure| with a pointer to a desired structure of type T at 157 // the given offset if the image is sufficiently large to contain 158 // |structure_size| bytes. Returns false if the structure does not fully fit 159 // within the image at the given offset. 160 template <typename T> GetStructureAt(size_t offset,size_t structure_size,const T ** structure)161 bool GetStructureAt(size_t offset, 162 size_t structure_size, 163 const T** structure) { 164 size_t remaining_bytes = 0; 165 if (!CheckSub(image_data_.size(), offset).AssignIfValid(&remaining_bytes)) { 166 return false; 167 } 168 if (structure_size > remaining_bytes) { 169 return false; 170 } 171 *structure = reinterpret_cast<const T*>(image_data_.subspan(offset).data()); 172 return true; 173 } 174 175 raw_span<const uint8_t> image_data_; 176 uint32_t validation_state_ = 0; 177 std::unique_ptr<OptionalHeader> optional_header_; 178 }; 179 180 } // namespace win 181 } // namespace base 182 183 #endif // BASE_WIN_PE_IMAGE_READER_H_ 184