diff --git a/src/plugins/HookSystem.cpp b/src/plugins/HookSystem.cpp index ca38fe95..3ac75c04 100644 --- a/src/plugins/HookSystem.cpp +++ b/src/plugins/HookSystem.cpp @@ -8,6 +8,7 @@ #include #include #include +#include CFunctionHook::CFunctionHook(HANDLE owner, void* source, void* destination) { m_pSource = source; @@ -64,10 +65,14 @@ CFunctionHook::SInstructionProbe CFunctionHook::probeMinimumJumpSize(void* start } CFunctionHook::SAssembly CFunctionHook::fixInstructionProbeRIPCalls(const SInstructionProbe& probe) { + SAssembly returns; + // analyze the code and fix what we know how to. uint64_t currentAddress = (uint64_t)m_pSource; // actually newline + 1 - size_t lastAsmNewline = 0; + size_t lastAsmNewline = 0; + // needle for destination binary + size_t currentDestinationOffset = 0; std::string assemblyBuilder; for (auto& len : probe.insSizes) { @@ -82,17 +87,27 @@ CFunctionHook::SAssembly CFunctionHook::fixInstructionProbeRIPCalls(const SInstr return {}; const uint64_t DESTINATION = currentAddress + OFFSET + len; - if (code.starts_with("mov")) { - // mov +0xdeadbeef(%rip), %rax - assemblyBuilder += std::format("movabs $0x{:x}, {}\n", DESTINATION, tokens[2]); - } else if (code.starts_with("call")) { + if (code.starts_with("call")) { // call +0xdeadbeef(%rip) assemblyBuilder += std::format("pushq %rax\nmovabs $0x{:x}, %rax\ncallq *%rax\npopq %rax\n", DESTINATION); + currentDestinationOffset += 14; } else if (code.starts_with("lea")) { // lea 0xdeadbeef(%rip), %rax assemblyBuilder += std::format("movabs $0x{:x}, {}\n", DESTINATION, tokens[2]); + currentDestinationOffset += 10; } else { - return {}; + auto ADDREND = code.find("(%rip)"); + auto ADDRSTART = (code.substr(0, ADDREND).find_last_of(' ')); + + if (ADDREND == std::string::npos || ADDRSTART == std::string::npos) + return {}; + + const uint64_t PREDICTEDRIP = (uint64_t)m_pTrampolineAddr + currentDestinationOffset + len; + const bool POSITIVE = DESTINATION > PREDICTEDRIP; + const uint64_t NEWRIPOFFSET = POSITIVE ? DESTINATION - PREDICTEDRIP : PREDICTEDRIP - DESTINATION; + + assemblyBuilder += std::format("{} {}0x{:x}{}\n", code.substr(0, ADDRSTART), POSITIVE ? '+' : '-', NEWRIPOFFSET, code.substr(ADDREND)); + currentDestinationOffset += len; } } else if (code.contains("invalid")) { std::vector bytes; @@ -101,6 +116,7 @@ CFunctionHook::SAssembly CFunctionHook::fixInstructionProbeRIPCalls(const SInstr if (len == 4 && bytes[0] == 0xF3 && bytes[1] == 0x0F && bytes[2] == 0x1E && bytes[3] == 0xFA) { // F3 0F 1E FA = endbr64, udis doesn't understand that one assemblyBuilder += "endbr64\n"; + currentDestinationOffset += 4; } else { // raise error, unknown op std::string strBytes; @@ -112,6 +128,7 @@ CFunctionHook::SAssembly CFunctionHook::fixInstructionProbeRIPCalls(const SInstr } } else { assemblyBuilder += code + "\n"; + currentDestinationOffset += len; } lastAsmNewline = probe.assembly.find("\n", lastAsmNewline) + 1; @@ -131,7 +148,7 @@ CFunctionHook::SAssembly CFunctionHook::fixInstructionProbeRIPCalls(const SInstr } std::ifstream ifs("/tmp/hypr/.hookbinary2.o", std::ios::binary); - SAssembly returns = {std::vector(std::istreambuf_iterator(ifs), {})}; + returns = {std::vector(std::istreambuf_iterator(ifs), {})}; ifs.close(); std::filesystem::remove("/tmp/hypr/.hookcode.asm"); std::filesystem::remove("/tmp/hypr/.hookbinary.o"); @@ -158,6 +175,10 @@ bool CFunctionHook::hook() { // nop static constexpr uint8_t NOP = 0x90; + // alloc trampoline + const auto MAX_TRAMPOLINE_SIZE = HOOK_TRAMPOLINE_MAX_SIZE; // we will never need more. + m_pTrampolineAddr = (void*)g_pFunctionHookSystem->getAddressForTrampo(); + // probe instructions to be trampolin'd SInstructionProbe probe; try { @@ -174,9 +195,12 @@ bool CFunctionHook::hook() { const size_t HOOKSIZE = PROBEFIXEDASM.bytes.size(); const size_t ORIGSIZE = probe.len; - // alloc trampoline - const auto TRAMPOLINE_SIZE = sizeof(ABSOLUTE_JMP_ADDRESS) + HOOKSIZE + sizeof(PUSH_RAX); - m_pTrampolineAddr = mmap(NULL, TRAMPOLINE_SIZE, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + const auto TRAMPOLINE_SIZE = sizeof(ABSOLUTE_JMP_ADDRESS) + HOOKSIZE + sizeof(PUSH_RAX); + + if (TRAMPOLINE_SIZE > MAX_TRAMPOLINE_SIZE) { + Debug::log(ERR, "[functionhook] failed, not enough space in trampo to alloc:\n{}", probe.assembly); + return false; + } m_pOriginalBytes = malloc(ORIGSIZE); memcpy(m_pOriginalBytes, m_pSource, ORIGSIZE); @@ -236,14 +260,11 @@ bool CFunctionHook::unhook() { // revert mprot mprotect((uint8_t*)m_pSource - ((uint64_t)m_pSource) % sysconf(_SC_PAGE_SIZE), sysconf(_SC_PAGE_SIZE), PROT_READ | PROT_EXEC); - // unmap - munmap(m_pTrampolineAddr, m_iTrampoLen); - // reset vars m_bActive = false; m_iHookLen = 0; m_iTrampoLen = 0; - m_pTrampolineAddr = nullptr; + m_pTrampolineAddr = nullptr; // no unmapping, it's managed by the HookSystem m_pOriginalBytes = nullptr; free(m_pOriginalBytes); @@ -263,3 +284,108 @@ bool CHookSystem::removeHook(CFunctionHook* hook) { void CHookSystem::removeAllHooksFrom(HANDLE handle) { std::erase_if(m_vHooks, [&](const auto& other) { return other->m_pOwner == handle; }); } + +static uintptr_t seekNewPageAddr() { + const uint64_t PAGESIZE_VAR = sysconf(_SC_PAGE_SIZE); + auto MAPS = std::ifstream("/proc/self/maps"); + + uint64_t lastStart = 0, lastEnd = 0; + + std::string line; + while (std::getline(MAPS, line)) { + CVarList props{line, 0, 's', true}; + + uint64_t start = 0, end = 0; + if (props[0].empty()) { + Debug::log(WARN, "seekNewPageAddr: unexpected line in self maps"); + continue; + } + + CVarList startEnd{props[0], 0, '-', true}; + + try { + start = std::stoull(startEnd[0], nullptr, 16); + end = std::stoull(startEnd[1], nullptr, 16); + } catch (std::exception& e) { + Debug::log(WARN, "seekNewPageAddr: unexpected line in self maps: {}", line); + continue; + } + + Debug::log(LOG, "seekNewPageAddr: page 0x{:x} - 0x{:x}", start, end); + + if (lastStart == 0) { + lastStart = start; + lastEnd = end; + continue; + } + + if (start - lastEnd > PAGESIZE_VAR * 2) { + Debug::log(LOG, "seekNewPageAddr: found gap: 0x{:x}-0x{:x} ({} bytes)", lastEnd, start, start - lastEnd); + MAPS.close(); + return lastEnd; + } + + lastStart = start; + lastEnd = end; + } + + MAPS.close(); + return 0; +} + +uint64_t CHookSystem::getAddressForTrampo() { + // yes, technically this creates a memory leak of 64B every hook creation. But I don't care. + // tracking all the users of the memory would be painful. + // Nobody will hook 100k times, and even if, that's only 640kB. Nothing. + + SAllocatedPage* page = nullptr; + for (auto& p : pages) { + if (p.used + HOOK_TRAMPOLINE_MAX_SIZE > p.len) + continue; + + page = &p; + break; + } + + if (!page) + page = &pages.emplace_back(); + + if (!page->addr) { + // allocate it + Debug::log(LOG, "getAddressForTrampo: Allocating new page for hooks"); + const uint64_t PAGESIZE_VAR = sysconf(_SC_PAGE_SIZE); + const auto BASEPAGEADDR = seekNewPageAddr(); + for (int attempt = 0; attempt < 2; ++attempt) { + for (int i = 2; i >= 0; --i) { + const auto PAGEADDR = BASEPAGEADDR + i * PAGESIZE_VAR; + + page->addr = (uint64_t)mmap((void*)PAGEADDR, PAGESIZE_VAR, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + page->len = PAGESIZE_VAR; + page->used = 0; + + Debug::log(LOG, "Attempted to allocate 0x{:x}, got 0x{:x}", PAGEADDR, page->addr); + + if (page->addr == (uint64_t)MAP_FAILED) + continue; + if (page->addr != PAGEADDR && attempt == 0) { + munmap((void*)page->addr, PAGESIZE_VAR); + page->addr = 0; + page->len = 0; + continue; + } + + break; + } + if (page->addr) + break; + } + } + + const auto ADDRFORCONSUMER = page->addr + page->used; + + page->used += HOOK_TRAMPOLINE_MAX_SIZE; + + Debug::log(LOG, "getAddressForTrampo: Returning addr 0x{:x} for page at 0x{:x}", ADDRFORCONSUMER, page->addr); + + return ADDRFORCONSUMER; +} \ No newline at end of file diff --git a/src/plugins/HookSystem.hpp b/src/plugins/HookSystem.hpp index 8714ac32..f0a5a11d 100644 --- a/src/plugins/HookSystem.hpp +++ b/src/plugins/HookSystem.hpp @@ -4,7 +4,8 @@ #include #include -#define HANDLE void* +#define HANDLE void* +#define HOOK_TRAMPOLINE_MAX_SIZE 64 class CFunctionHook { public: @@ -60,6 +61,18 @@ class CHookSystem { private: std::vector> m_vHooks; + + uint64_t getAddressForTrampo(); + + struct SAllocatedPage { + uint64_t addr = 0; + uint64_t len = 0; + uint64_t used = 0; + }; + + std::vector pages; + + friend class CFunctionHook; }; inline std::unique_ptr g_pFunctionHookSystem; \ No newline at end of file