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#ifndef CHROME_BROWSER_SAFE_BROWSING_PE_IMAGE_READER_WIN_H_
6#define CHROME_BROWSER_SAFE_BROWSING_PE_IMAGE_READER_WIN_H_
7
8#include <windows.h>
9
10#include "base/basictypes.h"
11#include "base/memory/scoped_ptr.h"
12
13namespace safe_browsing {
14
15// Parses headers and various data from a PE image. This parser is safe for use
16// on untrusted data.
17class PeImageReader {
18 public:
19  enum WordSize {
20    WORD_SIZE_32,
21    WORD_SIZE_64,
22  };
23
24  PeImageReader();
25  ~PeImageReader();
26
27  // Returns false if the given data does not appear to be a valid PE image.
28  bool Initialize(const uint8_t* image_data, size_t image_size);
29
30  // Returns the machine word size for the image.
31  WordSize GetWordSize();
32
33  const IMAGE_DOS_HEADER* GetDosHeader();
34  const IMAGE_FILE_HEADER* GetCoffFileHeader();
35
36  // Returns a pointer to the optional header and its size.
37  const uint8_t* GetOptionalHeaderData(size_t* optional_data_size);
38  size_t GetNumberOfSections();
39  const IMAGE_SECTION_HEADER* GetSectionHeaderAt(size_t index);
40
41  // Returns a pointer to the image's export data (.edata) section and its size,
42  // or NULL if the section is not present.
43  const uint8_t* GetExportSection(size_t* section_size);
44
45  size_t GetNumberOfDebugEntries();
46  const IMAGE_DEBUG_DIRECTORY* GetDebugEntry(size_t index,
47                                             const uint8_t** raw_data,
48                                             size_t* raw_data_size);
49
50 private:
51  // Bits indicating what portions of the image have been validated.
52  enum ValidationStages {
53    VALID_DOS_HEADER = 1 << 0,
54    VALID_PE_SIGNATURE = 1 << 1,
55    VALID_COFF_FILE_HEADER = 1 << 2,
56    VALID_OPTIONAL_HEADER = 1 << 3,
57    VALID_SECTION_HEADERS = 1 << 4,
58  };
59
60  // An interface to an image's optional header.
61  class OptionalHeader {
62   public:
63    virtual ~OptionalHeader() {}
64
65    virtual WordSize GetWordSize() = 0;
66
67    // Returns the offset of the DataDirectory member relative to the start of
68    // the optional header.
69    virtual size_t GetDataDirectoryOffset() = 0;
70
71    // Returns the number of entries in the data directory.
72    virtual DWORD GetDataDirectorySize() = 0;
73
74    // Returns a pointer to the first data directory entry.
75    virtual const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntries() = 0;
76  };
77
78  template<class OPTIONAL_HEADER_TYPE>
79  class OptionalHeaderImpl;
80
81  void Clear();
82  bool ValidateDosHeader();
83  bool ValidatePeSignature();
84  bool ValidateCoffFileHeader();
85  bool ValidateOptionalHeader();
86  bool ValidateSectionHeaders();
87
88  // Return a pointer to the first byte of the image's optional header.
89  const uint8_t* GetOptionalHeaderStart();
90  size_t GetOptionalHeaderSize();
91
92  // Returns the desired directory entry, or NULL if |index| is out of bounds.
93  const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntryAt(size_t index);
94
95  // Returns the header for the section that contains the given address, or NULL
96  // if the address is out of bounds or the image does not contain the section.
97  const IMAGE_SECTION_HEADER* FindSectionFromRva(uint32_t relative_address);
98
99  // Returns a pointer to the |data_length| bytes referenced by the |index|'th
100  // data directory entry.
101  const uint8_t* GetImageData(size_t index, size_t* data_length);
102
103  // Populates |structure| with a pointer to a desired structure of type T at
104  // the given offset if the image is sufficiently large to contain it. Returns
105  // false if the structure does not fully fit within the image at the given
106  // offset.
107  template<typename T> bool GetStructureAt(size_t offset, const T** structure) {
108    return GetStructureAt(offset, sizeof(**structure), structure);
109  }
110
111  // Populates |structure| with a pointer to a desired structure of type T at
112  // the given offset if the image is sufficiently large to contain
113  // |structure_size| bytes. Returns false if the structure does not fully fit
114  // within the image at the given offset.
115  template<typename T> bool GetStructureAt(size_t offset,
116                                           size_t structure_size,
117                                           const T** structure) {
118    if (offset > image_size_)
119      return false;
120    if (structure_size > image_size_ - offset)
121      return false;
122    *structure = reinterpret_cast<const T*>(image_data_ + offset);
123    return true;
124  }
125
126  const uint8_t* image_data_;
127  size_t image_size_;
128  uint32_t validation_state_;
129  scoped_ptr<OptionalHeader> optional_header_;
130  DISALLOW_COPY_AND_ASSIGN(PeImageReader);
131};
132
133}  // namespace safe_browsing
134
135#endif  // CHROME_BROWSER_SAFE_BROWSING_PE_IMAGE_READER_WIN_H_
136