diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/events.c | 112 | ||||
| -rw-r--r-- | src/main.c | 2 | ||||
| -rw-r--r-- | src/module.c | 45 | ||||
| -rw-r--r-- | src/parser.c | 6 | ||||
| -rw-r--r-- | src/pdb.c | 19 | ||||
| -rw-r--r-- | src/stack_walk.c | 45 |
6 files changed, 189 insertions, 40 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); @@ -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 { @@ -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"); } |