xref: /aosp_15_r20/external/cronet/base/win/pe_image_reader.h (revision 6777b5387eb2ff775bb5750e3f5d96f37fb7352b)
1 // Copyright 2014 The Chromium Authors
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 BASE_WIN_PE_IMAGE_READER_H_
6 #define BASE_WIN_PE_IMAGE_READER_H_
7 
8 #include <windows.h>
9 
10 #include <stddef.h>
11 #include <stdint.h>
12 
13 #include <memory>
14 
15 #include "base/base_export.h"
16 #include "base/containers/span.h"
17 #include "base/memory/raw_span.h"
18 #include "base/numerics/safe_math.h"
19 
20 namespace base {
21 namespace win {
22 
23 // Parses headers and various data from a PE image. This parser is safe for use
24 // on untrusted data and works on PE files with different bitness from the
25 // current process. The PeImageReader is initialized after construction by
26 // passing the address and size of a PE file that has been read into memory -
27 // not loaded by the OS as an image. Parsing of a PE file that has been loaded
28 // as an image can be done with PEImage.
29 class BASE_EXPORT PeImageReader {
30  public:
31   enum WordSize {
32     WORD_SIZE_32,
33     WORD_SIZE_64,
34   };
35 
36   // A callback invoked by EnumCertificates once for each attribute certificate
37   // entry in the image's attribute certificate table. |revision| and
38   // |certificate_type| identify the contents of |certificate_data| (which is of
39   // |certificate_data_size| bytes). |context| is the value provided by the
40   // caller to EnumCertificates(). Implementations must return true to continue
41   // the enumeration, or false to abort.
42   using EnumCertificatesCallback = bool (*)(uint16_t revision,
43                                             uint16_t certificate_type,
44                                             const uint8_t* certificate_data,
45                                             size_t certificate_data_size,
46                                             void* context);
47 
48   PeImageReader();
49 
50   PeImageReader(const PeImageReader&) = delete;
51   PeImageReader& operator=(const PeImageReader&) = delete;
52 
53   ~PeImageReader();
54 
55   // Returns false if the given data does not appear to be a valid PE image.
56   bool Initialize(span<const uint8_t> image_data);
57 
58   // Returns the machine word size for the image.
59   WordSize GetWordSize();
60 
61   const IMAGE_DOS_HEADER* GetDosHeader();
62   const IMAGE_FILE_HEADER* GetCoffFileHeader();
63 
64   // Returns the optional header data.
65   span<const uint8_t> GetOptionalHeaderData();
66   size_t GetNumberOfSections();
67   const IMAGE_SECTION_HEADER* GetSectionHeaderAt(size_t index);
68 
69   // Returns the image's export data (.edata) section, or an empty span if the
70   // section is not present.
71   span<const uint8_t> GetExportSection();
72 
73   size_t GetNumberOfDebugEntries();
74   // Returns a pointer to the |index|'th debug directory entry, or nullptr if
75   // |index| is out of bounds. |raw_data| is an out-param which will be filled
76   // with the corresponding raw data.
77   const IMAGE_DEBUG_DIRECTORY* GetDebugEntry(size_t index,
78                                              span<const uint8_t>& raw_data);
79 
80   // Invokes |callback| once per attribute certificate entry. |context| is a
81   // caller-specific value that is passed to |callback|. Returns true if all
82   // certificate entries are visited (even if there are no such entries) and
83   // |callback| returns true for each. Conversely, returns |false| if |callback|
84   // returns false or if the image is malformed in any way.
85   bool EnumCertificates(EnumCertificatesCallback callback, void* context);
86 
87   // Returns the size of the image file.
88   DWORD GetSizeOfImage();
89 
90  private:
91   // Bits indicating what portions of the image have been validated.
92   enum ValidationStages {
93     VALID_DOS_HEADER = 1 << 0,
94     VALID_PE_SIGNATURE = 1 << 1,
95     VALID_COFF_FILE_HEADER = 1 << 2,
96     VALID_OPTIONAL_HEADER = 1 << 3,
97     VALID_SECTION_HEADERS = 1 << 4,
98   };
99 
100   // An interface to an image's optional header.
101   class OptionalHeader {
102    public:
103     virtual ~OptionalHeader() = default;
104 
105     virtual WordSize GetWordSize() = 0;
106 
107     // Returns the offset of the DataDirectory member relative to the start of
108     // the optional header.
109     virtual size_t GetDataDirectoryOffset() = 0;
110 
111     // Returns the number of entries in the data directory.
112     virtual DWORD GetDataDirectorySize() = 0;
113 
114     // Returns a pointer to the first data directory entry.
115     virtual const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntries() = 0;
116 
117     // Returns the size of the image file.
118     virtual DWORD GetSizeOfImage() = 0;
119   };
120 
121   template <class OPTIONAL_HEADER_TYPE>
122   class OptionalHeaderImpl;
123 
124   void Clear();
125   bool ValidateDosHeader();
126   bool ValidatePeSignature();
127   bool ValidateCoffFileHeader();
128   bool ValidateOptionalHeader();
129   bool ValidateSectionHeaders();
130 
131   // Return a pointer to the first byte of the image's optional header.
132   const uint8_t* GetOptionalHeaderStart();
133   size_t GetOptionalHeaderSize();
134 
135   // Returns the desired directory entry, or nullptr if |index| is out of
136   // bounds.
137   const IMAGE_DATA_DIRECTORY* GetDataDirectoryEntryAt(size_t index);
138 
139   // Returns the header for the section that contains the given address, or
140   // nullptr if the address is out of bounds or the image does not contain the
141   // section.
142   const IMAGE_SECTION_HEADER* FindSectionFromRva(uint32_t relative_address);
143 
144   // Returns the bytes referenced by the |index|'th data directory entry.
145   span<const uint8_t> GetImageData(size_t index);
146 
147   // Populates |structure| with a pointer to a desired structure of type T at
148   // the given offset if the image is sufficiently large to contain it. Returns
149   // false if the structure does not fully fit within the image at the given
150   // offset.
151   template <typename T>
GetStructureAt(size_t offset,const T ** structure)152   bool GetStructureAt(size_t offset, const T** structure) {
153     return GetStructureAt(offset, sizeof(**structure), structure);
154   }
155 
156   // Populates |structure| with a pointer to a desired structure of type T at
157   // the given offset if the image is sufficiently large to contain
158   // |structure_size| bytes. Returns false if the structure does not fully fit
159   // within the image at the given offset.
160   template <typename T>
GetStructureAt(size_t offset,size_t structure_size,const T ** structure)161   bool GetStructureAt(size_t offset,
162                       size_t structure_size,
163                       const T** structure) {
164     size_t remaining_bytes = 0;
165     if (!CheckSub(image_data_.size(), offset).AssignIfValid(&remaining_bytes)) {
166       return false;
167     }
168     if (structure_size > remaining_bytes) {
169       return false;
170     }
171     *structure = reinterpret_cast<const T*>(image_data_.subspan(offset).data());
172     return true;
173   }
174 
175   raw_span<const uint8_t> image_data_;
176   uint32_t validation_state_ = 0;
177   std::unique_ptr<OptionalHeader> optional_header_;
178 };
179 
180 }  // namespace win
181 }  // namespace base
182 
183 #endif  // BASE_WIN_PE_IMAGE_READER_H_
184