Add a base_name field in module struct with module filenameHEADmaster

Implement function prologue unwind in stack walk code

3 files changed, 53 insertions, 25 deletions

diff --git a/src/events.c b/src/events.c
index ae78851..bef968a 100644
--- a/src/events.c
+++ b/src/events.c
@@ -312,12 +312,12 @@ void read_name_ansi(const void* name_addr, const char* name, size_t name_size) {
 
 void load_symbols(const void* base_addr, struct module* module) {
     if (!read_memory(base_addr, (void*)&module->dos_header, sizeof(module->dos_header))) {
-        print("Cannot read DOS header for module %s at address %x\n", (char*)module->name, base_addr);
+        print("Cannot read DOS header for module %s at address %x\n", (char*)module->base_name, base_addr);
         return;
     }
 
     if (!read_memory((const void*)((u64)base_addr + module->dos_header.e_lfanew), (void*)&module->nt_header, sizeof(module->nt_header))) {
-        print("Cannot read NT header for module %s at address %x\n", (char*)module->name, base_addr);
+        print("Cannot read NT header for module %s at address %x\n", (char*)module->base_name, base_addr);
         return;
     }
 
@@ -329,12 +329,12 @@ void load_symbols(const void* base_addr, struct module* module) {
     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);
+        print("Exception directory is empty in module %s\n", module->base_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);
+        print("Cannot read exception table for module %s at address %x\n", (char*)module->base_name, base_addr);
         return;
     }
 
@@ -360,28 +360,28 @@ void load_symbols(const void* base_addr, struct module* module) {
     const void* export_dir_end = (const void*)((u64)export_dir_start + export_table.Size);
     if (export_table.Size != 0U) {
         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);
+            print("Cannot read export dir for module %s at address %x\n", (char*)module->base_name, base_addr);
             return;
         }
 
         size_t ordinals_size = export_dir.NumberOfNames * sizeof(u16);
         u16* ordinals = VirtualAlloc(NULL, ordinals_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
         if (!read_memory((const void*)((u64)base_addr + (u64)export_dir.AddressOfNameOrdinals), (void*)ordinals, ordinals_size)) {
-            print("Cannot read name ordinals from export directory for module %s at address %x\n", (char*)module->name, base_addr);
+            print("Cannot read name ordinals from export directory for module %s at address %x\n", (char*)module->base_name, base_addr);
             return;
         }
 
         size_t names_size = export_dir.NumberOfNames * sizeof(u32);
         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);
+            print("Cannot read names from export directory for module %s at address %x\n", (char*)module->base_name, base_addr);
             return;
         }
 
         size_t addresses_size = export_dir.NumberOfFunctions * sizeof(u32);
         u32* addresses = VirtualAlloc(NULL, addresses_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
         if (!read_memory((const void*)((u64)base_addr + (u64)export_dir.AddressOfFunctions), (void*)addresses, addresses_size)) {
-            print("Cannot read functions addresses from export directory for module %s at address %x\n", (char*)module->name, base_addr);
+            print("Cannot read functions addresses from export directory for module %s at address %x\n", (char*)module->base_name, base_addr);
             return;
         }
 
@@ -468,7 +468,7 @@ void load_symbols(const void* base_addr, struct module* module) {
     IMAGE_DEBUG_DIRECTORY debug_entries[20U] = { 0U };
 
     if (debug_table.VirtualAddress && !read_memory((const void*)((u64)base_addr + (u64)debug_table.VirtualAddress), debug_entries, debug_table.Size)) {
-        print("Cannot read image debug infos for module %s at address %x\n", module->name, base_addr);
+        print("Cannot read image debug infos for module %s at address %x\n", module->base_name, base_addr);
         return;
     }
 
@@ -485,7 +485,7 @@ void load_symbols(const void* base_addr, struct module* module) {
             i8 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);
+            print("Cannot read code view header for module %s at address %x\n", (char*)module->base_name, base_addr);
         }
 
         print("PDB path : %s\n", &header.path);
@@ -625,8 +625,21 @@ void handle_create_process(void) {
     // module->base_addr = proc_dbg_info.lpBaseOfImage;
     // read_name_uni(proc_dbg_info.lpImageName, (const char*)module->name, sizeof(module->name));
     module->base_addr = proc_dbg_info.lpBaseOfImage;
-    GetFinalPathNameByHandleA(proc_dbg_info.hFile, module->name, sizeof(module->name), FILE_NAME_NORMALIZED);
-    print("%s%s\n", "Loading Process : ", (char*)module->name);
+    module->full_name_len = GetFinalPathNameByHandleA(proc_dbg_info.hFile, module->full_name, sizeof(module->full_name), FILE_NAME_NORMALIZED);
+
+    i8* cur_char = &module->full_name[module->full_name_len];
+    while (*cur_char != '\\' && cur_char != module->full_name) {
+        cur_char--;
+    }
+
+    if (cur_char == module->full_name) {
+        module->base_name = module->full_name;
+        module->base_name_len = module->full_name_len;
+    } else {
+        module->base_name = cur_char + 1U;
+        module->base_name_len = &module->full_name[module->full_name_len] - module->base_name;
+    }
+    print("%s%s\n", "Loading Process : ", (char*)module->base_name);
 
     load_symbols((const void*)proc_dbg_info.lpBaseOfImage, module);
 }
@@ -672,9 +685,23 @@ void handle_load_dll(void) {
 
     struct module* module = &modules[add_module()];
     module->base_addr = dll_base_addr;
-    read_name_uni(dll_name_addr, module->name, sizeof(module->name));
+    read_name_uni(dll_name_addr, module->full_name, sizeof(module->full_name));
+    module->full_name_len = str_len((char*)module->full_name);
+
+    i8* cur_char = &module->full_name[module->full_name_len];
+    while (*cur_char != '\\' && cur_char != module->full_name) {
+        cur_char--;
+    }
+
+    if (cur_char == module->full_name) {
+        module->base_name = module->full_name;
+        module->base_name_len = module->full_name_len;
+    } else {
+        module->base_name = cur_char + 1U;
+        module->base_name_len = &module->full_name[module->full_name_len] - module->base_name;
+    }
 
-    print("%s%s\n", "Loading DLL : ", (char*)module->name);
+    print("%s%s\n", "Loading DLL : ", (char*)module->base_name);
 
     load_symbols(dll_base_addr, module);
 }
diff --git a/src/parser.c b/src/parser.c
index 8dcbc76..4991852 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -377,7 +377,7 @@ void run(struct ast_node* root) {
         for (u8 index = 0U; index < 64U; index++) {
             if (((1ULL << index) & modules_list) == 0U) {
                 struct module* module = &modules[index];
-                print("Module Name : %s\n", module->name);
+                print("Module Name : %s\n", module->base_name);
                 print("Module Base : %x\n", module->base_addr);
                 print("Module Size : %x\n\n", module->nt_header.OptionalHeader.SizeOfImage);
             }
diff --git a/src/stack_walk.c b/src/stack_walk.c
index a161f6f..db5f236 100644
--- a/src/stack_walk.c
+++ b/src/stack_walk.c
@@ -266,7 +266,7 @@ void walk_stack(void* ctx) {
         u64 unwind_offset = module->functions_unwind[sym_index];
         u32 start = module->functions_start[sym_index];
 #endif
-        print("%s!%s\n", module->name, name);
+        print("%s!%s\n", module->base_name, name);
         // print("rsp value is %xu\n", rsp);
         // print("rbp value is %xu\n", rbp);
 
@@ -278,18 +278,19 @@ void walk_stack(void* ctx) {
         size_t total_size = sizeof(*u_info) + (u_info->count_of_unwind_codes - 1U) * sizeof(struct UNWIND_CODE);
         assert(total_size < 4096U);
 
-        // TODO: Do this once for the first function
-        if (start <= addr && addr < u_info->size_of_prolog) {
+        u64 func_start = (u64)module->base_addr + start;
+        u32 code_index = 0U;
+        // Undo prolog
+        if (func_start <= addr && addr < func_start + u_info->size_of_prolog) {
+            u32 RIP_offset = addr - func_start;
             // print("RIP is in prolog\n");
-        } else if (1) {
-            // print("RIP is in epilog\n");
-            // u8* byte_code = (u8*)addr;
-            addr = (u64)module->base_addr + start + u_info->size_of_prolog;
-        } else {
-            // print("RIP is in body\n");
+            struct UNWIND_CODE* cur_code = &u_info->unwind_code[code_index];
+            while (cur_code->offset_in_prolog != RIP_offset){
+                cur_code = &u_info->unwind_code[++code_index];
+            } 
         }
 
-        for (u32 code_index = 0U; code_index < u_info->count_of_unwind_codes;) {
+        for (; code_index < u_info->count_of_unwind_codes;) {
             struct UNWIND_CODE* cur_code = &u_info->unwind_code[code_index];
             switch (cur_code->unwind_operation_code) {
                 case UWOP_PUSH_NONVOL: