pe_image_reader.cc 12 KB

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