injector.cpp

#include <iostream>

#include "preprocessor.hpp"
#include "context.hpp"
#include "classdef.hpp"

#include <string>
#include <vector>
#include <set>
#include <unordered_map>
#include <regex>
#include <stack>
#include <queue>
#include <cassert>
#include <fstream>

static const unsigned char hde64_bin[] = {
    #embed "hde64.h"
};

const char * required_headers = R"(
#include <sys/mman.h>
#include <unistd.h>
#include <cstring>
#include <cstdint>
#include <iostream>
#include <cerrno>
#include <cassert>
)";

const char * str_allocate_nearby_memory = R"(
void* allocate_nearby_memory(void* target_addr, size_t size) {
    size_t page_size = sysconf(_SC_PAGESIZE);
    uintptr_t start_addr = (uintptr_t)target_addr;
    // Round down to page boundary
    start_addr &= ~(page_size - 1);

    // Search range: +/- 2GB, but let's be safe with 1GB
    // We search downwards first (often easier to find holes)
    for (size_t i = 0; i < 1024; i++) {
        uintptr_t try_addr = start_addr - (i * page_size);
        void* mem = mmap((void*)try_addr, size, PROT_READ | PROT_WRITE | PROT_EXEC,
                         MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED_NOREPLACE, -1, 0);
        if (mem != MAP_FAILED) return mem;
        
        // Try upwards
        try_addr = start_addr + (i * page_size);
        mem = mmap((void*)try_addr, size, PROT_READ | PROT_WRITE | PROT_EXEC,
                   MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED_NOREPLACE, -1, 0);
        if (mem != MAP_FAILED) return mem;
    }
    // Fallback (might fail relocation if too far)
    return mmap(NULL, size, PROT_READ | PROT_WRITE | PROT_EXEC, 
                MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
}
)";

const char * str_install_smart_hook = R"(
void install_smart_hook(void* original_func, void* hook_func, void** tramp_ptr_out) {
    uint8_t* src = (uint8_t*)original_func;
    uint8_t* dst;
    
    void* tramp_mem = allocate_nearby_memory(original_func, 4096);
    dst = (uint8_t*)tramp_mem;

    hde64s hs;
    size_t bytes_copied = 0;
    // We still only need to steal 12 bytes from the SOURCE to fit the hook there.
    // The trampoline TAIL can be larger (14 bytes) because we own that memory.
    const size_t HOOK_SIZE = 12; 

    while (bytes_copied < HOOK_SIZE) {
        unsigned int len = hde64_disasm(src, &hs);
        if (hs.flags & F_ERROR) return;

        std::memcpy(dst, src, len);

        // --- Relocation Logic (Keep this, it was correct) ---
        if ((hs.flags & F_MODRM) && (hs.modrm_mod == 0) && (hs.modrm_rm == 5)) {
            int32_t old_disp = hs.disp.disp32;
            uintptr_t next_rip_original = (uintptr_t)src + len;
            uintptr_t target_data_addr = next_rip_original + old_disp;
            uintptr_t next_rip_tramp = (uintptr_t)dst + len;
            int64_t new_disp = (int64_t)target_data_addr - (int64_t)next_rip_tramp;

            if (new_disp > 2147483647 || new_disp < -2147483648) {
                std::cerr << "[Hook Error] Relocation out of range!\n";
            } else {
                uint8_t* disp_ptr = dst + (len - 4 - (hs.flags & F_IMM32 ? 4 : 0) - (hs.flags & F_IMM16 ? 2 : 0) - (hs.flags & F_IMM8 ? 1 : 0));
                *(int32_t*)disp_ptr = (int32_t)new_disp;
            }
        }
        
        // Handle Relative Jumps (E8/E9)
        if (src[0] == 0xE8 || src[0] == 0xE9) {
             int32_t old_rel = *(int32_t*)(src + 1);
             uintptr_t target = (uintptr_t)src + len + old_rel;
             int64_t new_rel = (int64_t)target - ((int64_t)dst + len);
             *(int32_t*)(dst + 1) = (int32_t)new_rel;
        }

        src += len;
        dst += len;
        bytes_copied += len;
    }

    // --- STEP 3: REGISTER-SAFE JUMP BACK ---
    // We CANNOT use MOV RAX,... because it overwrites the return value of the getter!
    // We use JMP [RIP+0] followed by the address.
    // Opcode: FF 25 00 00 00 00 (JMP qword ptr [RIP+0])
    
    dst[0] = 0xFF; dst[1] = 0x25; dst[2] = 0x00; dst[3] = 0x00; dst[4] = 0x00; dst[5] = 0x00;
    *(uintptr_t*)(dst + 6) = (uintptr_t)src; // Jump back to where we stopped stealing
    
    *tramp_ptr_out = tramp_mem;

    // --- STEP 4: APPLY HOOK TO ORIGINAL ---
    // This part was fine. We can clobber RAX here because the function hasn't started yet.
    unprotect_memory(original_func, bytes_copied);
    
    uint8_t* hook_loc = (uint8_t*)original_func;
    
    hook_loc[0] = 0x48; hook_loc[1] = 0xB8;
    *(uintptr_t*)(hook_loc + 2) = (uintptr_t)hook_func;
    hook_loc[10] = 0xFF; hook_loc[11] = 0xE0;
    
    for (size_t i = 12; i < bytes_copied; i++) {
        hook_loc[i] = 0x90;
    }
})";



context_t extract_class_def_context(const std::string& class_name, const std::vector<std::string>& header_files) {
    std::string source_code = "#include <iostream>\n";
    for (const auto& header : header_files) {
        source_code += "#include <" + header + ">\n";
    }
    source_code += "\nint main() {\n";
    source_code += "    " + class_name + "* obj;\n";
    source_code += "    return 0;\n";
    source_code += "}\n";
    
    auto ctx = context_t(preprocess(source_code));
    return ctx;
}


std::string proxy_class_inner(const std::string& class_name, const context_t& context ) {
    std::string c = "";
    std::vector<std::string> parent_classes;

    // std::cerr << "Generating proxy for class: " << class_name << "\n";
    std::string class_abs_def = context.get_abs_def("global", class_name);
    assert(!class_abs_def.empty() && "Class definition not found in context");



    std::string class_scope = class_abs_def.substr(0, class_abs_def.find_last_of("::") - 1);

    // std::cerr << "Absolute class definition name: " << class_abs_def << "\n";
    std::string class_def = std::string(context.get_source_code_of(class_abs_def));
    // std::cerr << "Class Definition Extracted:\n" << class_def << "\n";

    { // extract parent classes
    size_t first_brace = class_def.find("{");
    size_t first_colon = class_def.find(":");

    if (first_colon != std::string::npos && first_colon < first_brace) {
        std::string parents_str = class_def.substr(first_colon + 1, first_brace - first_colon - 1);
        std::regex parent_regex(R"(\s*(?:public|protected|private)?\s*(?:virtual)?\s*([\w:<>]+))");        
        auto it_begin = std::sregex_iterator(parents_str.begin(), parents_str.end(), parent_regex);
        auto it_end = std::sregex_iterator();
        
        for (std::sregex_iterator it = it_begin; it != it_end; ++it) {
            std::string parent_class = it->str(1);
            if (parent_class.empty()) continue;
            
            std::string abs_scope = context.get_abs_scope(class_scope, parent_class);
            std::string simple_name = parent_class.substr(parent_class.find_last_of("::") + 1);
            parent_classes.push_back(abs_scope + "::" + simple_name);
        }

    }
    } // end extract parent classes


    class_def = strip_function_bodies(class_def);
    auto functions = analyze_class_def(class_def);


    std::string safe_cpp_class_name = class_name; 
    std::replace(safe_cpp_class_name.begin(), safe_cpp_class_name.end(), ':', '_');
    std::string class_namespace = class_name.substr(0, class_name.find_last_of("::") - 1);

    for (const auto& parent_class : parent_classes) {
        std::string proxy_res = proxy_class_inner(parent_class, context);
        c += proxy_res + "\n\n\n";
    }

    // normalize types
    for (auto& func : functions) {
        
        std::cerr << class_name << " ( " << AccessLevelToString(func.access_level) << " ) ";
        std::cerr << "func " << func.name << "(";
        for (auto& arg_type : func.arg_types) {
            std::cerr << arg_type << " ";
            arg_type = context.normalize_type(class_namespace, arg_type);
            std::cerr << "(" << arg_type << "), ";
        }

        std::cerr << ") -> " << func.return_type;
        func.return_type = context.normalize_type(class_namespace, func.return_type);
        std::cerr << " (" << func.return_type << ")\n";

    }


    // we should skip overloaded functions for simplicity
    std::map<std::string, int> func_name_count; 
    for (const auto& func : functions) {
        func_name_count[func.name]++;
    }

    // filter functions with empty return type or arg types
    functions.erase(std::remove_if(functions.begin(), functions.end(), [&func_name_count](const FuncSpec& func) {
        if (func.return_type.empty()) {
            return true;
        }
        for (const auto& arg_type : func.arg_types) {
            if (arg_type.empty()) {
                return true;
            }
        }
        return func_name_count[func.name] > 1 || func.access_level != AccessLevel::Public;
    }), functions.end());


    size_t virtual_count = 0;

    c += "namespace hooks {\n\n";
    c += "namespace " + safe_cpp_class_name + " {\n\n";

    for (const auto& func : functions) {
        if (func.type == FuncType::Static) {
            continue; // skip static functions
        }

        
        if (func.type == FuncType::Virtual || func.type == FuncType::PureVirtual) {
            virtual_count++;
        }

        std::string proxy_name = func.name;
        std::string proxy_type_name = "T" + proxy_name;
        std::string func_def;
        
        func_def += func.return_type + " " + proxy_name + "(" + class_name + "* this_ptr";
        for (size_t i = 0; i < func.arg_types.size(); ++i) {
            func_def += ", " + func.arg_types[i] + " arg" + std::to_string(i);
        }

        func_def += ")";

        // c += func_def + ";\n";
        // c += "using " + proxy_type_name + " = decltype(&" + proxy_name + ");\n";

        c += "using " + proxy_type_name + " = " + func.return_type + " (*)( " + class_name + "*";
        for (size_t i = 0; i < func.arg_types.size(); ++i) {
            c += ", " + func.arg_types[i];
        }
        c += ");\n";

        c += proxy_type_name + " original_" + proxy_name + " = nullptr;\n";

        c += func_def + " {\n";
        c += "\tassert(original_" + proxy_name + " != nullptr);\n";
        c += "\tstd::cout << \"[Proxy - " + class_name + "] Called " + func.name + "\\n\";\n";
        c += "\treturn original_" + proxy_name + "(this_ptr";
        for (size_t i = 0; i < func.arg_types.size(); ++i) {
            c += ", arg" + std::to_string(i);
        }
        c += ");\n";
        c += "}\n\n";
    }

    c += "} // namespace " + safe_cpp_class_name + "\n\n";
    c += "} // namespace hooks\n\n";

    c += "void inject_" + safe_cpp_class_name + "(" + class_name + "* original_obj_ptr) {\n";
   
    // call parent classes
    for (const auto& parent_class : parent_classes) {
        std::string safe_parent_class_name = parent_class; 
        std::replace(safe_parent_class_name.begin(), safe_parent_class_name.end(), ':', '_');
        c += "\tinject_" + safe_parent_class_name + "(original_obj_ptr);\n";
    }

    if (virtual_count > 0) {
        c += "\tusing vtable_t = void**;\n";
        c += "\tvtable_t* vtable_ptr = reinterpret_cast<vtable_t*>(original_obj_ptr);\n";
        c += "\tvtable_t vtable = *vtable_ptr;\n\n";

        // we need to remove memory protection to write into vtable
        c += "\n";
        c += "\t// Change memory protection to allow writing to vtable\n";
        // c += "\tuintptr_t latest_func_offset = reinterpret_cast<uintptr_t>(&" + safe_cpp_class_name + "::" + functions.back().name + ") - vtable_addr;\n";
        
        // c += "\tuintptr_t vtable_addr = reinterpret_cast<uintptr_t>(vtable);\n";
        // c += "\tuintptr_t latest_func_offset = reinterpret_cast<uintptr_t>(&hooks::" + safe_cpp_class_name + "::" + functions.back().name + ") - vtable_addr;\n";
        // c += "\tunprotect_memory(vtable_ptr, latest_func_offset + sizeof(void*));\n\n";

        // c += "\t// Change memory protection to allow writing to vtable\n";
        c += "\tsize_t page_size = sysconf(_SC_PAGESIZE);\n";
        c += "\tuintptr_t vtable_addr = reinterpret_cast<uintptr_t>(vtable);\n";
        c += "\tuintptr_t page_start = vtable_addr & ~(page_size - 1);\n";
        c += "\tif (mprotect(reinterpret_cast<void*>(page_start), page_size, PROT_READ | PROT_WRITE | PROT_EXEC) != 0) {\n";
        c += "\t\tperror(\"mprotect\");\n";
        c += "\t\treturn;\n";
        c += "\t}\n\n";
    }

    for (size_t i = 0; i < functions.size(); ++i) {
        const auto& func = functions[i];
        std::string proxy_name = func.name;
        std::string proxy_type_name = "hooks::" + safe_cpp_class_name + "::T" + proxy_name;
        

        if (func.type == FuncType::Static) {
            continue; // skip static functions
        }

        c += "\tusing MemberPtrType_" + proxy_name + " = ";
        c += func.return_type + " (" + class_name + "::*)(";
        for (size_t j = 0; j < func.arg_types.size(); ++j) {
            c += func.arg_types[j];
            if (j + 1 < func.arg_types.size()) {
                c += ", ";
            }
        }
        c += ") " + (func.is_const ? std::string("const") : std::string("")) + ";\n";
        c += "\tMemberPtrType_" + proxy_name + " member_ptr_" + proxy_name + " = &" + class_name + "::" + func.name + ";\n";

        if (functions[i].type != FuncType::Virtual && functions[i].type != FuncType::PureVirtual) {
            // c += "\thooks::" + safe_cpp_class_name + "::original_" + proxy_name + " = reinterpret_cast<" + proxy_type_name + ">(member_ptr_" + proxy_name + " );\n";
            std::string hook_ptr = "hook_ptr_" + proxy_name;
            std::string original_ptr = "original_ptr_" + proxy_name;
            std::string tramp_mem_ptr = "tramp_mem_ptr_" + proxy_name;
            
            
            c += "\tvoid* "+ hook_ptr + " = reinterpret_cast<void*>(&hooks::" + safe_cpp_class_name + "::" + proxy_name + ");\n";
            // c += "\tvoid* original_ptr = reinterpret_cast<void*>(member_ptr_" + proxy_name + ");\n";
            c += "\tvoid* " + original_ptr + " = reinterpret_cast<void*>(*(reinterpret_cast<uintptr_t*>(&member_ptr_" + proxy_name + ")));\n";
            
            // c += "\tvoid** tramp_mem_ptr = reinterpret_cast<void**>(&hooks::" + safe_cpp_class_name + "::original_" + proxy_name + ");\n";
            c += "\tvoid** " + tramp_mem_ptr + " = reinterpret_cast<void**>(&hooks::" + safe_cpp_class_name + "::original_" + proxy_name + ");\n";

            // c += "\tinjector::install_non_virtual_hook(original_ptr, hook_ptr, tramp_mem_ptr);\n\n";
            // c += "\tinjector::install_non_virtual_hook(" + original_ptr + ", " + hook_ptr + ", " + tramp_mem_ptr + ");\n\n";
            c += "\tinjector::install_smart_hook(" + original_ptr + ", " + hook_ptr + ", " + tramp_mem_ptr + ");\n\n";
            continue;
        } else {
            c += "\tlong vtable_index_" + proxy_name + " = *reinterpret_cast<long*>(&member_ptr_" + proxy_name + ") / sizeof(void*);\n";
            c += "\thooks::" + safe_cpp_class_name + "::original_" + proxy_name + " = reinterpret_cast<" + proxy_type_name + ">(vtable[vtable_index_" + proxy_name + "]);\n";
            c += "\tvtable[vtable_index_" + proxy_name + "] = reinterpret_cast<void*>(&hooks::" + safe_cpp_class_name + "::" + proxy_name + ");\n\n";
        }
    }


    c += "}\n";


    return c;
}

std::string proxy_class(const std::string& class_name, const std::vector<std::string>& header_files) {
    context_t context = extract_class_def_context(class_name, header_files);
    return proxy_class_inner(class_name, context);
}


// usage: injector
// $> injector tvm::runtime::vm::VirtualMachine(tvm/runtime/vm/vm.h)

const char * usage_message = R"(
Usage: injector class_name#header_file [class_name#header_file ...]
To add custom include paths, set the INCLUDE_PATHS environment variable with paths separated by colons.
example:
$> INCLUDE_PATHS="/custom/include/path:/another/path" injector tvm::runtime::vm::VirtualMachine(tvm/runtime/vm/vm.h)
)";

int main (int argc, char** argv) {

    if (argc < 2) {
        std::cerr << usage_message << "\n";
        return 1;
    }

    std::vector<std::string> args;
    for (int i = 1; i < argc; ++i) {
        args.push_back(argv[i]);
    }

    if (args[0] == "--help" || args[0] == "-h") {
        std::cerr << usage_message << "\n";
        return 0;
    }


    std::string gen_code = "";
    gen_code += "// This code is generated by injector by mamodev\n";
    gen_code += "// Do not edit manually\n\n";
    gen_code += "// This code is free to use and modify but keep this notice intact.\n\n";

    gen_code += required_headers;
    gen_code += "\n";

    gen_code += std::string(reinterpret_cast<const char*>(hde64_bin), sizeof(hde64_bin));

    gen_code += "\n";
    gen_code += "namespace injector {\n";

    gen_code += "\n";
    gen_code += "size_t get_page_size() {\n";
    gen_code += "\treturn sysconf(_SC_PAGESIZE);\n";
    gen_code += "}\n";

    gen_code += "\n";
    gen_code += "void unprotect_memory(void* addr, size_t length) {\n";
    gen_code += "\tuintptr_t page_start = reinterpret_cast<uintptr_t>(addr) & ~(get_page_size() - 1);\n";
    gen_code += "\tassert(mprotect(reinterpret_cast<void*>(page_start), length, PROT_READ | PROT_WRITE | PROT_EXEC) == 0 && \"mprotect failed\");\n";
    gen_code += "}\n";

    gen_code += "\n";

    gen_code += str_allocate_nearby_memory;
    gen_code += "\n";
    gen_code += str_install_smart_hook;
    gen_code += "\n";

    for (const auto& arg : args) {
        // Find the position of the hash sign (#)
        size_t hash_pos = arg.find("#");

        // Check if '#' was found and is not the very last character
        if (hash_pos == std::string::npos || hash_pos == arg.length() - 1 || hash_pos == 0) {
            std::cerr << "Invalid argument format: " << arg << " (Expected: class_name#header_file)\n";
            std::cerr << usage_message << "\n";
            return 1;
        }

        // The class name is everything before the '#'
        std::string class_name = arg.substr(0, hash_pos);
        
        // The header file is everything after the '#'
        std::string header_file = arg.substr(hash_pos + 1);

        gen_code += proxy_class(class_name, {header_file});
        gen_code += "\n\n";
    }

    gen_code += "\n} // namespace injector\n";

    std::cout << gen_code << "\n";
    return 0;
}