From d55874d936cf3fe6c620c9fffc7ca7941091cc2f Mon Sep 17 00:00:00 2001
From: hasherezade <hasherezade@gmail.com>
Date: Mon, 21 Oct 2024 17:42:23 -0700
Subject: [PATCH] [FEATURE] Allow to specify desired base of loading

---
 libpeconv/include/peconv/pe_loader.h |  8 ++++----
 libpeconv/src/pe_loader.cpp          | 13 ++++++-------
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/libpeconv/include/peconv/pe_loader.h b/libpeconv/include/peconv/pe_loader.h
index c16d2e3a2..3df5c9eb9 100644
--- a/libpeconv/include/peconv/pe_loader.h
+++ b/libpeconv/include/peconv/pe_loader.h
@@ -16,7 +16,7 @@ namespace peconv {
     If the relocate flag is true, applies relocations. Does not load imports.
     Automatically allocates buffer of the needed size (the size is returned in outputSize). The buffer can be freed by the function free_pe_buffer.
     */
-    BYTE* load_pe_module(BYTE* payload_raw, size_t r_size, OUT size_t &v_size, bool executable, bool relocate);
+    BYTE* load_pe_module(BYTE* payload_raw, size_t r_size, OUT size_t &v_size, bool executable, bool relocate, ULONGLONG desired_base = 0);
 
     /**
     Reads PE from the given file into memory and maps it into vitual format.
@@ -25,18 +25,18 @@ namespace peconv {
     If the relocate flag is true, applies relocations. Does not load imports.
     Automatically allocates buffer of the needed size (the size is returned in outputSize). The buffer can be freed by the function free_pe_buffer.
     */
-    BYTE* load_pe_module(LPCTSTR filename, OUT size_t &v_size, bool executable, bool relocate);
+    BYTE* load_pe_module(LPCTSTR filename, OUT size_t &v_size, bool executable, bool relocate, ULONGLONG desired_base = 0);
 
     /**
     Loads full PE from the raw buffer in a way in which it can be directly executed: remaps to virual format, applies relocations, loads imports.
     Allows for supplying custom function resolver.
     */
-    BYTE* load_pe_executable(BYTE* payload_raw, size_t r_size, OUT size_t &v_size, t_function_resolver* import_resolver=NULL);
+    BYTE* load_pe_executable(BYTE* payload_raw, size_t r_size, OUT size_t &v_size, t_function_resolver* import_resolver = nullptr, ULONGLONG desired_base = 0);
 
     /**
     Loads full PE from file in a way in which it can be directly executed: remaps to virtual format, applies relocations, loads imports.
     Allows for supplying custom function resolver.
     */
-    BYTE* load_pe_executable(LPCTSTR filename, OUT size_t &v_size, t_function_resolver* import_resolver=NULL);
+    BYTE* load_pe_executable(LPCTSTR filename, OUT size_t &v_size, t_function_resolver* import_resolver = nullptr);
 
 };// namespace peconv
diff --git a/libpeconv/src/pe_loader.cpp b/libpeconv/src/pe_loader.cpp
index ab8bed8a2..db405af46 100644
--- a/libpeconv/src/pe_loader.cpp
+++ b/libpeconv/src/pe_loader.cpp
@@ -31,7 +31,7 @@ namespace peconv {
     }
 };
 
-BYTE* peconv::load_pe_module(BYTE* dllRawData, size_t r_size, OUT size_t &v_size, bool executable, bool relocate)
+BYTE* peconv::load_pe_module(BYTE* dllRawData, size_t r_size, OUT size_t &v_size, bool executable, bool relocate, ULONGLONG desired_base)
 {
     if (!peconv::get_nt_hdrs(dllRawData, r_size)) {
         return NULL;
@@ -39,8 +39,7 @@ BYTE* peconv::load_pe_module(BYTE* dllRawData, size_t r_size, OUT size_t &v_size
     if (peconv::get_sections_count(dllRawData, r_size) == 0) {
         return load_no_sec_pe(dllRawData, r_size, v_size, executable);
     }
-    // by default, allow to load the PE at any base:
-    ULONGLONG desired_base = NULL;
+    // by default, allow to load the PE at the supplied base
     // if relocating is required, but the PE has no relocation table...
     if (relocate && !has_relocations(dllRawData)) {
         // ...enforce loading the PE image at its default base (so that it will need no relocations)
@@ -62,7 +61,7 @@ BYTE* peconv::load_pe_module(BYTE* dllRawData, size_t r_size, OUT size_t &v_size
     return mappedDLL;
 }
 
-BYTE* peconv::load_pe_module(LPCTSTR filename, OUT size_t &v_size, bool executable, bool relocate)
+BYTE* peconv::load_pe_module(LPCTSTR filename, OUT size_t &v_size, bool executable, bool relocate, ULONGLONG desired_base)
 {
     size_t r_size = 0;
     BYTE *dllRawData = load_file(filename, r_size);
@@ -72,14 +71,14 @@ BYTE* peconv::load_pe_module(LPCTSTR filename, OUT size_t &v_size, bool executab
 #endif
         return NULL;
     }
-    BYTE* mappedPE = load_pe_module(dllRawData, r_size, v_size, executable, relocate);
+    BYTE* mappedPE = load_pe_module(dllRawData, r_size, v_size, executable, relocate, desired_base);
     free_file(dllRawData);
     return mappedPE;
 }
 
-BYTE* peconv::load_pe_executable(BYTE* dllRawData, size_t r_size, OUT size_t &v_size, t_function_resolver* import_resolver)
+BYTE* peconv::load_pe_executable(BYTE* dllRawData, size_t r_size, OUT size_t &v_size, t_function_resolver* import_resolver, ULONGLONG desired_base)
 {
-    BYTE* loaded_pe = load_pe_module(dllRawData, r_size, v_size, true, true);
+    BYTE* loaded_pe = load_pe_module(dllRawData, r_size, v_size, true, true, desired_base);
     if (!loaded_pe) {
         std::cerr << "[-] Loading failed!\n";
         return NULL;