Fix runtime crashes due to loaded functions

1 files changed, 33 insertions, 34 deletions

diff --git a/src/loader.c b/src/loader.c
index da98364..2e79e67 100644
--- a/src/loader.c
+++ b/src/loader.c
@@ -114,11 +114,25 @@ typedef struct _IMAGE_EXPORT_DIRECTORY {
 #define KERNEL_BASE(name)   \
     PFN_##name name = NULL;
 
+#define NTDLL(name)         \
+    PFN_##name name = NULL;
+
 #include "funcs.h"
 
+#undef NTDLL
 #undef KERNEL_BASE
 #undef KERNEL32
 
+size_t str_len(const char* str) {
+    const char* c = str;
+    size_t count = 0U;
+    while (*c != '\0') {
+        c++;
+        count++;
+    }
+    return count;
+}
+
 i32 str_cmp(const i8* str1, size_t str1_len, const i8* str2, size_t str2_len) {
     u32 min_len = str1_len < str2_len ? str1_len : str2_len;
 
@@ -151,46 +165,23 @@ PEB* get_PEB(void) {
     return (PEB*)__readgsqword(__builtin_offsetof(TEB, ProcessEnvironmentBlock));
 }
 
-void* get_kernel32(void) {
-    const u16 kernel32_str[] = L"KERNEL32.DLL";
+void* get_module(u16* name, size_t name_size) {
     PEB* PEB = get_PEB();
-    LIST_ENTRY* last_entry = PEB->LoaderData->InMemoryOrderModuleList.Blink;
-    LIST_ENTRY* cur_entry = PEB->LoaderData->InMemoryOrderModuleList.Flink;
+    LIST_ENTRY* first_entry = PEB->LoaderData->InMemoryOrderModuleList.Flink;
+    LIST_ENTRY* cur_entry = first_entry;
 
-    while(cur_entry != last_entry) {
+    do {
         LDR_DATA_TABLE_ENTRY* table_entry = (LDR_DATA_TABLE_ENTRY*)((u8*)cur_entry - 0x10);
         UNICODE_STRING* BaseDllName =  &table_entry->BaseDllName;
 
         if (BaseDllName->Buffer) {
-            if (wstr_cmp(kernel32_str, sizeof(kernel32_str) / sizeof(u16), BaseDllName->Buffer, BaseDllName->Length / sizeof(u16)) == 0) {
+            if (wstr_cmp(name, name_size / sizeof(u16), BaseDllName->Buffer, BaseDllName->Length / sizeof(u16)) == 0) {
                 return table_entry->DllBase;
             }
         }
 
         cur_entry = cur_entry->Flink;
-    }
-
-    return NULL;
-}
-
-void* get_kernel_base(void) {
-    const u16 kernel32_str[] = L"KERNELBASE.dll";
-    PEB* PEB = get_PEB();
-    LIST_ENTRY* last_entry = PEB->LoaderData->InMemoryOrderModuleList.Blink;
-    LIST_ENTRY* cur_entry = PEB->LoaderData->InMemoryOrderModuleList.Flink;
-
-    while(cur_entry != last_entry) {
-        LDR_DATA_TABLE_ENTRY* table_entry = (LDR_DATA_TABLE_ENTRY*)((u8*)cur_entry - 0x10);
-        UNICODE_STRING* BaseDllName =  &table_entry->BaseDllName;
-
-        if (BaseDllName->Buffer) {
-            if (wstr_cmp(kernel32_str, sizeof(kernel32_str) / sizeof(u16), BaseDllName->Buffer, BaseDllName->Length / sizeof(u16)) == 0) {
-                return table_entry->DllBase;
-            }
-        }
-
-        cur_entry = cur_entry->Flink;
-    }
+    } while(cur_entry != first_entry);
 
     return NULL;
 }
@@ -207,7 +198,7 @@ void* load_func(u64 base_addr, const char* func_name, u32 func_name_len) {
 
     for (u32 name_index = 0U; name_index < export_dir->NumberOfNames; name_index++) {
         const char* name = (const char*)(base_addr + names_address[name_index]);
-        u32 length = str_len(name, 256U);
+        u32 length = str_len(name);
 
         if (str_cmp((i8*)name, length, (i8*)func_name, func_name_len) == 0) {
             return (void*)(base_addr + addresses[ordinals[name_index]]);
@@ -218,17 +209,25 @@ void* load_func(u64 base_addr, const char* func_name, u32 func_name_len) {
 }
 
 void load_funcs(void) {
-    void* kernel32 = get_kernel32();
-    void* kernel_base = get_kernel_base();
+    u16 kernel32_str[] = L"KERNEL32.DLL";
+    u16 kernel_base_str[] = L"KERNELBASE.dll";
+    u16 ntdll_str[] = L"ntdll.dll";
+    void* kernel32 = get_module(kernel32_str, sizeof(kernel32_str));
+    void* kernel_base = get_module(kernel_base_str, sizeof(kernel_base_str));
+    void* ntdll = get_module(ntdll_str, sizeof(ntdll_str));
 
-#define KERNEL32(name)                                                          \
+#define KERNEL32(name)                                                    \
     name = (PFN_##name)load_func((u64)kernel32, #name, sizeof(#name));
 
-#define KERNEL_BASE(name)                                                       \
+#define KERNEL_BASE(name)                                                 \
     name = (PFN_##name)load_func((u64)kernel_base, #name, sizeof(#name));
 
+#define NTDLL(name)                                                       \
+    name = (PFN_##name)load_func((u64)ntdll, #name, sizeof(#name));
+
 #include "funcs.h"
 
+#undef NTDLL
 #undef KERNEL_BASE
 #undef KERNEL32