From e80a11d730b2fb8fe0e86402c767721713f496c8 Mon Sep 17 00:00:00 2001 From: Jeremy <51220084+jeremy-rifkin@users.noreply.github.com> Date: Sat, 9 Mar 2024 21:25:57 -0600 Subject: [PATCH] Logic for PE, and a couple tweaks --- src/binary/mach-o.hpp | 2 +- src/binary/module_base.hpp | 8 +++-- src/binary/object.hpp | 2 +- src/binary/pe.hpp | 63 +++++++++++++++++++++++++------------- 4 files changed, 49 insertions(+), 26 deletions(-) diff --git a/src/binary/mach-o.hpp b/src/binary/mach-o.hpp index cca495c..ff140e4 100644 --- a/src/binary/mach-o.hpp +++ b/src/binary/mach-o.hpp @@ -583,7 +583,7 @@ namespace detail { return monostate{}; } // If this is reached... something went wrong. The cpu we're on wasn't found. - throw internal_error("Couldn't find appropriate architecture in fat Mach-O"); + return internal_error("Couldn't find appropriate architecture in fat Mach-O"); } template diff --git a/src/binary/module_base.hpp b/src/binary/module_base.hpp index af2e6b8..133682d 100644 --- a/src/binary/module_base.hpp +++ b/src/binary/module_base.hpp @@ -56,8 +56,8 @@ namespace detail { if(it == cache.end()) { // arguably it'd be better to release the lock while computing this, but also arguably it's good to not // have two threads try to do the same computation - // TODO: Cache the error auto obj = mach_o::open_mach_o(object_path); + // TODO: Cache the error if(!obj) { return obj.unwrap_error(); } @@ -81,7 +81,11 @@ namespace detail { // arguably it'd be better to release the lock while computing this, but also arguably it's good to not // have two threads try to do the same computation auto base = pe_get_module_image_base(object_path); - cache.insert(it, {object_path, base}); + // TODO: Cache the error + if(!base) { + return base.unwrap_error(); + } + cache.insert(it, {object_path, base.unwrap_value()}); return base; } else { return it->second; diff --git a/src/binary/object.hpp b/src/binary/object.hpp index 3f71875..6da981c 100644 --- a/src/binary/object.hpp +++ b/src/binary/object.hpp @@ -138,7 +138,7 @@ namespace detail { inline object_frame resolve_safe_object_frame(const safe_object_frame& frame) { auto base = get_module_image_base(frame.object_path); if(base.is_error()) { - throw base.unwrap_error(); + throw base.unwrap_error(); // This throw is intentional } return { frame.raw_address, diff --git a/src/binary/pe.hpp b/src/binary/pe.hpp index 452f3b1..de442e0 100644 --- a/src/binary/pe.hpp +++ b/src/binary/pe.hpp @@ -26,7 +26,7 @@ namespace detail { } } - inline std::uintptr_t pe_get_module_image_base(const std::string& object_path) { + inline Result pe_get_module_image_base(const std::string& object_path) { // https://drive.google.com/file/d/0B3_wGJkuWLytbnIxY1J5WUs4MEk/view?pli=1&resourcekey=0-n5zZ2UW39xVTH8ZSu6C2aQ // https://0xrick.github.io/win-internals/pe3/ // Endianness should always be little for dos and pe headers @@ -37,18 +37,37 @@ namespace detail { throw file_error("Unable to read object file " + object_path); } auto magic = load_bytes>(file, 0); - VERIFY(std::memcmp(magic.data(), "MZ", 2) == 0, "File is not a PE file " + object_path); - DWORD e_lfanew = pe_byteswap_if_needed(load_bytes(file, 0x3c)); // dos header + 0x3c - DWORD nt_header_offset = e_lfanew; + if(!magic) { + return magic.unwrap_error(); + } + if(std::memcmp(magic.unwrap_value().data(), "MZ", 2) != 0) { + return internal_error("File is not a PE file " + object_path); + } + auto e_lfanew = load_bytes(file, 0x3c); // dos header + 0x3c + if(!e_lfanew) { + return e_lfanew.unwrap_error(); + } + DWORD nt_header_offset = pe_byteswap_if_needed(e_lfanew.unwrap_value()); auto signature = load_bytes>(file, nt_header_offset); // nt header + 0 - VERIFY(std::memcmp(signature.data(), "PE\0\0", 4) == 0, "File is not a PE file " + object_path); - WORD size_of_optional_header = pe_byteswap_if_needed( - load_bytes(file, nt_header_offset + 4 + 0x10) // file header + 0x10 - ); - VERIFY(size_of_optional_header != 0); - WORD optional_header_magic = pe_byteswap_if_needed( - load_bytes(file, nt_header_offset + 0x18) // optional header + 0x0 - ); + if(!signature) { + return signature.unwrap_error(); + } + if(std::memcmp(signature.unwrap_value().data(), "PE\0\0", 4) != 0) { + return internal_error("File is not a PE file " + object_path); + } + auto size_of_optional_header_raw = load_bytes(file, nt_header_offset + 4 + 0x10); // file header + 0x10 + if(!size_of_optional_header_raw) { + return size_of_optional_header_raw.unwrap_error(); + } + WORD size_of_optional_header = pe_byteswap_if_needed(size_of_optional_header_raw.unwrap_value()); + if(size_of_optional_header == 0) { + return internal_error("Unexpected optional header size for PE file"); + } + auto optional_header_magic_raw = load_bytes(file, nt_header_offset + 0x18); // optional header + 0x0 + if(!optional_header_magic_raw) { + return optional_header_magic_raw.unwrap_error(); + } + WORD optional_header_magic = pe_byteswap_if_needed(optional_header_magic_raw.unwrap_value()); VERIFY( optional_header_magic == IMAGE_NT_OPTIONAL_HDR_MAGIC, "PE file does not match expected bit-mode " + object_path @@ -56,19 +75,19 @@ namespace detail { // finally get image base if(optional_header_magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) { // 32 bit - return to( - pe_byteswap_if_needed( - load_bytes(file, nt_header_offset + 0x18 + 0x1c) // optional header + 0x1c - ) - ); + auto bytes = load_bytes(file, nt_header_offset + 0x18 + 0x1c); // optional header + 0x1c + if(!bytes) { + return bytes.unwrap_error(); + } + return to(pe_byteswap_if_needed(bytes.unwrap_value())); } else { // 64 bit // I get an "error: 'QWORD' was not declared in this scope" for some reason when using QWORD - return to( - pe_byteswap_if_needed( - load_bytes(file, nt_header_offset + 0x18 + 0x18) // optional header + 0x18 - ) - ); + auto bytes = load_bytes(file, nt_header_offset + 0x18 + 0x18); // optional header + 0x18 + if(!bytes) { + return bytes.unwrap_error(); + } + return to(pe_byteswap_if_needed(bytes.unwrap_value())); } } }