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 334