Log dll name on loading

Print output debug string
Refactored code
Improved print
Add a function to read memory from the deubgged process
Fixed a read out of bounds

1 files changed, 72 insertions, 320 deletions

diff --git a/src/main.c b/src/main.c
index 83d4acd..94fdd37 100644
--- a/src/main.c
+++ b/src/main.c
@@ -1,4 +1,5 @@
 #include <Windows.h>
+#include <intrin.h>
 #include <strsafe.h>
 
 #include "types.h"
@@ -6,288 +7,101 @@
 #include "utils.c"
 
 static HANDLE h_input = NULL;
+static HANDLE h_output = NULL;
 static u8 is_open = TRUE;
 static u8 process_commands = TRUE;
 static u8 expect_single_step = FALSE;
 static DWORD continue_status = DBG_EXCEPTION_NOT_HANDLED;
 static DEBUG_EVENT dbg_event;
-static PROCESS_INFORMATION proc_info;
+static PROCESS_INFORMATION proc_info = { 0 };
 
-enum TOKEN_TYPE {
-    // Commands
-    TOKEN_CONTINUE,
-    TOKEN_PRINT_REG,
-    TOKEN_STEP_INTO,
-    TOKEN_DUMP_MEMORY,
-    TOKEN_QUIT,
-
-    // Primaries
-    TOKEN_DECIMAL_NUMBER,
-    TOKEN_HEXADECIMAL_NUMBER,
-
-    TOKEN_UNKNOWN,
+struct module {
+    u8 fix_compilation;
 };
 
-struct token {
-    enum TOKEN_TYPE type;
-    const char* lexeme;
-    u8 size;
-};
+static struct module mods[64U] = { 0U };
+static u64 mods_list = 0U;
 
-struct command {
-    const char* start;
-    u8 size;
-};
+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];
 
-struct ast_node {
-    struct token* token;
-    struct ast_node* left;
-    struct ast_node* right;
-};
 
-static struct command commands[] = {
-    { .start = "g", .size = 1 },
-    { .start = "r", .size = 1 },
-    { .start = "t", .size = 1 },
-    { .start = "db", .size = 2 },
-    { .start = "q", .size = 1 },
-};
 
-_Static_assert((sizeof(commands) / sizeof(commands[0])) == TOKEN_QUIT + 1, "Missing commands");
+    return mod_index;
+};
 
-u8 is_end(char c) {
-    return c == '\0';
+void remove_mod(u8 mod_index) {
+    mods_list ^= (u64)(0x1 << (((i8)mod_index) - 1U));
 }
 
-u8 is_alpha(char c) {
-    return ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z');
-}
+/* Read debugged process memory at address */
+u8 read_memory(const void* addr, void* buffer, size_t size) {
+    size_t b_read = 0U;
+    HRESULT hr = ReadProcessMemory(proc_info.hProcess, addr, buffer, size, &b_read);
 
-u8 is_decimal_digit(char c) {
-    return '0' <= c && c <= '9';
+    return hr > 0 && b_read == size;
 }
 
-u8 is_hexadecimal_digit(char c) {
-    return ('0' <= c && c <= '9') || ('a' <= c && c <= 'f') || ('A' <= c && c <= 'F');
-}
-
-void lexer(const char* str_cmd, size_t str_cmd_len, struct token* tokens) {
-    const char* cur = str_cmd;
-    struct token* cur_token = tokens;
-    str_cmd_len = 0;
-
-    while(*cur != '\0') {
-        cur_token->type = TOKEN_UNKNOWN;
-        if(*cur == '0' && cur[1] == 'x') {
-            const char* start = cur;
-            cur += 2;
+#include "parser.c"
+#include "events.c"
 
-            while(is_hexadecimal_digit(*cur) && !is_end(*cur)) {
-                cur++;
-            }
-
-            cur_token->lexeme = start;
-            cur_token->type = TOKEN_HEXADECIMAL_NUMBER;
-            cur_token->size = (u8)(cur - start);
-            cur_token++;
-        } else if(is_decimal_digit(*cur)) {
-            const char* start = cur;
-            cur++;
-
-            while(is_decimal_digit(*cur) && !is_end(*cur)) {
-                cur++;
-            }
-
-            cur_token->lexeme = start;
-            cur_token->type = TOKEN_DECIMAL_NUMBER;
-            cur_token->size = (u8)(cur - start);
-            cur_token++;
-        } else if(is_alpha(*cur)) {
-            const char* start = cur;
-            cur++;
-            while(is_alpha(*cur) && !is_end(*cur)) {
-                cur++;
-            }
+void dbg_loop() {
+    while(is_open) {
+        WaitForDebugEventEx(&dbg_event, INFINITE);
 
-            cur_token->lexeme = start;
-            cur_token->size = (u8)(cur - start);
+        continue_status = DBG_EXCEPTION_NOT_HANDLED;
 
-            for(int i = 0; i < sizeof(commands) / sizeof(commands[0]); i++) {
-                struct command cur_cmd = commands[i];
-                if (CompareStringA(LOCALE_SYSTEM_DEFAULT, 0, cur_cmd.start, cur_cmd.size, cur_token->lexeme, cur_token->size) == CSTR_EQUAL) {
-                    cur_token->type = i;
-                    cur_token++;
-                    break;
-                }
-            }
-        } else {
-            cur++;
+        switch (dbg_event.dwDebugEventCode) {
+            case CREATE_PROCESS_DEBUG_EVENT:
+                handle_create_process();
+                break;
+            case CREATE_THREAD_DEBUG_EVENT:
+                handle_create_thread();
+                break;
+            case EXCEPTION_DEBUG_EVENT:
+                handle_exception();
+                break;
+            case EXIT_PROCESS_DEBUG_EVENT:
+                handle_exit_process();
+                break;
+            case EXIT_THREAD_DEBUG_EVENT:
+                handle_exit_thread();
+                break;
+            case LOAD_DLL_DEBUG_EVENT:
+                handle_load_dll();
+                break;
+            case OUTPUT_DEBUG_STRING_EVENT:
+                handle_output_debug_string();
+                break;
+            case RIP_EVENT:
+                handle_rip();
+                break;
+            case UNLOAD_DLL_DEBUG_EVENT:
+                handle_unload_dll();
+                break;
         }
-    }
-}
-
-struct ast_node* parse_primary(struct token* tokens, struct ast_node* nodes) {
-    struct ast_node* cur_node = &nodes[0];
-    struct token* cur_token = &tokens[0];
-    if (cur_token->type == TOKEN_HEXADECIMAL_NUMBER ||
-        cur_token->type == TOKEN_DECIMAL_NUMBER) {
-        cur_node->token = cur_token;
-        cur_node->left = NULL;
-        cur_node->right = NULL;
-        return cur_node;
-    } else {
-        print("Expecting a primary token\n");
-        return NULL;
-    }
-}
-
-struct ast_node* parse_command(struct token* tokens, struct ast_node* nodes) {
-    struct ast_node* cur_node = &nodes[0];
-    struct token* cur_token = &tokens[0];
-    switch(cur_token->type) {
-        case TOKEN_CONTINUE:
-            tokens++;
-            print("TOKEN_CONTINUE\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");
-            cur_node->token = cur_token;
-            cur_node->left = NULL;
-            cur_node->right = NULL;
-            return cur_node;
-        case TOKEN_STEP_INTO:
-            tokens++;
-            print("TOKEN_STEP_INTO\n");
-            cur_node->token = cur_token;
-            cur_node->left = NULL;
-            cur_node->right = NULL;
-            return cur_node;
-        case TOKEN_DUMP_MEMORY:
-            tokens++;
-            print("TOKEN_DUMP_MEMORY\n");
-            cur_node->token = cur_token;
-            cur_node->left = NULL;
-            cur_node->right = parse_primary(tokens, ++nodes);
-            return cur_node;
-        case TOKEN_QUIT:
-            tokens++;
-            print("TOKEN_QUIT\n");
-            cur_node->token = cur_token;
-            cur_node->left = NULL;
-            cur_node->right = NULL;
-            return cur_node;
-        default:
-            print("%s must be a command\n", cur_token->lexeme);
-            return NULL;
-    }
-}
-
-u64 eval_number(const struct ast_node* node) {
-    if (node->token->type == TOKEN_DECIMAL_NUMBER ||
-        node->token->type == TOKEN_HEXADECIMAL_NUMBER) {
-        return atou64(node->token->lexeme, node->token->size);
-    }
-    return ~0U;
-}
 
-void run(struct ast_node* root) {
-    struct ast_node cur_node = root[0];
-    CONTEXT ctx = {
-        .ContextFlags = CONTEXT_ALL,
-    };
-    HANDLE h_thread = OpenThread(
-        THREAD_GET_CONTEXT | THREAD_SET_CONTEXT,
-        FALSE,
-        dbg_event.dwThreadId
-    );
-    GetThreadContext(h_thread, &ctx);
-
-    if (cur_node.token->type == TOKEN_CONTINUE) {
-        ContinueDebugEvent(
-            dbg_event.dwProcessId,
-            dbg_event.dwThreadId,
-            continue_status
-        );
-        process_commands = FALSE;
-    } 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);
         process_commands = TRUE;
-    } else if (cur_node.token->type == TOKEN_STEP_INTO) {
-        ctx.EFlags |= 0x100; // TRAP_FLAG
-        if(SetThreadContext(h_thread, &ctx) == 0) {
-            __debugbreak();
-        }
-        ContinueDebugEvent(
-            dbg_event.dwProcessId,
-            dbg_event.dwThreadId,
-            continue_status
-        );
-        expect_single_step = TRUE;
-        process_commands = FALSE;
-    } else if (cur_node.token->type == TOKEN_QUIT) {
-        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 };
-        size_t b_read = 0U;
-
-        HRESULT hr = ReadProcessMemory(
-            proc_info.hProcess,
-            (void*)addr,
-            buf,
-            sizeof(buf),
-            &b_read
-        );
+        while(process_commands) {
+            const char write_buffer[] = "> ";
+            DWORD c_written = 0U;
+            WriteConsole(h_output, write_buffer, sizeof(write_buffer), &c_written, NULL);
+            char input_buffer[64] = { 0 };
+            DWORD input_size = 0U;
+            ReadConsole(h_input, input_buffer, sizeof(input_buffer), &input_size, NULL);
 
-        if (hr > 0) {
-            print("%b %b %b %b %b %b %b %b %b %b %b %b %b %b %b %b\n", buf[0], buf[1], buf[2], buf[3], buf[4], buf[5], buf[6], buf[7], buf[8], buf[9], buf[10], buf[11], buf[12], buf[13], buf[14], buf[15]);
-        } else {
-            print("Unable to read memory at this address\n");
+            interpreter(input_buffer, input_size);
         }
-
-    } else {
-        print("Command unknown\n");
-        process_commands = TRUE;
-    }
-
-    CloseHandle(h_thread);
-}
-
-u8 interpreter(const char* str_cmd, size_t str_cmd_len) {
-    struct token tokens[64U];
-    struct ast_node nodes[64U];
-
-    lexer(str_cmd, str_cmd_len, tokens);
-
-    struct ast_node* root = parse_command(tokens, nodes);
-
-    if(root == NULL) {
-        print("Ill formed command\n");
-        return FALSE;
     }
-
-    run(root);
-    return TRUE;
 }
 
 int mainCRTStartup() {
 
-    STARTUPINFOW startup_info = {
-        .cb = sizeof(startup_info),
-    };
+    STARTUPINFOW startup_info = { 0 };
+    startup_info.cb = sizeof(startup_info);
 
     BOOL ret = CreateProcessW(
         NULL,
@@ -307,6 +121,12 @@ int mainCRTStartup() {
         return 0;
     }
 
+    h_output = GetStdHandle(STD_OUTPUT_HANDLE);
+    if (h_output == NULL) {
+        TerminateProcess(proc_info.hProcess, 1U);
+        return 1U;
+    }
+
     h_input = GetStdHandle(STD_INPUT_HANDLE);
     if (h_input == NULL) {
         TerminateProcess(proc_info.hProcess, 1U);
@@ -315,75 +135,7 @@ int mainCRTStartup() {
     DWORD console_mode = ENABLE_ECHO_INPUT | ENABLE_LINE_INPUT | ENABLE_PROCESSED_INPUT;
     SetConsoleMode(h_input, console_mode);
 
-    while(is_open) {
-        WaitForDebugEventEx(&dbg_event, INFINITE);
-
-        continue_status = DBG_EXCEPTION_NOT_HANDLED;
-
-        switch (dbg_event.dwDebugEventCode) {
-            case CREATE_PROCESS_DEBUG_EVENT: {
-                OpenThread(THREAD_GET_CONTEXT | THREAD_SET_CONTEXT, FALSE, dbg_event.dwThreadId);
-                print("%s\n", "CreateProcessDebugEvent");
-                break;
-            }
-            case CREATE_THREAD_DEBUG_EVENT: {
-                print("%s\n", "CreateThreadDebugEvent");
-                break;
-            }
-            case EXCEPTION_DEBUG_EVENT: {
-                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);
-                }
-
-                break;
-            }
-            case EXIT_PROCESS_DEBUG_EVENT: {
-                print("%s\n", "ExitProcessDebugEvent");
-                TerminateProcess(NULL, 0U);
-                return 0;
-                break;
-            }
-            case EXIT_THREAD_DEBUG_EVENT: {
-                print("%s\n", "ExitThreadDebugEvent");
-                break;
-            }
-            case LOAD_DLL_DEBUG_EVENT: {
-                // print("%s%s\n", "Loading DLL : ", (char*)dbg_event.u.LoadDll.lpImageName);
-                print("%s%s\n", "Loading DLL : ", NULL);
-                break;
-            }
-            case OUTPUT_DEBUG_STRING_EVENT: {
-                print("%s\n", "OutputDebugStringDebugEvent");
-                break;
-            }
-            case RIP_EVENT: {
-                print("%s\n", "RIPEvent");
-                break;
-            }
-            case UNLOAD_DLL_DEBUG_EVENT: {
-                print("%s\n", "UnloadDLLDebugEvent");
-                break;
-            }
-        }
-
-        process_commands = TRUE;
-        while(process_commands) {
-            char input_buffer[16];
-            DWORD input_size = 0U;
-            ReadConsole(h_input, input_buffer, sizeof(input_buffer), &input_size, NULL);
-
-            interpreter(input_buffer, input_size);
-        }
-    }
+    dbg_loop();
 
     TerminateProcess(proc_info.hProcess, 0U);
     return 0;