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 |
---|