Store loaded modules info

Add command to list loaded modules
Add command to print callstack
Load modules of pdb files

5 files changed, 219 insertions, 79 deletions

diff --git a/src/events.c b/src/events.c
index 8870571..e57587c 100644
--- a/src/events.c
+++ b/src/events.c
@@ -1,13 +1,29 @@
+struct module modules[64U] = { 0U };
+u64 modules_list = ~0ULL;
+
+u8 add_module() {
+    assert(!modules_list);
+    u32 module_index = 0U;
+    _BitScanForward64((unsigned long*)&module_index, modules_list);
+    modules_list ^= (u64)(0x1 << module_index);
+
+    return (u8)module_index;
+};
+
+void remove_module(u8 module_index) {
+    modules_list |= (u64)(0x1 << (((i8)module_index) - 1U));
+}
+
 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\n", "Loading DLL : error reading dll name");
+        print("%s %u\n", "Loading DLL : error reading dll name : ", GetLastError());
         return;
     }
 
-    wchar_t name_uni[1024U] = { 0U };
+    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\n");
+        print("Cannot read dll name from proc memory %u\n", GetLastError());
         return;
     }
 
@@ -33,42 +49,43 @@ void read_name_ansi(const void* name_addr, const char* name, size_t name_size) {
         print("Cannot read dll name from proc memory\n");
         return;
     }
-
 }
-void load_symbols(const void* base_addr, const char* name) {
-    IMAGE_DOS_HEADER dos_header = { 0U };
-    if (!read_memory(base_addr, (void*)&dos_header, sizeof(dos_header))) {
-        print("Cannot read DOS header for module %s at address %x\n", (char*)name, base_addr);
+
+void load_symbols(const void* base_addr, const 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;
     }
 
-    IMAGE_NT_HEADERS64 nt_header = { 0U };
-    if (!read_memory((const void*)((u64)base_addr + dos_header.e_lfanew), (void*)&nt_header, sizeof(nt_header))) {
-        print("Cannot read NT header for module %s at address %x\n", (char*)name, base_addr);
+    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 = nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
+    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))) {
-            print("Cannot read export dir for module %s at address %x\n", (char*)name, base_addr);
+            print("Cannot read export dir for module %s at address %x\n", (char*)module->name, base_addr);
             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]))) {
-            print("Cannot read name ordinals from export directory for module %s at address %x\n", (char*)name, base_addr);
+            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]))) {
-            print("Cannot read names from export directory for module %s at address %x\n", (char*)name, base_addr);
+            print("Cannot read names from export directory for module %s at address %x\n", (char*)module->name, base_addr);
             return;
         }
 
@@ -85,11 +102,11 @@ void load_symbols(const void* base_addr, const char* name) {
 
 
     // Private symbols
-    IMAGE_DATA_DIRECTORY debug_table = nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG];
+    IMAGE_DATA_DIRECTORY debug_table = module->nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG];
     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", name, base_addr);
+        print("Cannot read image debug infos for module %s at address %x\n", module->name, base_addr);
         return;
     }
 
@@ -108,7 +125,7 @@ void load_symbols(const void* base_addr, const char* name) {
 
         struct codeview_header 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*)name, base_addr);
+            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);
@@ -120,11 +137,21 @@ void handle_create_process() {
     print("%s\n", "CreateProcessDebugEvent");
 
     CREATE_PROCESS_DEBUG_INFO proc_dbg_info = dbg_event.u.CreateProcessInfo;
-    u8 proc_name[1024U] = { 0U };
-    read_name_uni(proc_dbg_info.lpImageName, (const char*)proc_name, sizeof(proc_name));
-    print("%s%s\n", "Loading Process : ", (char*)proc_name);
 
-    load_symbols((const void*)proc_dbg_info.lpBaseOfImage, (const char*)proc_name);
+    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() {
@@ -158,18 +185,19 @@ void handle_exit_thread() {
 void handle_load_dll() {
     LOAD_DLL_DEBUG_INFO dll_info = dbg_event.u.LoadDll;
     const void* dll_name_addr = dll_info.lpImageName;
-    const void* dll_base_addr = dll_info.lpBaseOfDll;
+    void* dll_base_addr = dll_info.lpBaseOfDll;
     if (dll_name_addr == NULL) {
-        print("%s\n", "Loading DLL : error read dll name");
+        print("%s\n", "Loading DLL : error reading dll name");
         return;
     }
 
-    u8 dll_name[1024U] = { 0U };
-    read_name_uni(dll_name_addr, (const char*)dll_name, sizeof(dll_name));
+    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*)dll_name);
+    print("%s%s\n", "Loading DLL : ", (char*)module->name);
 
-    load_symbols(dll_base_addr, (const char*)dll_name);
+    load_symbols(dll_base_addr, module);
 }
 
 void handle_unload_dll() {
@@ -182,8 +210,8 @@ void handle_output_debug_string() {
     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\n", "OutputDebugStringDebugEvent : error read debug string");
+    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) {
diff --git a/src/main.c b/src/main.c
index 5ff1e42..b08c378 100644
--- a/src/main.c
+++ b/src/main.c
@@ -3,6 +3,7 @@
 #include <strsafe.h>
 
 #include "types.h"
+#include "module.h"
 
 #include "utils.c"
 
@@ -15,28 +16,6 @@ static DWORD continue_status = DBG_EXCEPTION_NOT_HANDLED;
 static DEBUG_EVENT dbg_event;
 static PROCESS_INFORMATION proc_info = { 0 };
 
-struct module {
-    u8 fix_compilation;
-};
-
-static struct module mods[64U] = { 0U };
-static u64 mods_list = 0U;
-
-u8 add_mod() {
-    u64 free_mods = ~mods_list;
-    assert(!free_mods);
-    u8 mod_index = (u8)(free_mods >> (_lzcnt_u64(free_mods) + 1U));
-    // struct module* mod = &mods[mod_index];
-
-
-
-    return mod_index;
-};
-
-void remove_mod(u8 mod_index) {
-    mods_list ^= (u64)(0x1 << (((i8)mod_index) - 1U));
-}
-
 /* Read debugged process memory at address */
 u8 read_memory(const void* addr, void* buffer, size_t size) {
     size_t b_read = 0U;
diff --git a/src/parser.c b/src/parser.c
index c42c2cc..db3a958 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -1,6 +1,8 @@
 enum TOKEN_TYPE {
     // Commands
     TOKEN_CONTINUE,
+    TOKEN_CALLSTACK,
+    TOKEN_LISTMODULE,
     TOKEN_PRINT_REG,
     TOKEN_STEP_INTO,
     TOKEN_DUMP_MEMORY,
@@ -32,6 +34,8 @@ struct ast_node {
 
 static struct command commands[] = {
     { .start = "g", .size = 1 },
+    { .start = "k", .size = 1 },
+    { .start = "lm", .size = 2 },
     { .start = "r", .size = 1 },
     { .start = "t", .size = 1 },
     { .start = "db", .size = 2 },
@@ -122,6 +126,20 @@ 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_CALLSTACK:
+            tokens++;
+            print("TOKEN_CALLSTACK\n");
+            cur_node->token = cur_token;
+            cur_node->left = NULL;
+            cur_node->right = NULL;
+            return cur_node;
+        case TOKEN_LISTMODULE:
+            tokens++;
+            print("TOKEN_LISTMODULE\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");
@@ -183,16 +201,34 @@ void run(struct ast_node* root) {
             continue_status
         );
         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("Function is in module name %s\n", module->name);
+                }
+            }
+        }
+    } else if (cur_node.token->type == TOKEN_LISTMODULE) {
+        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 Base : %xu\n", module->base_addr);
+                print("Module Size : %xu\n\n", module->nt_header.OptionalHeader.SizeOfImage);
+            }
+        }
     } else if (cur_node.token->type == TOKEN_PRINT_REG) {
-        print("RIP : %x\n", ctx.Rip);
-        print("RAX : %x\n", ctx.Rax);
-        print("RCX : %x\n", ctx.Rcx);
-        print("RDX : %x\n", ctx.Rdx);
-        print("RBX : %x\n", ctx.Rbx);
-        print("RSP : %x\n", ctx.Rsp);
-        print("RBP : %x\n", ctx.Rbp);
-        print("RSI : %x\n", ctx.Rsi);
-        print("RDI : %x\n", ctx.Rdi);
+        print("RIP : %xu\n", ctx.Rip);
+        print("RAX : %xu\n", ctx.Rax);
+        print("RCX : %xu\n", ctx.Rcx);
+        print("RDX : %xu\n", ctx.Rdx);
+        print("RBX : %xu\n", ctx.Rbx);
+        print("RSP : %xu\n", ctx.Rsp);
+        print("RBP : %xu\n", ctx.Rbp);
+        print("RSI : %xu\n", ctx.Rsi);
+        print("RDI : %xu\n", ctx.Rdi);
         process_commands = TRUE;
     } else if (cur_node.token->type == TOKEN_STEP_INTO) {
         ctx.EFlags |= 0x100; // TRAP_FLAG
diff --git a/src/pdb.c b/src/pdb.c
index dfa4c74..a28a201 100644
--- a/src/pdb.c
+++ b/src/pdb.c
@@ -1,8 +1,26 @@
+#include "cvinfo.h"
+
+enum STREAM_INDEX {
+    OLD_DIRECTORY,
+    PDB,
+    TPI,
+    DBI,
+    IPI,
+};
+
+enum SYMBOL_KIND {
+    OBJNAME = 0x1101,
+    PUBLIC = 0x110e,
+    GPROC32 = 0x1110,
+    COMPIL3 = 0x113C,
+};
+
 void open_pdb(const char* pdb_path) {
     HANDLE pdb_file = CreateFileA(pdb_path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
 
     if (pdb_file == INVALID_HANDLE_VALUE) {
         print("Cannot open pdb file at path %s due to error %u\n", pdb_path, GetLastError());
+        return;
     }
 
     // MSF superblock
@@ -24,24 +42,27 @@ void open_pdb(const char* pdb_path) {
     u32 dir_addr = (superblock.block_map_addr - 1U) * sizeof(block);
     SetFilePointer(pdb_file, dir_addr, NULL, FILE_BEGIN);
     ReadFile(pdb_file, (void*)block, superblock.num_directory_bytes, (LPDWORD)&b_read, NULL);
+
+    // Loog through streams
     u32 num_streams = ((u32*)block)[0U];
     u32* stream_sizes = &((u32*)block)[1U];
     u32* stream_blocks = stream_sizes + num_streams;
-    u32* pdb_info_stream_blocks = NULL;
-    u32* pdb_DBI_stream_blocks = NULL;
+    u32* stream_blocks_indices[1024U] = { 0U };
 
     u32 block_index = 0U;
     for(u32 stream_id = 0u; stream_id < num_streams; stream_id++) {
         print("PDB stream size : %u\n", stream_sizes[stream_id]);
-        if (stream_id == 1U) {
-            pdb_info_stream_blocks = &stream_blocks[block_index];
-        }
-        if (stream_id == 3U) {
-            pdb_DBI_stream_blocks = &stream_blocks[block_index];
-        }
+
         if (stream_sizes[stream_id] == (u32)-1) {
+            print("PDB stream blocks : NONE\n");
+            continue;
+        }
+
+        if (stream_sizes[stream_id] == 0U) {
             continue;
         }
+
+        stream_blocks_indices[stream_id] = &stream_blocks[block_index];
         for(u32 id = 0u; id < (stream_sizes[stream_id] / sizeof(block)) + 1U; id++) {
             print("PDB stream blocks : %u\n", stream_blocks[block_index]);
             block_index++;
@@ -57,12 +78,11 @@ void open_pdb(const char* pdb_path) {
     };
 
     struct pdb_info_stream pdb_info_stream = { 0U };
-    SetFilePointer(pdb_file, pdb_info_stream_blocks[0U] * sizeof(block), NULL, FILE_BEGIN);
+    SetFilePointer(pdb_file, stream_blocks_indices[PDB][0U] * sizeof(block), NULL, FILE_BEGIN);
     ReadFile(pdb_file, (void*)&pdb_info_stream, sizeof(pdb_info_stream), (LPDWORD)&b_read, NULL);
-    pdb_info_stream;
 
     // DBI Stream
-    struct pdb_DBI_stream {
+    struct DBI_stream_header {
         i32 version_signature;
         u32 version_header;
         u32 age;
@@ -84,13 +104,89 @@ void open_pdb(const char* pdb_path) {
         u16 machine;
         u32 padding;
     };
-    struct pdb_DBI_stream pdb_DBI_stream = { 0U };
-    SetFilePointer(pdb_file, pdb_DBI_stream_blocks[0U] * sizeof(block), NULL, FILE_BEGIN);
-    ReadFile(pdb_file, (void*)&pdb_DBI_stream, sizeof(pdb_DBI_stream), (LPDWORD)&b_read, NULL);
-    print("");
-    pdb_DBI_stream;
+
+    struct mod_info {
+        u32 Unused1;
+        struct SectionContribEntry {
+            u16 Section;
+            u8 Padding1[2];
+            i32 Offset;
+            i32 Size;
+            u32 Characteristics;
+            u16 ModuleIndex;
+            u8 Padding2[2];
+            u32 DataCrc;
+            u32 RelocCrc;
+        } SectionContr;
+        u16 Flags;
+        u16 ModuleSymStream;
+        u32 SymByteSize;
+        u32 C11ByteSize;
+        u32 C13ByteSize;
+        u16 SourceFileCount;
+        u8 Padding[2];
+        u32 Unused2;
+        u32 SourceFileNameIndex;
+        u32 PdbFilePathNameIndex;
+    } *modules_info[64U];
+    u8* modules_name[64U] = { 0U };
+    u8* obj_files_name[64U] = { 0U };
+    u8 DBI_block[4096U] = { 0U };
+    SetFilePointer(pdb_file, stream_blocks_indices[DBI][0U] * sizeof(block), NULL, FILE_BEGIN);
+    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;
+
+    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];
+        modules_name[module_index] = &DBI_block[byte_index + sizeof(struct mod_info)];
+        print("Module symbole stream index %xu\n", modules_info[module_index]->ModuleSymStream);
+        print("%s\n", modules_name[module_index]);
+        size_t module_name_length = 0U;
+        StringCchLengthA((STRSAFE_LPSTR)modules_name[module_index], STRSAFE_MAX_CCH, &module_name_length);
+        obj_files_name[module_index] = &DBI_block[byte_index + sizeof(struct mod_info) + module_name_length + 1U];
+
+        print("%s\n", obj_files_name[module_index]);
+        size_t obj_file_name_length = 0U;
+        StringCchLengthA((STRSAFE_LPSTR)obj_files_name[module_index], STRSAFE_MAX_CCH, &obj_file_name_length);
+
+        u32 module_info_size = sizeof(*modules_info[0U]) + (i32)module_name_length + 1 + (i32)obj_file_name_length + 1;
+        module_info_size = ((module_info_size + 3U) * 4U) / 4U;
+        byte_index += module_info_size;
+
+        // 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) {
+            struct SYMTYPE* rec = (struct SYMTYPE*)cur_sym;
+            if (rec->rectyp == OBJNAME) {
+                struct OBJNAMESYM* objname = (struct OBJNAMESYM*)cur_sym;
+                print("OBJNAME SYMBOL %s\n", objname->name);
+            } else if (rec->rectyp == PUBLIC) {
+                struct PUBSYM32* pub = (struct PUBSYM32*)cur_sym;
+                print("PUBLIC SYMBOL %s\n", pub->name);
+            } else if (rec->rectyp == GPROC32) {
+                struct PROCSYM32* proc = (struct PROCSYM32*)cur_sym;
+                print("PROC32 SYMBOL %s\n", proc->name);
+            } else if (rec->rectyp == COMPIL3) {
+                struct COMPILESYM3* com = (struct COMPILESYM3*)cur_sym;
+                print("COMPILER SYMBOL %s\n", com->verSz);
+            } else {
+                print("Unkown symbole kind %xu\n", rec->rectyp);
+            }
+            cur_sym += rec->reclen + 2U;
+        }
+    }
 
     // Global Symbol Stream
-    SetFilePointer(pdb_file, pdb_DBI_stream.global_stream_index * sizeof(block), NULL, FILE_BEGIN);
-    ReadFile(pdb_file, (void*)&pdb_DBI_stream, sizeof(pdb_DBI_stream), (LPDWORD)&b_read, NULL);
+    SetFilePointer(pdb_file, stream_blocks_indices[DBI_stream_header->global_stream_index][0U] * sizeof(block), NULL, FILE_BEGIN);
+    ReadFile(pdb_file, (void*)&block, sizeof(DBI_stream_header), (LPDWORD)&b_read, NULL);
+
+    // 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);
 }
diff --git a/src/test.c b/src/test.c
index 85c9136..48c0168 100644
--- a/src/test.c
+++ b/src/test.c
@@ -3,6 +3,7 @@ int mainCRTStartup() {
     unsigned long long int i = 0;
 
     while (1) {
+        __debugbreak();
         OutputDebugStringA("Hello World !");
         i++;
     }