From 6c69ae0fc3ec7453bb4c0119bc23b7449526be74 Mon Sep 17 00:00:00 2001 From: ladevieq Date: Tue, 21 Jan 2025 23:22:56 +0100 Subject: Stack walking --- src/events.c | 169 ++++++++++++++++++++++++++++++--- src/hashmap.c | 113 ++++++++++++++++++++++ src/main.c | 16 +++- src/module.c | 5 + src/parser.c | 47 ++++++---- src/pdb.c | 14 ++- src/stack_walk.c | 282 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/test.c | 1 + 8 files changed, 609 insertions(+), 38 deletions(-) create mode 100644 src/hashmap.c create mode 100644 src/stack_walk.c diff --git a/src/events.c b/src/events.c index 61cc9dd..7438309 100644 --- a/src/events.c +++ b/src/events.c @@ -1,3 +1,5 @@ +const char* dummy_name = "Function_dummy_name"; + void read_name_uni(const void* name_addr, const char* name, size_t name_size) { u8* addr = NULL; if (!read_memory(name_addr, (void*)&addr, 8U)) { @@ -35,7 +37,7 @@ void read_name_ansi(const void* name_addr, const char* name, size_t name_size) { } } -void load_symbols(const void* base_addr, const struct module* module) { +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); return; @@ -61,29 +63,56 @@ void load_symbols(const void* base_addr, const struct module* module) { return; } - u16 name_ordinals[8192]; - if (!read_memory((const void*)((u64)base_addr + (u64)export_dir.AddressOfNameOrdinals), (void*)&name_ordinals, export_dir.NumberOfNames * sizeof(name_ordinals[0]))) { + 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); return; } - u32 names[8192]; - if (!read_memory((const void*)((u64)base_addr + (u64)export_dir.AddressOfNames), (void*)&names, export_dir.NumberOfNames * sizeof(names[0]))) { + 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)) { print("Cannot read names from export directory for module %s at address %x\n", (char*)module->name, base_addr); return; } - print("Found %u functions\n", export_dir.NumberOfNames); - for (u32 i = 0U; i < export_dir.NumberOfNames; i++) { - const char* func_name = (const char*)((u64)base_addr + names[name_ordinals[i]]); - if (export_dir_end < (void*)func_name || (void*)func_name < export_dir_start) { - print("%u : Fowarded function\n", i); + 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); + return; + } + + print("Found %u functions\n", export_dir.NumberOfFunctions); + for (u32 i = 0U; i < export_dir.NumberOfFunctions; i++) { + char func_name[256U] = { 0 }; + u64 func_address = (u64)base_addr + addresses[i]; + u32 name_index = 0U; + for (; name_index < export_dir.NumberOfNames; name_index++) { + if (ordinals[name_index] == i) { + break; + } + } + + u64 name_address = (u64)base_addr + names[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) { + print("%u : Fowarded function : %s\n", i, func_name); + } else if (name_index && name_index < export_dir.NumberOfNames) { + print("%u : Function name : %s, function address %xu\n", i, func_name, addresses[i]); + insert_elem(&module->symbols, addresses[i], 0U, name_address); } else { - print("%u : Function name : %s\n", i, func_name); + print("%u : Function has no name, function address %xu\n", i, addresses[i]); + insert_elem(&module->symbols, addresses[i], 0U, (u64)dummy_name); } } - } + VirtualFree(names, names_size, MEM_RELEASE); + VirtualFree(addresses, addresses_size, MEM_RELEASE); + VirtualFree(ordinals, ordinals_size, MEM_RELEASE); + } // Private symbols IMAGE_DATA_DIRECTORY debug_table = module->nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG]; @@ -113,8 +142,122 @@ void load_symbols(const void* base_addr, const struct module* module) { } print("PDB path : %s\n", &header.path); - open_pdb(&header.path); + open_pdb(&header.path, module); + } + + + // Expection 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; + } + + // 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++) { + RUNTIME_FUNCTION cur_entry = exception_table[index]; + + struct element* elem = find_elem(&module->symbols, cur_entry.BeginAddress); + if (elem == NULL) { + elem = insert_elem(&module->symbols, cur_entry.BeginAddress, cur_entry.EndAddress, 0ULL); + } else { + if (elem->name == 0U) { + print("function name %s\n", elem->name); + } + + elem->end = cur_entry.EndAddress; + } + + elem->unwind_offset = cur_entry.UnwindData; + + // if (!read_memory((const void*)((u64)base_addr + (u64)cur_entry.UnwindData), (void*)u_info, 4096U)) { + // print("Cannot read unwind info in for function %s\n", elem->name); + // continue; + // } + + // size_t total_size = sizeof(*u_info) + (u_info->count_of_unwind_codes - 1U) * sizeof(struct UNWIND_CODE); + // assert(total_size < 4096U); + + // for (u32 code_index = 0U; code_index < u_info->count_of_unwind_codes; code_index++) { + // struct UNWIND_CODE* cur_code = &u_info->unwind_code[code_index]; + // switch (cur_code->unwind_operation_code) { + // case UWOP_PUSH_NONVOL: + // print("push non volatile\n"); + // print("push "); + // print_reg(cur_code->op_info); + // break; + // case UWOP_ALLOC_LARGE: + // print("alloc large\n"); + // if (cur_code->op_info == 0U) { + // u16 size = *(u16*)&u_info->unwind_code[code_index + 1U]; + // if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + // } else { + // print("sub rsp %xu\n", size); + // } + // code_index++; + // } else { + // u32 size = *(u32*)&u_info->unwind_code[code_index + 1U]; + // if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + // } else { + // print("sub rsp %xu\n", size); + // } + // code_index += 2U; + // } + // break; + // case UWOP_ALLOC_SMALL: + // print("alloc small\n"); + // u32 size = cur_code->op_info * 8U + 8U; + // if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + // } else { + // print("sub rsp %xu\n", size); + // } + // break; + // case UWOP_SET_FPREG: + // print("set frame point reg\n"); + // print("mov rbp rsp - %xu\n", u_info->frame_register_offset * 16U); + // break; + // case UWOP_SAVE_NONVOL: + // print("save non volatile\n"); + // if (u_info->frame_register) { + // print("mov rbp "); + // } else { + // print("mov rsp "); + // } + // print_reg(cur_code->op_info); + // code_index++; + // break; + // case UWOP_SAVE_NONVOL_FAR: + // print("save non volatile far\n"); + // code_index += 2U; + // break; + // case UWOP_SAVE_XMM128: + // print("save xmm128\n"); + // code_index++; + // break; + // case UWOP_SAVE_XMM128_FAR: + // print("save xmm128 far\n"); + // code_index += 2U; + // break; + // case UWOP_PUSH_MACHFRAME: + // print("push mach frame\n"); + // break; + // default: + // print("bad operation code\n"); + // } + // } } + + // VirtualFree(u_info, 4096U, MEM_RELEASE); + VirtualFree(exception_table, exception_entry.Size, MEM_RELEASE); } void handle_create_process() { diff --git a/src/hashmap.c b/src/hashmap.c new file mode 100644 index 0000000..5b2d5a0 --- /dev/null +++ b/src/hashmap.c @@ -0,0 +1,113 @@ +struct element { + u64 name; + u32 start; + u32 end; + u32 unwind_offset; + u32 psl; +}; + +struct hashmap { + u32 capacity; + u32 size; + + struct element* arr; +}; + +static const u32 initial_size = 4096U * 100U; +static const u32 empty_psl = (u32)-1; +static const float max_load_factor = 0.75f; + +void resize(struct hashmap* map) { + void* alloc = VirtualAlloc((void*)(map->arr + 1U), map->capacity * sizeof(struct element), MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + + if (alloc == NULL) { + alloc = VirtualAlloc(NULL, map->capacity * 2U * sizeof(struct element), MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + map->capacity *= 2U; + // copy elements + rehashing + } else { + map->capacity *= 2U; + // rehashing + } +} + +struct hashmap create_unordered_flat() { + struct hashmap map = { + .capacity = initial_size / sizeof(struct element), + .size = 0U, + .arr = (struct element*)VirtualAlloc(NULL, initial_size * sizeof(struct element), MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE), + }; + + for (u32 index = 0U; index < map.capacity; index++) { + map.arr[index].name = 0ULL; + map.arr[index].start = 0U; + map.arr[index].end = 0U; + map.arr[index].psl = 0U; + } + return map; +} + +struct element* insert_elem(struct hashmap* map, u32 start, u32 end, u64 name) { + if (map->size >= map->capacity) { + return NULL; + } + float load = map->size / (float)map->capacity; + if (load >= max_load_factor) { + resize(map); + } + + u32 index = start % map->capacity; + u32 psl = 0U; + + while(map->arr[index + psl].name != 0ULL || (map->arr[index + psl].psl >= psl && map->arr[index + psl].psl != 0U)) { + psl++; + } + + map->arr[index + psl].name = name; + map->arr[index + psl].start = start; + map->arr[index + psl].end = end; + map->arr[index + psl].psl = psl; + map->size++; + + return &map->arr[index + psl]; +} + +void remove_elem(struct hashmap* map, u32 start) { + u32 index = start % map->capacity; + + while (map->arr[index].start != start) { + index++; + } + + while (map->arr[index].name != 0U && map->arr[index].psl != empty_psl && index < map->capacity) { + map->arr[index - 1U].psl = map->arr[index].psl - 1U; + map->arr[index - 1U].start = map->arr[index].start; + map->arr[index - 1U].end = map->arr[index].end; + map->arr[index - 1U].name = map->arr[index].name; + index++; + } + + map->arr[index].psl = empty_psl; +} + +struct element* find_elem(struct hashmap* map, u32 start) { + u32 index = start % map->capacity; + + while (index < map->capacity) { + if (map->arr[index].start == start) { + return &map->arr[index]; + } + index++; + } + + return NULL; +} + +void iterate(struct hashmap* map) { + u32 counter = 0U; + for (u32 index = 0U; index < map->capacity && counter < map->size; index++) { + if (map->arr[index].name) { + print("%s\n", map->arr[index].name); + counter++; + } + } +} diff --git a/src/main.c b/src/main.c index e71b1b7..32e4fa8 100644 --- a/src/main.c +++ b/src/main.c @@ -5,6 +5,9 @@ #include "types.h" #include "utils.c" +#include "hashmap.c" + +extern u32 _fltused = TRUE; static HANDLE h_input = NULL; static HANDLE h_output = NULL; @@ -15,6 +18,10 @@ static DWORD continue_status = DBG_EXCEPTION_NOT_HANDLED; static DEBUG_EVENT dbg_event; static PROCESS_INFORMATION proc_info = { 0 }; +// TODO: Perfect hashing table ? +u8* string_stack_base = NULL; +u8* string_stack_current = NULL; + /* Read debugged process memory at address */ u8 read_memory(const void* addr, void* buffer, size_t size) { size_t b_read = 0U; @@ -24,6 +31,7 @@ u8 read_memory(const void* addr, void* buffer, size_t size) { } #include "module.c" +#include "stack_walk.c" #include "parser.c" #include "pdb.c" #include "events.c" @@ -80,12 +88,14 @@ void dbg_loop() { int mainCRTStartup() { - STARTUPINFOW startup_info = { 0 }; + string_stack_current = string_stack_base = VirtualAlloc(NULL, 4096U * 16U, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + + STARTUPINFOA startup_info = { 0 }; startup_info.cb = sizeof(startup_info); - BOOL ret = CreateProcessW( + BOOL ret = CreateProcessA( NULL, - L"./test.exe", + "./test.exe", NULL, NULL, FALSE, diff --git a/src/module.c b/src/module.c index 8188665..ffff0ca 100644 --- a/src/module.c +++ b/src/module.c @@ -2,12 +2,15 @@ struct module { void* base_addr; IMAGE_DOS_HEADER dos_header; IMAGE_NT_HEADERS64 nt_header; + struct hashmap symbols; u8 name[256U]; }; struct module modules[64U] = { 0U }; u64 modules_list = ~0ULL; +// u64 module_symbols + u8 add_module() { // struct module* add_module() { assert(!modules_list); @@ -15,6 +18,8 @@ u8 add_module() { _BitScanForward64((unsigned long*)&module_index, modules_list); modules_list ^= (u64)(0x1 << module_index); + modules[module_index].symbols = create_unordered_flat(); + return (u8)module_index; }; diff --git a/src/parser.c b/src/parser.c index 5ec6b0d..937d640 100644 --- a/src/parser.c +++ b/src/parser.c @@ -3,6 +3,7 @@ enum TOKEN_TYPE { TOKEN_CONTINUE, TOKEN_CALLSTACK, TOKEN_LISTMODULE, + TOKEN_DUMP_SYMBOLS, TOKEN_PRINT_REG, TOKEN_STEP_INTO, TOKEN_DUMP_MEMORY, @@ -32,10 +33,13 @@ struct ast_node { struct ast_node* right; }; +#define DECLARE_COMMAND(x) { .start = x, .size = sizeof(x) } + static struct command commands[] = { { .start = "g", .size = 1 }, { .start = "k", .size = 1 }, { .start = "lm", .size = 2 }, + DECLARE_COMMAND("x"), { .start = "r", .size = 1 }, { .start = "t", .size = 1 }, { .start = "db", .size = 2 }, @@ -140,6 +144,13 @@ struct ast_node* parse_command(struct token* tokens, struct ast_node* nodes) { cur_node->left = NULL; cur_node->right = NULL; return cur_node; + case TOKEN_DUMP_SYMBOLS: + tokens++; + print("TOKEN_DUMP_SYMBOLS\n"); + cur_node->token = cur_token; + cur_node->left = NULL; + cur_node->right = NULL; + return cur_node; case TOKEN_PRINT_REG: tokens++; print("TOKEN_PRINT_REG\n"); @@ -158,8 +169,8 @@ struct ast_node* parse_command(struct token* tokens, struct ast_node* nodes) { tokens++; print("TOKEN_DUMP_MEMORY\n"); cur_node->token = cur_token; - cur_node->left = NULL; - cur_node->right = parse_primary(tokens, ++nodes); + cur_node->left = parse_primary(tokens++, ++nodes); + cur_node->right = parse_primary(tokens++, ++nodes); return cur_node; case TOKEN_QUIT: tokens++; @@ -202,18 +213,7 @@ void run(struct ast_node* root) { ); process_commands = FALSE; } else if (cur_node.token->type == TOKEN_CALLSTACK) { - for (u8 index = 0U; index < 64U; index++) { - if (((1ULL << index) & modules_list) == 0U) { - struct module* module = &modules[index]; - if ((u64)module->base_addr < ctx.Rip && ctx.Rip < ((u64)module->base_addr + module->nt_header.OptionalHeader.SizeOfImage)) { - print("Module name %s\n", module->name); - print("Module address %xu\n", (u64)module->base_addr); - print("Module size %xu\n", module->nt_header.OptionalHeader.SizeOfImage); - print("Function address %xu\n", ctx.Rip); - print("Function offset in module %xu\n", ctx.Rip - (u64)module->base_addr); - } - } - } + walk_stack(&ctx); } else if (cur_node.token->type == TOKEN_LISTMODULE) { for (u8 index = 0U; index < 64U; index++) { if (((1ULL << index) & modules_list) == 0U) { @@ -250,12 +250,13 @@ void run(struct ast_node* root) { is_open = FALSE; process_commands = FALSE; } else if (cur_node.token->type == TOKEN_DUMP_MEMORY) { - u64 addr = eval_number(cur_node.right); - u8 buf[16U] = { 0 }; + u64 addr = eval_number(cur_node.left); + u64 size = eval_number(cur_node.right); + u8* buf = VirtualAlloc(NULL, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); - if(read_memory((const void*)addr, (void*)buf, sizeof(buf))) { - for (u32 index = 0U; index < sizeof(buf); index++) { - print("%b ", buf[index]); + if(read_memory((const void*)addr, (void*)buf, size)) { + for (u32 index = 0U; index < size; index += 8U) { + print("%xb %xb %xb %xb %xb %xb %xb %xb\n", buf[index], buf[index + 1U], buf[index + 2U], buf[index + 3U], buf[index + 4U], buf[index + 5U], buf[index + 6U], buf[index + 7U]); } print("\n"); } else { @@ -263,6 +264,14 @@ void run(struct ast_node* root) { print("Unable to read memory at this address, error %u\n", GetLastError()); } + VirtualFree(buf, size, MEM_RELEASE); + } else if (cur_node.token->type == TOKEN_DUMP_SYMBOLS) { + for (u8 index = 0U; index < 64U; index++) { + if (((1ULL << index) & modules_list) == 0U) { + struct module* module = &modules[index]; + iterate(&module->symbols); + } + } } else { print("Command unknown\n"); process_commands = TRUE; diff --git a/src/pdb.c b/src/pdb.c index f646a85..3f89611 100644 --- a/src/pdb.c +++ b/src/pdb.c @@ -8,7 +8,7 @@ enum STREAM_INDEX { IPI, }; -void open_pdb(const char* pdb_path) { +void open_pdb(const char* pdb_path, struct module* module) { HANDLE pdb_file = CreateFileA(pdb_path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); if (pdb_file == INVALID_HANDLE_VALUE) { @@ -174,6 +174,14 @@ void open_pdb(const char* pdb_path) { } else if (rec->rectyp == S_GPROC32) { struct PROCSYM32* proc = (struct PROCSYM32*)cur_sym; print("PROC32 SYMBOL %s\n", proc->name); + // TODO: Add string management code + assert(string_stack_current - 16U * 4096U < string_stack_base); + size_t length = 0U; + StringCchLengthA((PCNZCH)proc->name, STRSAFE_MAX_CCH, &length); + StringCchCopyA((LPSTR)string_stack_current, STRSAFE_MAX_CCH, (LPCSTR)proc->name); + string_stack_current[length] = '\0'; + insert_elem(&module->symbols, proc->off + module->nt_header.OptionalHeader.BaseOfCode, proc->len + module->nt_header.OptionalHeader.BaseOfCode, (u64)string_stack_current); + string_stack_current += length + 1U; } else if (rec->rectyp == S_REGREL32) { struct REGREL32* regrel = (struct REGREL32*)cur_sym; print("REGREL32 SYMBOL %s\n", regrel->name); @@ -183,7 +191,7 @@ void open_pdb(const char* pdb_path) { print("COMPILE2 SYMBOL\n"); } else if (rec->rectyp == S_SECTION) { struct SECTIONSYM* section = (struct SECTIONSYM*)cur_sym; - print("SECTION SYMBOL %s\n", section->name); + print("SECTION SYMBOL index %u, name %s, rva %u, size %u\n", section->isec, section->name, section->rva, section->cb); } else if (rec->rectyp == S_COFFGROUP) { struct COFFGROUPSYM* coff = (struct COFFGROUPSYM*)cur_sym; print("COFFGROUP SYMBOL %s\n", coff->name); @@ -199,7 +207,7 @@ void open_pdb(const char* pdb_path) { buildinfo; print("BUILDINFO SYMBOL\n"); } else { - print("Unkown symbole kind %xu\n", rec->rectyp); + print("Unknown symbol kind %xu\n", rec->rectyp); } cur_sym = (u8*)NextSym((SYMTYPE*)cur_sym); } diff --git a/src/stack_walk.c b/src/stack_walk.c new file mode 100644 index 0000000..94a76e1 --- /dev/null +++ b/src/stack_walk.c @@ -0,0 +1,282 @@ +// enum UNWIND_INFO_FLAGS { +// UNW_FLAG_EHANDLER, +// UNW_FLAG_UHANDLER, +// UNW_FLAG_CHAININFO, +// }; + +enum UNWIND_OP_TYPE { + UWOP_PUSH_NONVOL, + UWOP_ALLOC_LARGE, + UWOP_ALLOC_SMALL, + UWOP_SET_FPREG, + UWOP_SAVE_NONVOL, + UWOP_SAVE_NONVOL_FAR, + UWOP_SAVE_XMM128 = 8U, + UWOP_SAVE_XMM128_FAR, + UWOP_PUSH_MACHFRAME, +}; + +struct UNWIND_CODE { + u8 offset_in_prolog; + u8 unwind_operation_code : 4; + u8 op_info : 4; +}; + +struct UNWIND_INFO { + u8 version : 3; + u8 flags : 5; + u8 size_of_prolog; + u8 count_of_unwind_codes; + u8 frame_register : 4; + u8 frame_register_offset : 4; + struct UNWIND_CODE unwind_code[1U]; + // Optional Exception handler of chained unwind info + // + // Exception handler + // ULONG address_of_exception_handler; + // language specific handler data; + // + // Chained unwind info + // ULONG function_start; + // ULONG function_end; + // ULONG unwind_info_address; +}; + +enum UNWIND_OP_INFO_TYPE { + RAX, + RCX, + RDX, + RBX, + RSP, + RBP, + RSI, + RDI, + R8, + R9, + R10, + R11, + R12, + R13, + R14, + R15, +}; + +struct module* get_addr_module(u64 addr) { + for (u8 index = 0U; index < 64U; index++) { + if (((1ULL << index) & modules_list) != 0U) { + continue; + } + + struct module* module = &modules[index]; + u64 offset = addr - (u64)module->base_addr; + if (0U <= offset && offset < module->nt_header.OptionalHeader.SizeOfImage) { + return module; + } + } + + return NULL; +} + +void print_reg(enum UNWIND_OP_INFO_TYPE op_info) { + switch (op_info) { + case RAX: + print("RAX\n"); + break; + case RCX: + print("RCX\n"); + break; + case RDX: + print("RDX\n"); + break; + case RBX: + print("RBX\n"); + break; + case RSP: + print("RSP\n"); + break; + case RBP: + print("RBP\n"); + break; + case RSI: + print("RSI\n"); + break; + case RDI: + print("RDI\n"); + break; + case R8: + print("R8\n"); + break; + case R9: + print("R9\n"); + break; + case R10: + print("R10\n"); + break; + case R11: + print("R11\n"); + break; + case R12: + print("R12\n"); + break; + case R13: + print("R13\n"); + break; + case R14: + print("R14\n"); + break; + case R15: + print("R15\n"); + break; + } +} + +struct element* find_function(struct module* module, u64 rip_offset) { + for (u32 sym_index = 0U; sym_index < module->symbols.capacity; sym_index++) { + struct element* cur_sym = &module->symbols.arr[sym_index]; + if (cur_sym->start != 0U && cur_sym->start <= rip_offset && rip_offset < cur_sym->end) { + return cur_sym; + } + } + return NULL; +} + +void walk_stack(CONTEXT* ctx) { + u64 addr = ctx->Rip; + u64 rsp = ctx->Rsp; + u64 rbp = ctx->Rbp; + struct module* module = NULL; + struct UNWIND_INFO* u_info = VirtualAlloc(NULL, 4096U, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); + while (addr) { + module = get_addr_module(addr); + // print("Module name %s\n", module->name); + // print("Module address %xu\n", (u64)module->base_addr); + // print("Module size %xu\n", module->nt_header.OptionalHeader.SizeOfImage); + // print("Module code address %xu\n", (u64)module->base_addr + module->nt_header.OptionalHeader.BaseOfCode); + // print("Module code size %xu\n", module->nt_header.OptionalHeader.SizeOfCode); + // print("Function address %xu\n", addr); + // print("Function offset in module %xu\n", addr - (u64)module->base_addr); + + struct element* sym = find_function(module, addr - (u64)module->base_addr); + if (sym == NULL) { + print("No symbol found at address %xu\n", addr); + + addr = *(u64*)rsp; + continue; + } + + u8 buffer[4096U]; + read_memory((void*)sym->name, buffer, 4096U); + // print("Buffer is %s!%s\n", module->name, buffer); + print("%s!%s\n", module->name, sym->name); + // print("rsp value is %xu\n", rsp); + // print("rbp value is %xu\n", rbp); + + if (!read_memory((const void*)((u64)module->base_addr + (u64)sym->unwind_offset), (void*)u_info, 4096U)) { + print("Cannot read unwind info in for function %s\n", sym->name); + continue; + } + + 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 (sym->start <= addr && addr < u_info->size_of_prolog) { + // 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 + sym->start + u_info->size_of_prolog; + } else { + // print("RIP is in body\n"); + } + + for (u32 code_index = 0U; 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: + // print("push non volatile\n"); + // print("push "); + // print_reg(cur_code->op_info); + rsp += 8U; + code_index++; + break; + case UWOP_ALLOC_LARGE: + // print("alloc large\n"); + if (cur_code->op_info == 0U) { + u16 size = *(u16*)&u_info->unwind_code[code_index + 1U]; + if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + rbp += size; + } else { + // print("sub rsp %xu\n", size); + rsp += size; + } + code_index += 2U; + } else { + u32 size = *(u32*)&u_info->unwind_code[code_index + 1U]; + if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + rbp += size; + } else { + // print("sub rsp %xu\n", size); + rsp += size; + } + code_index += 3U; + } + break; + case UWOP_ALLOC_SMALL: + // print("alloc small\n"); + u32 size = cur_code->op_info * 8U + 8U; + if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + rbp += size; + } else { + // print("sub rsp %xu\n", size); + rsp += size; + } + code_index++; + break; + case UWOP_SET_FPREG: + // print("set frame point reg\n"); + // print("mov rbp rsp - %xu\n", u_info->frame_register_offset * 16U); + code_index++; + break; + case UWOP_SAVE_NONVOL: + // print("save non volatile\n"); + if (u_info->frame_register) { + // print("mov rbp "); + } else { + // print("mov rsp "); + } + // print_reg(cur_code->op_info); + code_index += 2U; + break; + case UWOP_SAVE_NONVOL_FAR: + // print("save non volatile far\n"); + code_index += 3U; + break; + case UWOP_SAVE_XMM128: + // print("save xmm128\n"); + code_index += 2U; + break; + case UWOP_SAVE_XMM128_FAR: + // print("save xmm128 far\n"); + code_index += 3U; + break; + case UWOP_PUSH_MACHFRAME: + // print("push mach frame\n"); + code_index++; + break; + default: + // print("bad operation code\n"); + code_index++; + } + } + u64 val; + read_memory((void*)rsp, &val, sizeof(u64)); + rsp += 8U; + // print("called address is %xu\n", val); + addr = val; + } + VirtualFree(u_info, 4096U, MEM_RELEASE); +} diff --git a/src/test.c b/src/test.c index 48c0168..b53fc55 100644 --- a/src/test.c +++ b/src/test.c @@ -1,5 +1,6 @@ #include int mainCRTStartup() { + __debugbreak(); unsigned long long int i = 0; while (1) { -- cgit v1.3 From 70208630ba98a9f2531a86830e514b1030f4fcde Mon Sep 17 00:00:00 2001 From: ladevieq Date: Mon, 3 Feb 2025 01:00:48 +0100 Subject: Start refactoring symbols management --- ] | 420 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ build.bat | 3 +- src/events.c | 112 ++++++++++++--- src/main.c | 2 +- src/module.c | 45 ++++++ src/parser.c | 6 + src/pdb.c | 19 ++- src/stack_walk.c | 45 +++--- 8 files changed, 611 insertions(+), 41 deletions(-) create mode 100644 ] diff --git a/] b/] new file mode 100644 index 0000000..1254912 --- /dev/null +++ b/] @@ -0,0 +1,420 @@ +const char* dummy_name = "Function_dummy_name"; + +void read_name_uni(const void* name_addr, const char* name, size_t name_size) { + u8* addr = NULL; + if (!read_memory(name_addr, (void*)&addr, 8U)) { + print("%s %u\n", "Loading DLL : error reading dll name : ", GetLastError()); + return; + } + + wchar_t name_uni[256U] = { 0U }; + if (!read_memory((const void*)addr, name_uni, sizeof(name_uni))) { + print("Cannot read dll name from proc memory %u\n", GetLastError()); + return; + } + + WideCharToMultiByte( + CP_ACP, + 0, + name_uni, + -1, + (char*)name, + (int)name_size, + NULL, + NULL + ); +} +void read_name_ansi(const void* name_addr, const char* name, size_t name_size) { + u8* addr = NULL; + if (!read_memory(name_addr, (void*)&addr, 8U)) { + print("%s\n", "Loading DLL : error reading dll name"); + return; + } + + if (!read_memory((const void*)addr, (char*)name, name_size)) { + print("Cannot read dll name from proc memory\n"); + return; + } +} + +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); + 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); + return; + } + + print("Module Base Address %xu\n", base_addr); + print("Module Size %xu\n", module->nt_header.OptionalHeader.SizeOfImage); + + // 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); + +#ifndef HASHMAP + u8** public_names = NULL; + u32* public_addresses = NULL; +#endif + if (export_table.Size != 0U) { + if (!read_memory((const void*)((u64)base_addr + (u64)export_table.VirtualAddress), (void*)&export_dir, sizeof(export_dir))) { + print("Cannot read export dir for module %s at address %x\n", (char*)module->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); + return; + } + + size_t names_size = export_dir.NumberOfNames * sizeof(u32); + u32* names_address = VirtualAlloc(NULL, names_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); +#ifndef HASHMAP + external_names = VirtualAlloc(NULL, export_dir.NumberOfNames * sizeof(external_names[0U]), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); +#endif + 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; + } + + size_t addresses_size = export_dir.NumberOfFunctions * sizeof(u32); + u32* addresses = VirtualAlloc(NULL, addresses_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); +#ifndef HASHMAP + external_addresses = addresses; +#endif + 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); + return; + } + +#ifndef HASHMAP + for (u32 i = 0U; i < export_dir.NumberOfFunctions; i++) { + for (u32 name_index = 0U; 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); + size_t length = 0U; + void* name_address = (void*)((u64)module->base_addr + names_address[name_index]); + read_memory(name_address, string_stack_current, 256U); + StringCchLengthA((PCNZCH)string_stack_current, STRSAFE_MAX_CCH, &length); + string_stack_current[length] = '\0'; + external_names[i] = string_stack_current; + string_stack_current += length + 1U; + break; + } + } + } + + init_module_functions(module, external_names, external_addresses, NULL, NULL, export_dir.NumberOfFunctions); +#endif + + print("Found %u functions\n", export_dir.NumberOfFunctions); + for (u32 i = 0U; i < export_dir.NumberOfFunctions; i++) { + char func_name[256U] = { 0 }; + u64 func_address = (u64)base_addr + addresses[i]; + u32 name_index = 0U; + for (; name_index < export_dir.NumberOfNames; name_index++) { + if (ordinals[name_index] == i) { + break; + } + } + + 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) { + print("%u : Fowarded function : %s\n", i, func_name); + } else if (name_index && name_index < export_dir.NumberOfNames) { + print("%u : Function name : %s, function address %xu\n", i, func_name, addresses[i]); +#ifdef HASHMAP + insert_elem(&module->symbols, addresses[i], 0U, name_address); +#endif + } else { + print("%u : Function has no name, function address %xu\n", i, addresses[i]); +#ifdef HASHMAP + insert_elem(&module->symbols, addresses[i], 0U, (u64)dummy_name); +#endif + } + } + + VirtualFree(names_address, names_size, MEM_RELEASE); + VirtualFree(ordinals, ordinals_size, MEM_RELEASE); +#ifdef HASHMAP + VirtualFree(addresses, addresses_size, MEM_RELEASE); +#endif + } + + // Private symbols + IMAGE_DATA_DIRECTORY debug_table = module->nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG]; + IMAGE_DEBUG_DIRECTORY debug_entries[20U] = { 0U }; + + u8** private_names = NULL; + u32* private_starts = NULL; + u32* private_ends = NULL; + 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); + return; + } + + 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 { + 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); + } + + print("PDB path : %s\n", &header.path); + open_pdb(&header.path, module); + } + + + // 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; + } + + // TODO: Merge private and public symbols + + u32 pdata_entries = exception_entry.Size / sizeof(exception_table[0U]); + u32* unwind_offsets = VirtualAlloc(NULL, pdata_entries * sizeof(unwind_offsets[0U]), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); + // Find unwind info offset + for (u32 index = 0U; index < pdata_entries; index++) { + RUNTIME_FUNCTION cur_entry = exception_table[index]; + +#ifdef HASHMAP + struct element* elem = find_elem(&module->symbols, cur_entry.BeginAddress); + if (elem == NULL) { + elem = insert_elem(&module->symbols, cur_entry.BeginAddress, cur_entry.EndAddress, 0ULL); + } else { + if (elem->name == 0U) { + print("function name %s\n", elem->name); + } + + elem->end = cur_entry.EndAddress; + } + + elem->unwind_offset = cur_entry.UnwindData; +#else + if (module->functions_start == cur_entry.BeginAddress) { + pdata_entries[i] = cur_entry.UnwindData; + } + // u32 sym_index = find_function(module, cur_entry.BeginAddress); + // if (sym_index == (u32)-1) { + // // TODO: Add element + // add_functions(module, NULL, (u32*)&cur_entry.BeginAddress, (u32*)&cur_entry.EndAddress, (u32*)&cur_entry.UnwindData, 1U); + // assert(FALSE); + // } else { + // module->functions_unwind[sym_index] = cur_entry.UnwindData; + // } +#endif + + // if (!read_memory((const void*)((u64)base_addr + (u64)cur_entry.UnwindData), (void*)u_info, 4096U)) { + // print("Cannot read unwind info in for function %s\n", elem->name); + // continue; + // } + + // size_t total_size = sizeof(*u_info) + (u_info->count_of_unwind_codes - 1U) * sizeof(struct UNWIND_CODE); + // assert(total_size < 4096U); + + // for (u32 code_index = 0U; code_index < u_info->count_of_unwind_codes; code_index++) { + // struct UNWIND_CODE* cur_code = &u_info->unwind_code[code_index]; + // switch (cur_code->unwind_operation_code) { + // case UWOP_PUSH_NONVOL: + // print("push non volatile\n"); + // print("push "); + // print_reg(cur_code->op_info); + // break; + // case UWOP_ALLOC_LARGE: + // print("alloc large\n"); + // if (cur_code->op_info == 0U) { + // u16 size = *(u16*)&u_info->unwind_code[code_index + 1U]; + // if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + // } else { + // print("sub rsp %xu\n", size); + // } + // code_index++; + // } else { + // u32 size = *(u32*)&u_info->unwind_code[code_index + 1U]; + // if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + // } else { + // print("sub rsp %xu\n", size); + // } + // code_index += 2U; + // } + // break; + // case UWOP_ALLOC_SMALL: + // print("alloc small\n"); + // u32 size = cur_code->op_info * 8U + 8U; + // if (u_info->frame_register) { + // print("sub rbp %xu\n", size); + // } else { + // print("sub rsp %xu\n", size); + // } + // break; + // case UWOP_SET_FPREG: + // print("set frame point reg\n"); + // print("mov rbp rsp - %xu\n", u_info->frame_register_offset * 16U); + // break; + // case UWOP_SAVE_NONVOL: + // print("save non volatile\n"); + // if (u_info->frame_register) { + // print("mov rbp "); + // } else { + // print("mov rsp "); + // } + // print_reg(cur_code->op_info); + // code_index++; + // break; + // case UWOP_SAVE_NONVOL_FAR: + // print("save non volatile far\n"); + // code_index += 2U; + // break; + // case UWOP_SAVE_XMM128: + // print("save xmm128\n"); + // code_index++; + // break; + // case UWOP_SAVE_XMM128_FAR: + // print("save xmm128 far\n"); + // code_index += 2U; + // break; + // case UWOP_PUSH_MACHFRAME: + // print("push mach frame\n"); + // break; + // default: + // print("bad operation code\n"); + // } + // } + } + + // VirtualFree(u_info, 4096U, MEM_RELEASE); + VirtualFree(exception_table, exception_entry.Size, MEM_RELEASE); +} + +void handle_create_process() { + print("%s\n", "CreateProcessDebugEvent"); + + CREATE_PROCESS_DEBUG_INFO proc_dbg_info = dbg_event.u.CreateProcessInfo; + + HANDLE hproc = OpenProcess(PROCESS_VM_READ, FALSE, dbg_event.dwProcessId); + hproc; + + struct module* module = &modules[add_module()]; + // https://learn.microsoft.com/en-us/windows/win32/api/minwinbase/ns-minwinbase-create_process_debug_info + // Debuggers must be prepared to handle the case where lpImageName is NULL or *lpImageName (in the address space of the process being debugged) is NULL. + // Specifically, the system does not provide an image name for a create process event + // 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, (LPSTR)module->name, sizeof(module->name), FILE_NAME_NORMALIZED); + print("%s%s\n", "Loading Process : ", (char*)module->name); + + load_symbols((const void*)proc_dbg_info.lpBaseOfImage, module); +} + +void handle_create_thread() { + print("%s\n", "CreateThreadDebugEvent"); +} + +void handle_exception() { + print("%s : ", "ExceptionDebugEvent"); + + if (dbg_event.u.Exception.ExceptionRecord.ExceptionCode == EXCEPTION_SINGLE_STEP && expect_single_step) { + print("%s\n", "Single Step"); + continue_status = DBG_CONTINUE; + expect_single_step = FALSE; + } else if (dbg_event.u.Exception.ExceptionRecord.ExceptionCode == EXCEPTION_BREAKPOINT) { + print("%s\n", "Breakpoint"); + continue_status = DBG_CONTINUE; + } else { + print("%u\n", dbg_event.u.Exception.ExceptionRecord.ExceptionCode); + } +} + +void handle_exit_process() { + print("%s\n", "ExitProcessDebugEvent"); + TerminateProcess(NULL, 0U); +} + +void handle_exit_thread() { + print("%s\n", "ExitThreadDebugEvent"); +} + +void handle_load_dll() { + LOAD_DLL_DEBUG_INFO dll_info = dbg_event.u.LoadDll; + const void* dll_name_addr = dll_info.lpImageName; + void* dll_base_addr = dll_info.lpBaseOfDll; + if (dll_name_addr == NULL) { + print("%s\n", "Loading DLL : error reading dll name"); + return; + } + + struct module* module = &modules[add_module()]; + module->base_addr = dll_base_addr; + read_name_uni(dll_name_addr, (const char*)module->name, sizeof(module->name)); + + print("%s%s\n", "Loading DLL : ", (char*)module->name); + + load_symbols(dll_base_addr, module); +} + +void handle_unload_dll() { + print("%s\n", "UnloadDLLDebugEvent"); +} + +void handle_output_debug_string() { + OUTPUT_DEBUG_STRING_INFO dbg_str_event = dbg_event.u.DebugString; + char* str_addr = dbg_str_event.lpDebugStringData; + size_t str_size = (size_t)dbg_str_event.nDebugStringLength; + // assert(str_size > 1024U); + u8 dbg_str[1024U] = { 0U }; + if (!read_memory(str_addr, dbg_str, str_size)) { + print("%s %u\n", "OutputDebugStringDebugEvent : error reading debug string : ", GetLastError()); + return; + } + if (dbg_str_event.fUnicode) { + u8 ansi_str[1024U] = { 0U }; + WideCharToMultiByte( + CP_ACP, + 0, + (wchar_t*)dbg_str, + -1, + (char*)ansi_str, + sizeof(ansi_str), + NULL, + NULL + ); + print("%s%s\n", "OutputDebugStringDebugEvent : ", ansi_str); + } else { + print("%s%s\n", "OutputDebugStringDebugEvent : ", dbg_str); + } +} + +void handle_rip() { + print("%s\n", "RIPEvent"); +} diff --git a/build.bat b/build.bat index 57d60ee..1d48f79 100644 --- a/build.bat +++ b/build.bat @@ -25,7 +25,8 @@ if EXIST %VS2022_PRO% ( :compilation -set C_FLAGS=/nologo /W4 /WX /Zi /GS- /GR- /Gs1000000000 /Fo:%BUILD_DIR%\ /Iinclude /Isrc /std:c17 /c /Tc +set C_FLAGS=/DHASHMAP /nologo /W4 /WX /Zi /GS- /GR- /Gs1000000000 /Fo:%BUILD_DIR%\ /Iinclude /Isrc /std:c17 /c /Tc +REM set C_FLAGS=/nologo /W4 /WX /Zi /GS- /GR- /Gs1000000000 /Fo:%BUILD_DIR%\ /Iinclude /Isrc /std:c17 /c /Tc set L_FLAGS=/WX /SUBSYSTEM:CONSOLE /NODEFAULTLIB /stack:0x100000,100000 if "%release%"=="0" ( 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); diff --git a/src/main.c b/src/main.c index 32e4fa8..ea94bc9 100644 --- a/src/main.c +++ b/src/main.c @@ -88,7 +88,7 @@ void dbg_loop() { int mainCRTStartup() { - string_stack_current = string_stack_base = VirtualAlloc(NULL, 4096U * 16U, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + string_stack_current = string_stack_base = VirtualAlloc(NULL, 4096U * 64U, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); STARTUPINFOA startup_info = { 0 }; startup_info.cb = sizeof(startup_info); diff --git a/src/module.c b/src/module.c index ffff0ca..9148aa8 100644 --- a/src/module.c +++ b/src/module.c @@ -1,8 +1,21 @@ struct module { void* base_addr; + u8** functions_name; + u32* functions_start; + u32* functions_end; + u32* functions_unwind; + u32 functions_count; + u32 functions_capacity; + // u32 functions_name_capacity; + // u32 function_start_capacity; + // u32 function_end_capacity; + // u32 function_unwind_capacity; + IMAGE_DOS_HEADER dos_header; IMAGE_NT_HEADERS64 nt_header; +#ifdef HASHMAP struct hashmap symbols; +#endif u8 name[256U]; }; @@ -18,7 +31,9 @@ u8 add_module() { _BitScanForward64((unsigned long*)&module_index, modules_list); modules_list ^= (u64)(0x1 << module_index); +#ifdef HASHMAP modules[module_index].symbols = create_unordered_flat(); +#endif return (u8)module_index; }; @@ -26,3 +41,33 @@ u8 add_module() { void remove_module(u8 index) { modules_list |= (u64)(0x1 << (((i8)index) - 1U)); } + +#ifdef HASHMAP +struct element* find_function(struct module* module, u32 rip_offset) { + for (u32 sym_index = 0U; sym_index < module->symbols.capacity; sym_index++) { + struct element* cur_sym = &module->symbols.arr[sym_index]; + if (cur_sym->start != 0U && cur_sym->start <= rip_offset && rip_offset < cur_sym->end) { + return cur_sym; + } + } + return NULL; +} +#else +u32 find_function(struct module* module, u64 rip_offset) { + if (module->functions_count == 0U) { + return (u32)-1; + } + u32 min = 0U, max = module->functions_count - 1U; + + while(min != max) { + u32 mid = (min + max) / 2U; + if (module->functions_start[mid] < rip_offset) { + min = mid + 1U; + } else { + max = mid - 1U; + } + } + + return min; +} +#endif diff --git a/src/parser.c b/src/parser.c index 937d640..ab9bde1 100644 --- a/src/parser.c +++ b/src/parser.c @@ -269,7 +269,13 @@ void run(struct ast_node* root) { for (u8 index = 0U; index < 64U; index++) { if (((1ULL << index) & modules_list) == 0U) { struct module* module = &modules[index]; +#ifdef HASHMAP iterate(&module->symbols); +#else + for (u32 sym_index = 0U; sym_index < module->functions_count; sym_index++) { + print("%s\n", module->functions_name[index]); + } +#endif } } } else { diff --git a/src/pdb.c b/src/pdb.c index 3f89611..7f37c34 100644 --- a/src/pdb.c +++ b/src/pdb.c @@ -129,6 +129,7 @@ void open_pdb(const char* pdb_path, struct module* module) { ReadFile(pdb_file, (void*)&DBI_block, sizeof(DBI_block), (LPDWORD)&b_read, NULL); struct DBI_stream_header* DBI_stream_header = (struct DBI_stream_header*)DBI_block; + // Read Module Info Substreams i32 byte_index = sizeof(*DBI_stream_header); for (u32 module_index = 0U; byte_index < DBI_stream_header->mod_info_size; module_index++) { modules_info[module_index] = (struct mod_info*)&DBI_block[byte_index]; @@ -147,14 +148,14 @@ void open_pdb(const char* pdb_path, struct module* module) { module_info_size = ((module_info_size + 3U) / 4U) * 4U; byte_index += module_info_size; - // Module Info Stream + // Read curren Module Info Stream u8 module_info_block[4096U] = { 0U }; SetFilePointer(pdb_file, stream_blocks_indices[modules_info[module_index]->ModuleSymStream][0U] * sizeof(module_info_block), NULL, FILE_BEGIN); ReadFile(pdb_file, (void*)&module_info_block, sizeof(module_info_block), (LPDWORD)&b_read, NULL); u8* symbols = &module_info_block[4U]; u8* cur_sym = symbols; - while ((u32)(cur_sym - symbols) < modules_info[module_index]->SymByteSize) { + while ((u32)(cur_sym - symbols) < modules_info[module_index]->SymByteSize - 4U) { struct SYMTYPE* rec = (struct SYMTYPE*)cur_sym; if (rec->rectyp == S_END) { print("END SYMBOL\n"); @@ -180,7 +181,17 @@ void open_pdb(const char* pdb_path, struct module* module) { StringCchLengthA((PCNZCH)proc->name, STRSAFE_MAX_CCH, &length); StringCchCopyA((LPSTR)string_stack_current, STRSAFE_MAX_CCH, (LPCSTR)proc->name); string_stack_current[length] = '\0'; - insert_elem(&module->symbols, proc->off + module->nt_header.OptionalHeader.BaseOfCode, proc->len + module->nt_header.OptionalHeader.BaseOfCode, (u64)string_stack_current); + u32 start_address = proc->off + module->nt_header.OptionalHeader.BaseOfCode; +#ifdef HASHMAP + insert_elem(&module->symbols, start_address, proc->len + module->nt_header.OptionalHeader.BaseOfCode, (u64)string_stack_current); +#else + for (u32 index = 0U; index < module->functions_count; index++) { + if (module->functions_start[index] == start_address) { + module->functions_name[index] = string_stack_current; + break; + } + } +#endif string_stack_current += length + 1U; } else if (rec->rectyp == S_REGREL32) { struct REGREL32* regrel = (struct REGREL32*)cur_sym; @@ -220,4 +231,6 @@ void open_pdb(const char* pdb_path, struct module* module) { // Public Symbol Stream SetFilePointer(pdb_file, stream_blocks_indices[DBI_stream_header->public_stream_index][0U] * sizeof(block), NULL, FILE_BEGIN); ReadFile(pdb_file, (void*)&block, sizeof(block), (LPDWORD)&b_read, NULL); + + CloseHandle(pdb_file); } diff --git a/src/stack_walk.c b/src/stack_walk.c index 94a76e1..bb4f4be 100644 --- a/src/stack_walk.c +++ b/src/stack_walk.c @@ -130,16 +130,6 @@ void print_reg(enum UNWIND_OP_INFO_TYPE op_info) { } } -struct element* find_function(struct module* module, u64 rip_offset) { - for (u32 sym_index = 0U; sym_index < module->symbols.capacity; sym_index++) { - struct element* cur_sym = &module->symbols.arr[sym_index]; - if (cur_sym->start != 0U && cur_sym->start <= rip_offset && rip_offset < cur_sym->end) { - return cur_sym; - } - } - return NULL; -} - void walk_stack(CONTEXT* ctx) { u64 addr = ctx->Rip; u64 rsp = ctx->Rsp; @@ -148,6 +138,11 @@ void walk_stack(CONTEXT* ctx) { struct UNWIND_INFO* u_info = VirtualAlloc(NULL, 4096U, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE); while (addr) { module = get_addr_module(addr); + + if (module == NULL) { + print("Module not found!\n"); + return; + } // print("Module name %s\n", module->name); // print("Module address %xu\n", (u64)module->base_addr); // print("Module size %xu\n", module->nt_header.OptionalHeader.SizeOfImage); @@ -156,23 +151,37 @@ void walk_stack(CONTEXT* ctx) { // print("Function address %xu\n", addr); // print("Function offset in module %xu\n", addr - (u64)module->base_addr); - struct element* sym = find_function(module, addr - (u64)module->base_addr); +#ifdef HASHMAP + struct element* sym = find_function(module, (u32)(addr - (u64)module->base_addr)); if (sym == NULL) { +#else + u32 sym_index = find_function(module, addr - (u64)module->base_addr); + if (sym_index == (u32)-1) { +#endif print("No symbol found at address %xu\n", addr); addr = *(u64*)rsp; continue; } +#ifdef HASHMAP + void* name = (void*)sym->name; + u64 unwind_offset = sym->unwind_offset; + u32 start = sym->start; u8 buffer[4096U]; - read_memory((void*)sym->name, buffer, 4096U); - // print("Buffer is %s!%s\n", module->name, buffer); - print("%s!%s\n", module->name, sym->name); + read_memory(name, buffer, 4096U); + print("Buffer is %s!%s\n", module->name, buffer); +#else + u8* name = module->functions_name[sym_index]; + u64 unwind_offset = module->functions_unwind[sym_index]; + u32 start = module->functions_start[sym_index]; +#endif + print("%s!%s\n", module->name, name); // print("rsp value is %xu\n", rsp); // print("rbp value is %xu\n", rbp); - if (!read_memory((const void*)((u64)module->base_addr + (u64)sym->unwind_offset), (void*)u_info, 4096U)) { - print("Cannot read unwind info in for function %s\n", sym->name); + if (!read_memory((const void*)((u64)module->base_addr + unwind_offset), (void*)u_info, 4096U)) { + print("Cannot read unwind info in for function %s\n", name); continue; } @@ -180,12 +189,12 @@ void walk_stack(CONTEXT* ctx) { assert(total_size < 4096U); // TODO: Do this once for the first function - if (sym->start <= addr && addr < u_info->size_of_prolog) { + if (start <= addr && addr < u_info->size_of_prolog) { // 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 + sym->start + u_info->size_of_prolog; + addr = (u64)module->base_addr + start + u_info->size_of_prolog; } else { // print("RIP is in body\n"); } -- cgit v1.3