Start refactoring symbols management

1 files changed, 94 insertions, 18 deletions

diff --git a/src/events.c b/src/events.c
index 7438309..4f2115a 100644
--- a/src/events.c
+++ b/src/events.c
@@ -51,14 +51,42 @@ void load_symbols(const void* base_addr, struct module* module) {
     print("Module Base Address %xu\n", base_addr);
     print("Module Size %xu\n", module->nt_header.OptionalHeader.SizeOfImage);
 
+#ifndef HASHMAP
+    // Exception directory
+    IMAGE_DATA_DIRECTORY exception_entry = module->nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
+
+    if (exception_entry.Size == 0U) {
+        print("Exception directory is empty in module %s\n", module->name);
+    }
+
+    RUNTIME_FUNCTION* exception_table = VirtualAlloc(NULL, exception_entry.Size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
+    if (!read_memory((const void*)((u64)base_addr + (u64)exception_entry.VirtualAddress), (void*)exception_table, exception_entry.Size)) {
+        print("Cannot read exception table for module %s at address %x\n", (char*)module->name, base_addr);
+        return;
+    }
+
+    u32 pdata_entries = exception_entry.Size / sizeof(exception_table[0U]);
+    module->functions_name = VirtualAlloc(NULL, pdata_entries * sizeof(module->functions_name[0U]), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    module->functions_start = VirtualAlloc(NULL, pdata_entries * sizeof(module->functions_start[0U]), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    module->functions_end = VirtualAlloc(NULL, pdata_entries * sizeof(module->functions_end[0U]), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    module->functions_unwind = VirtualAlloc(NULL, pdata_entries * sizeof(module->functions_unwind[0U]), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    module->functions_count = pdata_entries;
+    for (u32 index = 0U; index < pdata_entries; index++) {
+        RUNTIME_FUNCTION cur_entry = exception_table[index];
+        module->functions_start[index] = cur_entry.BeginAddress;
+        module->functions_end[index] = cur_entry.EndAddress;
+        module->functions_unwind[index] = cur_entry.UnwindData;
+        module->functions_name[index] = (u8*)dummy_name;
+    }
+#endif
+
     // Export table
     IMAGE_DATA_DIRECTORY export_table = module->nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
     IMAGE_EXPORT_DIRECTORY export_dir = { 0U };
     const void* export_dir_start = (const void*)((u64)base_addr + (u64)export_table.VirtualAddress);
     const void* export_dir_end = (const void*)((u64)export_dir_start + export_table.Size);
-
     if (export_table.Size != 0U) {
-        if (!read_memory((const void*)((u64)base_addr + (u64)export_table.VirtualAddress), (void*)&export_dir, sizeof(export_dir))) {
+        if (!read_memory(export_dir_start, (void*)&export_dir, sizeof(export_dir))) {
             print("Cannot read export dir for module %s at address %x\n", (char*)module->name, base_addr);
             return;
         }
@@ -71,8 +99,8 @@ void load_symbols(const void* base_addr, struct module* module) {
         }
 
         size_t names_size = export_dir.NumberOfNames * sizeof(u32);
-        u32* names = VirtualAlloc(NULL, names_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
-        if (!read_memory((const void*)((u64)base_addr + (u64)export_dir.AddressOfNames), (void*)names, names_size)) {
+        u32* names_address = VirtualAlloc(NULL, names_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+        if (!read_memory((const void*)((u64)base_addr + (u64)export_dir.AddressOfNames), (void*)names_address, names_size)) {
             print("Cannot read names from export directory for module %s at address %x\n", (char*)module->name, base_addr);
             return;
         }
@@ -84,6 +112,52 @@ void load_symbols(const void* base_addr, struct module* module) {
             return;
         }
 
+#ifndef HASHMAP
+        // u32 global_functions_index = 0U;
+        for (u32 i = 0U; i < export_dir.NumberOfFunctions; i++) {
+            u64 address = (u64)base_addr + addresses[i];
+            u32 name_index = 0U;
+            u8* name = string_stack_current;
+            size_t length = 0U;
+            for (; name_index < export_dir.NumberOfNames; name_index++) {
+                if (ordinals[name_index] == i) {
+                    // TODO: Add string management code
+                    assert(string_stack_current - 16U * 4096U < string_stack_base);
+                    void* name_address = (void*)((u64)module->base_addr + names_address[name_index]);
+                    read_memory(name_address, name, 256U);
+                    StringCchLengthA((PCNZCH)name, STRSAFE_MAX_CCH, &length);
+                    break;
+                }
+            }
+
+            if (name_index == export_dir.NumberOfNames) {
+                print("%u : Function has no name, function address %xu\n", i, addresses[i]);
+                continue;
+            }
+
+            if (export_dir_start < (void*)address && (void*)address < export_dir_end) {
+                print("%u : Fowarded function : %s\n", i, name);
+                continue;
+            }
+
+            u32 global_functions_index = 0U;
+            for(; global_functions_index < module->functions_count; global_functions_index++) {
+                if (module->functions_start[global_functions_index] == addresses[i]) {
+                    string_stack_current[length] = '\0';
+                    module->functions_name[global_functions_index] = name;
+                    string_stack_current += length + 1U;
+                    print("%u : Function name : %s, function address %xu\n", i, name, addresses[i]);
+                    break;
+                }
+            }
+
+            if (global_functions_index == module->functions_count) {
+                print("Function not found : %s!\n", string_stack_current);
+            }
+        }
+#endif
+
+#ifdef HASHMAP
         print("Found %u functions\n", export_dir.NumberOfFunctions);
         for (u32 i = 0U; i < export_dir.NumberOfFunctions; i++) {
             char func_name[256U] = { 0 };
@@ -95,7 +169,7 @@ void load_symbols(const void* base_addr, struct module* module) {
                 }
             }
 
-            u64 name_address = (u64)base_addr + names[name_index];
+            u64 name_address = (u64)base_addr + names_address[name_index];
             read_memory((void*)name_address, func_name, sizeof(func_name));
 
             if (export_dir_start < (void*)func_address && (void*)func_address < export_dir_end) {
@@ -108,10 +182,11 @@ void load_symbols(const void* base_addr, struct module* module) {
                 insert_elem(&module->symbols, addresses[i], 0U, (u64)dummy_name);
             }
         }
+#endif
 
-        VirtualFree(names, names_size, MEM_RELEASE);
-        VirtualFree(addresses, addresses_size, MEM_RELEASE);
+        VirtualFree(names_address, names_size, MEM_RELEASE);
         VirtualFree(ordinals, ordinals_size, MEM_RELEASE);
+        VirtualFree(addresses, addresses_size, MEM_RELEASE);
     }
 
     // Private symbols
@@ -123,20 +198,18 @@ void load_symbols(const void* base_addr, struct module* module) {
         return;
     }
 
-    struct codeview_header {
-        u32 signature;
-        u32 guid[4U];
-        u32 age;
-        char path;
-    };
-
     for (u32 i = 0U; i < debug_table.Size / sizeof(debug_entries[0U]); i++) {
         if (debug_entries[i].Type != IMAGE_DEBUG_TYPE_CODEVIEW) {
             print("Image debug type different from CodeView\n");
             continue;
         }
 
-        struct codeview_header header = { 0U };
+        struct codeview_header {
+            u32 signature;
+            u32 guid[4U];
+            u32 age;
+            char path;
+        } header = { 0U };
         if (!read_memory((const void*)((u64)base_addr + (u64)debug_entries[i].AddressOfRawData), (void*)&header, debug_entries[i].SizeOfData)) {
             print("Cannot read code view header for module %s at address %x\n", (char*)module->name, base_addr);
         }
@@ -146,7 +219,8 @@ void load_symbols(const void* base_addr, struct module* module) {
     }
 
 
-    // Expection directory
+#ifdef HASHMAP
+    // Exception directory
     IMAGE_DATA_DIRECTORY exception_entry = module->nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
 
     if (exception_entry.Size == 0U) {
@@ -159,8 +233,9 @@ void load_symbols(const void* base_addr, struct module* module) {
         return;
     }
 
-    // struct UNWIND_INFO* u_info = VirtualAlloc(NULL, 4096U, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
-    for (u32 index = 0U; index < exception_entry.Size / sizeof(exception_table[0U]); index++) {
+    // Find unwind info offset
+    u32 pdata_entries = exception_entry.Size / sizeof(exception_table[0U]);
+    for (u32 index = 0U; index < pdata_entries; index++) {
         RUNTIME_FUNCTION cur_entry = exception_table[index];
 
         struct element* elem = find_elem(&module->symbols, cur_entry.BeginAddress);
@@ -255,6 +330,7 @@ void load_symbols(const void* base_addr, struct module* module) {
         //     }
         // }
     }
+#endif
 
     // VirtualFree(u_info, 4096U, MEM_RELEASE);
     VirtualFree(exception_table, exception_entry.Size, MEM_RELEASE);