[2573] | 1 | #ifndef PE_HPP |
---|
| 2 | #define PE_HPP |
---|
| 3 | |
---|
| 4 | #include "common.hpp" |
---|
| 5 | |
---|
| 6 | #if IS_WINDOWS |
---|
| 7 | #include <array> |
---|
| 8 | #include <cstddef> |
---|
| 9 | #include <cstdio> |
---|
| 10 | #include <cstring> |
---|
| 11 | #include <string> |
---|
| 12 | |
---|
| 13 | #include <windows.h> |
---|
| 14 | |
---|
| 15 | template<typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0> |
---|
| 16 | T pe_byteswap_if_needed(T value) { |
---|
| 17 | // PE header values are little endian |
---|
| 18 | if(!is_little_endian()) { |
---|
| 19 | return byteswap(value); |
---|
| 20 | } else { |
---|
| 21 | return value; |
---|
| 22 | } |
---|
| 23 | } |
---|
| 24 | |
---|
| 25 | static uintptr_t pe_get_module_image_base(const std::string& obj_path) { |
---|
| 26 | FILE* file; |
---|
| 27 | errno_t ret = fopen_s(&file, obj_path.c_str(), "rb"); |
---|
| 28 | if(ret != 0 || file == nullptr) { |
---|
| 29 | throw file_error(); |
---|
| 30 | return 0; |
---|
| 31 | } |
---|
| 32 | auto magic = load_bytes<std::array<char, 2>>(file, 0); |
---|
| 33 | internal_verify(memcmp(magic.data(), "MZ", 2) == 0); |
---|
| 34 | DWORD e_lfanew = pe_byteswap_if_needed(load_bytes<DWORD>(file, 0x3c)); // dos header + 0x3c |
---|
| 35 | long nt_header_offset = e_lfanew; |
---|
| 36 | auto signature = load_bytes<std::array<char, 4>>(file, nt_header_offset); // nt header + 0 |
---|
| 37 | internal_verify(memcmp(signature.data(), "PE\0\0", 4) == 0); |
---|
| 38 | WORD size_of_optional_header = pe_byteswap_if_needed( |
---|
| 39 | load_bytes<WORD>(file, nt_header_offset + 4 + 0x10) // file header + 0x10 |
---|
| 40 | ); |
---|
| 41 | internal_verify(size_of_optional_header != 0); |
---|
| 42 | WORD optional_header_magic = pe_byteswap_if_needed( |
---|
| 43 | load_bytes<WORD>(file, nt_header_offset + 0x18) // optional header + 0x0 |
---|
| 44 | ); |
---|
| 45 | internal_verify(optional_header_magic == IMAGE_NT_OPTIONAL_HDR_MAGIC); |
---|
| 46 | uintptr_t image_base; |
---|
| 47 | if(optional_header_magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) { |
---|
| 48 | // 32 bit |
---|
| 49 | image_base = pe_byteswap_if_needed( |
---|
| 50 | load_bytes<DWORD>(file, nt_header_offset + 0x18 + 0x1c) // optional header + 0x1c |
---|
| 51 | ); |
---|
| 52 | } else { |
---|
| 53 | // 64 bit |
---|
| 54 | // I get an "error: 'QWORD' was not declared in this scope" for some reason when using QWORD |
---|
| 55 | image_base = pe_byteswap_if_needed( |
---|
| 56 | load_bytes<unsigned __int64>(file, nt_header_offset + 0x18 + 0x18) // optional header + 0x18 |
---|
| 57 | ); |
---|
| 58 | } |
---|
| 59 | fclose(file); |
---|
| 60 | return image_base; |
---|
| 61 | } |
---|
| 62 | #endif |
---|
| 63 | |
---|
| 64 | #endif |
---|