src/chrome/browser/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/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                         &section_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