chrome/common/safe_browsing/pe_image_reader_win.cc

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