| OLD | NEW |
| (Empty) |
| 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 #include "chrome/browser/safe_browsing/pe_image_reader_win.h" | |
| 6 | |
| 7 #include "base/logging.h" | |
| 8 | |
| 9 namespace safe_browsing { | |
| 10 | |
| 11 // A class template of traits pertaining to IMAGE_OPTIONAL_HEADER{32,64}. | |
| 12 template<class HEADER_TYPE> | |
| 13 struct OptionalHeaderTraits { | |
| 14 }; | |
| 15 | |
| 16 template<> | |
| 17 struct OptionalHeaderTraits<IMAGE_OPTIONAL_HEADER32> { | |
| 18 static const PeImageReader::WordSize word_size = PeImageReader::WORD_SIZE_32; | |
| 19 }; | |
| 20 | |
| 21 template<> | |
| 22 struct OptionalHeaderTraits<IMAGE_OPTIONAL_HEADER64> { | |
| 23 static const PeImageReader::WordSize word_size = PeImageReader::WORD_SIZE_64; | |
| 24 }; | |
| 25 | |
| 26 // A template for type-specific optional header implementations. This, in | |
| 27 // conjunction with the OptionalHeader interface, effectively erases the | |
| 28 // underlying structure type from the point of view of the PeImageReader. | |
| 29 template<class OPTIONAL_HEADER_TYPE> | |
| 30 class PeImageReader::OptionalHeaderImpl : public PeImageReader::OptionalHeader { | |
| 31 public: | |
| 32 typedef OptionalHeaderTraits<OPTIONAL_HEADER_TYPE> TraitsType; | |
| 33 | |
| 34 explicit OptionalHeaderImpl(const uint8_t* optional_header_start) | |
| 35 : optional_header_(reinterpret_cast<const OPTIONAL_HEADER_TYPE*>( | |
| 36 optional_header_start)) {} | |
| 37 | |
| 38 virtual WordSize GetWordSize() override { | |
| 39 return TraitsType::word_size; | |
| 40 } | |
| 41 | |
| 42 virtual size_t GetDataDirectoryOffset() override { | |
| 43 return offsetof(OPTIONAL_HEADER_TYPE, DataDirectory); | |
| 44 } | |
| 45 | |
| 46 virtual DWORD GetDataDirectorySize() override { | |
| 47 return optional_header_->NumberOfRvaAndSizes; | |
| 48 } | |
| 49 | |
| 50 virtual const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntries() override { | |
| 51 return &optional_header_->DataDirectory[0]; | |
| 52 } | |
| 53 | |
| 54 private: | |
| 55 const OPTIONAL_HEADER_TYPE* optional_header_; | |
| 56 DISALLOW_COPY_AND_ASSIGN(OptionalHeaderImpl); | |
| 57 }; | |
| 58 | |
| 59 PeImageReader::PeImageReader() | |
| 60 : image_data_(), | |
| 61 image_size_(), | |
| 62 validation_state_() {} | |
| 63 | |
| 64 PeImageReader::~PeImageReader() { | |
| 65 Clear(); | |
| 66 } | |
| 67 | |
| 68 bool PeImageReader::Initialize(const uint8_t* image_data, size_t image_size) { | |
| 69 image_data_ = image_data; | |
| 70 image_size_ = image_size; | |
| 71 | |
| 72 if (!ValidateDosHeader() || | |
| 73 !ValidatePeSignature() || | |
| 74 !ValidateCoffFileHeader() || | |
| 75 !ValidateOptionalHeader() || | |
| 76 !ValidateSectionHeaders()) { | |
| 77 Clear(); | |
| 78 return false; | |
| 79 } | |
| 80 | |
| 81 return true; | |
| 82 } | |
| 83 | |
| 84 PeImageReader::WordSize PeImageReader::GetWordSize() { | |
| 85 return optional_header_->GetWordSize(); | |
| 86 } | |
| 87 | |
| 88 const IMAGE_DOS_HEADER* PeImageReader::GetDosHeader() { | |
| 89 DCHECK_NE((validation_state_ & VALID_DOS_HEADER), 0U); | |
| 90 return reinterpret_cast<const IMAGE_DOS_HEADER*>(image_data_); | |
| 91 } | |
| 92 | |
| 93 const IMAGE_FILE_HEADER* PeImageReader::GetCoffFileHeader() { | |
| 94 DCHECK_NE((validation_state_ & VALID_COFF_FILE_HEADER), 0U); | |
| 95 return reinterpret_cast<const IMAGE_FILE_HEADER*>( | |
| 96 image_data_ + GetDosHeader()->e_lfanew + sizeof(DWORD)); | |
| 97 } | |
| 98 | |
| 99 const uint8_t* PeImageReader::GetOptionalHeaderData( | |
| 100 size_t* optional_header_size) { | |
| 101 *optional_header_size = GetOptionalHeaderSize(); | |
| 102 return GetOptionalHeaderStart(); | |
| 103 } | |
| 104 | |
| 105 size_t PeImageReader::GetNumberOfSections() { | |
| 106 return GetCoffFileHeader()->NumberOfSections; | |
| 107 } | |
| 108 | |
| 109 const IMAGE_SECTION_HEADER* PeImageReader::GetSectionHeaderAt(size_t index) { | |
| 110 DCHECK_NE((validation_state_ & VALID_SECTION_HEADERS), 0U); | |
| 111 DCHECK_LT(index, GetNumberOfSections()); | |
| 112 return reinterpret_cast<const IMAGE_SECTION_HEADER*>( | |
| 113 GetOptionalHeaderStart() + | |
| 114 GetOptionalHeaderSize() + | |
| 115 (sizeof(IMAGE_SECTION_HEADER) * index)); | |
| 116 } | |
| 117 | |
| 118 const uint8_t* PeImageReader::GetExportSection(size_t* section_size) { | |
| 119 size_t data_size = 0; | |
| 120 const uint8_t* data = GetImageData(IMAGE_DIRECTORY_ENTRY_EXPORT, &data_size); | |
| 121 | |
| 122 // The export section data must be big enough for the export directory. | |
| 123 if (!data || data_size < sizeof(IMAGE_EXPORT_DIRECTORY)) | |
| 124 return NULL; | |
| 125 | |
| 126 *section_size = data_size; | |
| 127 return data; | |
| 128 } | |
| 129 | |
| 130 size_t PeImageReader::GetNumberOfDebugEntries() { | |
| 131 size_t data_size = 0; | |
| 132 const uint8_t* data = GetImageData(IMAGE_DIRECTORY_ENTRY_DEBUG, &data_size); | |
| 133 return data ? (data_size / sizeof(IMAGE_DEBUG_DIRECTORY)) : 0; | |
| 134 } | |
| 135 | |
| 136 const IMAGE_DEBUG_DIRECTORY* PeImageReader::GetDebugEntry( | |
| 137 size_t index, | |
| 138 const uint8_t** raw_data, | |
| 139 size_t* raw_data_size) { | |
| 140 DCHECK_LT(index, GetNumberOfDebugEntries()); | |
| 141 | |
| 142 // Get the debug directory. | |
| 143 size_t debug_directory_size = 0; | |
| 144 const IMAGE_DEBUG_DIRECTORY* entries = | |
| 145 reinterpret_cast<const IMAGE_DEBUG_DIRECTORY*>( | |
| 146 GetImageData(IMAGE_DIRECTORY_ENTRY_DEBUG, &debug_directory_size)); | |
| 147 if (!entries) | |
| 148 return NULL; | |
| 149 | |
| 150 const IMAGE_DEBUG_DIRECTORY& entry = entries[index]; | |
| 151 const uint8_t* debug_data = NULL; | |
| 152 if (GetStructureAt(entry.PointerToRawData, entry.SizeOfData, &debug_data)) { | |
| 153 *raw_data = debug_data; | |
| 154 *raw_data_size = entry.SizeOfData; | |
| 155 } | |
| 156 return &entry; | |
| 157 } | |
| 158 | |
| 159 void PeImageReader::Clear() { | |
| 160 image_data_ = NULL; | |
| 161 image_size_ = 0; | |
| 162 validation_state_ = 0; | |
| 163 optional_header_.reset(); | |
| 164 } | |
| 165 | |
| 166 bool PeImageReader::ValidateDosHeader() { | |
| 167 const IMAGE_DOS_HEADER* dos_header = NULL; | |
| 168 if (!GetStructureAt(0, &dos_header) || | |
| 169 dos_header->e_magic != IMAGE_DOS_SIGNATURE || | |
| 170 dos_header->e_lfanew < 0) { | |
| 171 return false; | |
| 172 } | |
| 173 | |
| 174 validation_state_ |= VALID_DOS_HEADER; | |
| 175 return true; | |
| 176 } | |
| 177 | |
| 178 bool PeImageReader::ValidatePeSignature() { | |
| 179 const DWORD* signature = NULL; | |
| 180 if (!GetStructureAt(GetDosHeader()->e_lfanew, &signature) || | |
| 181 *signature != IMAGE_NT_SIGNATURE) { | |
| 182 return false; | |
| 183 } | |
| 184 | |
| 185 validation_state_ |= VALID_PE_SIGNATURE; | |
| 186 return true; | |
| 187 } | |
| 188 | |
| 189 bool PeImageReader::ValidateCoffFileHeader() { | |
| 190 DCHECK_NE((validation_state_ & VALID_PE_SIGNATURE), 0U); | |
| 191 const IMAGE_FILE_HEADER* file_header = NULL; | |
| 192 if (!GetStructureAt(GetDosHeader()->e_lfanew + | |
| 193 offsetof(IMAGE_NT_HEADERS32, FileHeader), | |
| 194 &file_header)) { | |
| 195 return false; | |
| 196 } | |
| 197 | |
| 198 validation_state_ |= VALID_COFF_FILE_HEADER; | |
| 199 return true; | |
| 200 } | |
| 201 | |
| 202 bool PeImageReader::ValidateOptionalHeader() { | |
| 203 const IMAGE_FILE_HEADER* file_header = GetCoffFileHeader(); | |
| 204 const size_t optional_header_offset = | |
| 205 GetDosHeader()->e_lfanew + offsetof(IMAGE_NT_HEADERS32, OptionalHeader); | |
| 206 const size_t optional_header_size = file_header->SizeOfOptionalHeader; | |
| 207 const WORD* optional_header_magic = NULL; | |
| 208 | |
| 209 if (optional_header_size < sizeof(*optional_header_magic) || | |
| 210 !GetStructureAt(optional_header_offset, &optional_header_magic)) { | |
| 211 return false; | |
| 212 } | |
| 213 | |
| 214 scoped_ptr<OptionalHeader> optional_header; | |
| 215 if (*optional_header_magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) { | |
| 216 optional_header.reset(new OptionalHeaderImpl<IMAGE_OPTIONAL_HEADER32>( | |
| 217 image_data_ + optional_header_offset)); | |
| 218 } else if (*optional_header_magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC) { | |
| 219 optional_header.reset(new OptionalHeaderImpl<IMAGE_OPTIONAL_HEADER64>( | |
| 220 image_data_ + optional_header_offset)); | |
| 221 } else { | |
| 222 return false; | |
| 223 } | |
| 224 | |
| 225 // Does all of the claimed optional header fit in the image? | |
| 226 if (optional_header_size > image_size_ - optional_header_offset) | |
| 227 return false; | |
| 228 | |
| 229 // Is the claimed optional header big enough for everything but the dir? | |
| 230 if (optional_header->GetDataDirectoryOffset() > optional_header_size) | |
| 231 return false; | |
| 232 | |
| 233 // Is there enough room for all of the claimed directory entries? | |
| 234 if (optional_header->GetDataDirectorySize() > | |
| 235 ((optional_header_size - optional_header->GetDataDirectoryOffset()) / | |
| 236 sizeof(IMAGE_DATA_DIRECTORY))) { | |
| 237 return false; | |
| 238 } | |
| 239 | |
| 240 optional_header_.swap(optional_header); | |
| 241 validation_state_ |= VALID_OPTIONAL_HEADER; | |
| 242 return true; | |
| 243 } | |
| 244 | |
| 245 bool PeImageReader::ValidateSectionHeaders() { | |
| 246 const uint8_t* first_section_header = | |
| 247 GetOptionalHeaderStart() + GetOptionalHeaderSize(); | |
| 248 const size_t number_of_sections = GetNumberOfSections(); | |
| 249 | |
| 250 // Do all section headers fit in the image? | |
| 251 if (!GetStructureAt(first_section_header - image_data_, | |
| 252 number_of_sections * sizeof(IMAGE_SECTION_HEADER), | |
| 253 &first_section_header)) { | |
| 254 return false; | |
| 255 } | |
| 256 | |
| 257 validation_state_ |= VALID_SECTION_HEADERS; | |
| 258 return true; | |
| 259 } | |
| 260 | |
| 261 const uint8_t* PeImageReader::GetOptionalHeaderStart() { | |
| 262 DCHECK_NE((validation_state_ & VALID_OPTIONAL_HEADER), 0U); | |
| 263 return (image_data_ + | |
| 264 GetDosHeader()->e_lfanew + | |
| 265 offsetof(IMAGE_NT_HEADERS32, OptionalHeader)); | |
| 266 } | |
| 267 | |
| 268 size_t PeImageReader::GetOptionalHeaderSize() { | |
| 269 return GetCoffFileHeader()->SizeOfOptionalHeader; | |
| 270 } | |
| 271 | |
| 272 const IMAGE_DATA_DIRECTORY* PeImageReader::GetDataDirectoryEntryAt( | |
| 273 size_t index) { | |
| 274 DCHECK_NE((validation_state_ & VALID_OPTIONAL_HEADER), 0U); | |
| 275 if (index >= optional_header_->GetDataDirectorySize()) | |
| 276 return NULL; | |
| 277 return &optional_header_->GetDataDirectoryEntries()[index]; | |
| 278 } | |
| 279 | |
| 280 const IMAGE_SECTION_HEADER* PeImageReader::FindSectionFromRva( | |
| 281 uint32_t relative_address) { | |
| 282 const size_t number_of_sections = GetNumberOfSections(); | |
| 283 for (size_t i = 0; i < number_of_sections; ++i) { | |
| 284 const IMAGE_SECTION_HEADER* section_header = GetSectionHeaderAt(i); | |
| 285 // Is the raw data present in the image? If no, optimistically keep looking. | |
| 286 const uint8_t* section_data = NULL; | |
| 287 if (!GetStructureAt(section_header->PointerToRawData, | |
| 288 section_header->SizeOfRawData, | |
| 289 §ion_data)) { | |
| 290 continue; | |
| 291 } | |
| 292 // Does the RVA lie on or after this section's start when mapped? If no, | |
| 293 // bail. | |
| 294 if (section_header->VirtualAddress > relative_address) | |
| 295 break; | |
| 296 // Does the RVA lie within the section when mapped? If no, keep looking. | |
| 297 size_t address_offset = relative_address - section_header->VirtualAddress; | |
| 298 if (address_offset > section_header->Misc.VirtualSize) | |
| 299 continue; | |
| 300 // We have a winner. | |
| 301 return section_header; | |
| 302 } | |
| 303 return NULL; | |
| 304 } | |
| 305 | |
| 306 const uint8_t* PeImageReader::GetImageData(size_t index, size_t* data_length) { | |
| 307 // Get the requested directory entry. | |
| 308 const IMAGE_DATA_DIRECTORY* entry = GetDataDirectoryEntryAt(index); | |
| 309 if (!entry) | |
| 310 return NULL; | |
| 311 | |
| 312 // Find the section containing the data. | |
| 313 const IMAGE_SECTION_HEADER* header = | |
| 314 FindSectionFromRva(entry->VirtualAddress); | |
| 315 if (!header) | |
| 316 return NULL; | |
| 317 | |
| 318 // Does the data fit within the section when mapped? | |
| 319 size_t data_offset = entry->VirtualAddress - header->VirtualAddress; | |
| 320 if (entry->Size > (header->Misc.VirtualSize - data_offset)) | |
| 321 return NULL; | |
| 322 | |
| 323 // Is the data entirely present on disk (if not it's zeroed out when loaded)? | |
| 324 if (data_offset >= header->SizeOfRawData || | |
| 325 header->SizeOfRawData - data_offset < entry->Size) { | |
| 326 return NULL; | |
| 327 } | |
| 328 | |
| 329 *data_length = entry->Size; | |
| 330 return image_data_ + header->PointerToRawData + data_offset; | |
| 331 } | |
| 332 | |
| 333 } // namespace safe_browsing | |
| OLD | NEW |