1 files changed, 10 insertions, 13 deletions

diff --git a/src/loader.c b/src/loader.c
index 2e79e67..8f2ac28 100644
--- a/src/loader.c
+++ b/src/loader.c
@@ -133,10 +133,8 @@ size_t str_len(const char* str) {
     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;
-
-    for (u32 index = 0U; index < min_len; index++) {
+i32 str_cmpn(const i8* str1, const i8* str2, size_t count) {
+    for (u32 index = 0U; index < count; index++) {
         if (str1[index] < str2[index]) {
             return -1;
         } else if (str1[index] > str2[index]) {
@@ -147,10 +145,8 @@ i32 str_cmp(const i8* str1, size_t str1_len, const i8* str2, size_t str2_len) {
     return 0;
 }
 
-i32 wstr_cmp(const u16* str1, size_t str1_len, const u16* str2, size_t str2_len) {
-    u32 min_len = str1_len < str2_len ? str1_len : str2_len;
-
-    for (u32 index = 0U; index < min_len; index++) {
+i32 wstr_cmpn(const u16* str1, const u16* str2, size_t count) {
+    for (u32 index = 0U; index < count; index++) {
         if (str1[index] < str2[index]) {
             return -1;
         } else if (str1[index] > str2[index]) {
@@ -174,8 +170,9 @@ void* get_module(u16* name, size_t name_size) {
         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(name, name_size / sizeof(u16), BaseDllName->Buffer, BaseDllName->Length / sizeof(u16)) == 0) {
+        if (BaseDllName->Buffer && ((name_size / sizeof(u16)) - 1) == (BaseDllName->Length / sizeof(u16))) {
+            if (wstr_cmpn(name, BaseDllName->Buffer, BaseDllName->Length / sizeof(u16)) == 0) {
+
                 return table_entry->DllBase;
             }
         }
@@ -183,6 +180,7 @@ void* get_module(u16* name, size_t name_size) {
         cur_entry = cur_entry->Flink;
     } while(cur_entry != first_entry);
 
+    assert(FALSE);
     return NULL;
 }
 
@@ -198,13 +196,12 @@ 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);
-
-        if (str_cmp((i8*)name, length, (i8*)func_name, func_name_len) == 0) {
+        if (str_cmpn((i8*)name, (i8*)func_name, func_name_len - 1U) == 0) {
             return (void*)(base_addr + addresses[ordinals[name_index]]);
         }
     }
 
+    assert(FALSE);
     return NULL;
 }