Logic for PE, and a couple tweaks

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<std::size_t Bits>

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;

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,

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<std::uintptr_t, internal_error> 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<std::array<char, 2>>(file, 0);
-        VERIFY(std::memcmp(magic.data(), "MZ", 2) == 0, "File is not a PE file " + object_path);
+        if(!magic) {
-        DWORD e_lfanew = pe_byteswap_if_needed(load_bytes<DWORD>(file, 0x3c)); // dos header + 0x3c
+            return magic.unwrap_error();
-        DWORD nt_header_offset = e_lfanew;
+        }
+        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<DWORD>(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<std::array<char, 4>>(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);
+        if(!signature) {
-        WORD size_of_optional_header = pe_byteswap_if_needed(
+            return signature.unwrap_error();
-            load_bytes<WORD>(file, nt_header_offset + 4 + 0x10) // file header + 0x10
+        }
-        );
+        if(std::memcmp(signature.unwrap_value().data(), "PE\0\0", 4) != 0) {
-        VERIFY(size_of_optional_header != 0);
+            return internal_error("File is not a PE file " + object_path);
-        WORD optional_header_magic = pe_byteswap_if_needed(
+        }
-            load_bytes<WORD>(file, nt_header_offset + 0x18) // optional header + 0x0
+        auto size_of_optional_header_raw = load_bytes<WORD>(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<WORD>(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<std::uintptr_t>(
+            auto bytes = load_bytes<DWORD>(file, nt_header_offset + 0x18 + 0x1c); // optional header + 0x1c
-                pe_byteswap_if_needed(
+            if(!bytes) {
-                    load_bytes<DWORD>(file, nt_header_offset + 0x18 + 0x1c) // optional header + 0x1c
+                return bytes.unwrap_error();
-                )
+            }
-            );
+            return to<std::uintptr_t>(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<std::uintptr_t>(
+            auto bytes = load_bytes<unsigned __int64>(file, nt_header_offset + 0x18 + 0x18); // optional header + 0x18
-                pe_byteswap_if_needed(
+            if(!bytes) {
-                    load_bytes<unsigned __int64>(file, nt_header_offset + 0x18 + 0x18) // optional header + 0x18
+                return bytes.unwrap_error();
-                )
+            }
-            );
+            return to<std::uintptr_t>(pe_byteswap_if_needed(bytes.unwrap_value()));
         }
     }
 }