From bbc7ed08f821176e03defc46bacc4f464067d4c6 Mon Sep 17 00:00:00 2001 From: kin4stat Date: Mon, 19 Sep 2022 01:44:51 +0300 Subject: [PATCH] fix x86 build issuefixed bugs with context saving; (perfectly saved for now) added more functionality to naked hooks; add more tests --- CMakeLists.txt | 6 + include/kthook/x64/kthook_detail.hpp | 2 +- include/kthook/x64/kthook_impl.hpp | 220 ++++++--- include/kthook/x86/kthook_detail.hpp | 8 +- include/kthook/x86/kthook_impl.hpp | 209 ++++++--- .../kthook/x86_64/kthook_x86_64_detail.hpp | 125 +++-- tests/CMakeLists.txt | 3 + tests/aggregates_test.cpp | 27 +- tests/create_context_test.cpp | 182 +++++++ tests/lots_of_arguments_test.cpp | 15 +- tests/non_executable_test.cpp | 9 +- tests/return_value_test.cpp | 443 ++++++++++++++++++ tests/simple_test.cpp | 93 ++-- tests/test_common.hpp | 65 ++- 14 files changed, 1148 insertions(+), 259 deletions(-) create mode 100644 tests/create_context_test.cpp create mode 100644 tests/return_value_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index dc63be4..226d5fa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -16,6 +16,12 @@ target_include_directories(${PROJECT_NAME} INTERFACE $) { auto res = - (!(std::is_trivial_v && std::is_standard_layout_v && (sizeof(Ret) % 2 == 0) && sizeof(Ret) <= 8)); + (!(std::is_trivial_v && std::is_standard_layout_v && (sizeof(Ret) % 2 == 0 || sizeof(Ret) == 1) && sizeof(Ret) <= 8)); if (res) used = true; used_integral_registers += res; } diff --git a/include/kthook/x64/kthook_impl.hpp b/include/kthook/x64/kthook_impl.hpp index 71b5afe..41c5a0c 100644 --- a/include/kthook/x64/kthook_impl.hpp +++ b/include/kthook/x64/kthook_impl.hpp @@ -4,23 +4,6 @@ namespace kthook { #pragma pack(push, 1) struct cpu_ctx { - std::uintptr_t rax; - std::uintptr_t rbx; - std::uintptr_t rcx; - std::uintptr_t rdx; - std::uintptr_t rsp; - std::uintptr_t rbp; - std::uintptr_t rsi; - std::uintptr_t rdi; - std::uintptr_t r8; - std::uintptr_t r9; - std::uintptr_t r10; - std::uintptr_t r11; - std::uintptr_t r12; - std::uintptr_t r13; - std::uintptr_t r14; - std::uintptr_t r15; - struct eflags { public: std::uint32_t CF : 1; @@ -64,9 +47,25 @@ struct cpu_ctx { private: std::uint32_t reserved5 : 10; std::uint32_t reserved6 : 32; - } flags; + }; - std::uint8_t align; + std::uintptr_t rax; + std::uintptr_t rbx; + std::uintptr_t rcx; + std::uintptr_t rdx; + std::uintptr_t rsp; + std::uintptr_t rbp; + std::uintptr_t rsi; + std::uintptr_t rdi; + std::uintptr_t r8; + std::uintptr_t r9; + std::uintptr_t r10; + std::uintptr_t r11; + std::uintptr_t r12; + std::uintptr_t r13; + std::uintptr_t r14; + std::uintptr_t r15; + eflags* flags; }; #pragma pack(pop) @@ -379,11 +378,25 @@ class kthook_simple { detail::create_trampoline(hook_address, jump_gen); jump_gen->L(UserCode); - jump_gen->mov(rax, rsp); - jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); if constexpr (create_context) { - jump_gen->mov(rsp, reinterpret_cast(&context.align)); jump_gen->pushfq(); + + jump_gen->push(rax); + jump_gen->push(rbx); + jump_gen->mov(rax, ptr[rsp + sizeof(std::uintptr_t)]); + jump_gen->xchg(rbx, rax); + jump_gen->mov(rax, reinterpret_cast(&context.flags)); + jump_gen->mov(ptr[rax], rbx); + jump_gen->pop(rbx); + jump_gen->pop(rax); + jump_gen->add(rsp, sizeof(cpu_ctx::eflags)); + + jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); + jump_gen->mov(rax, rsp); + jump_gen->mov(ptr[reinterpret_cast(&context.rsp)], rax); + jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); + + jump_gen->mov(rsp, reinterpret_cast(&context.flags)); jump_gen->push(r15); jump_gen->push(r14); jump_gen->push(r13); @@ -395,14 +408,18 @@ class kthook_simple { jump_gen->push(rdi); jump_gen->push(rsi); jump_gen->push(rbp); - jump_gen->push(rsp); + jump_gen->sub(rsp, sizeof(std::uintptr_t)); jump_gen->push(rdx); jump_gen->push(rcx); jump_gen->push(rbx); - jump_gen->mov(rax, ptr[reinterpret_cast(&last_return_address)]); + + jump_gen->mov(rax, ptr[reinterpret_cast(&context.rsp)]); jump_gen->mov(rsp, rax); - jump_gen->mov(ptr[reinterpret_cast(&context.rsp)], rax); + } else { + jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); } + jump_gen->mov(rax, rsp); + jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); #if defined(KTHOOK_64_WIN) constexpr std::array registers{rcx, rdx, r8, r9}; @@ -429,7 +446,7 @@ class kthook_simple { #ifdef KTHOOK_64_WIN jump_gen->mov(rax, reinterpret_cast(this)); // set rsp to next stack argument pointer - jump_gen->add(rsp, static_cast(sizeof(void*) * (registers.size()))); + jump_gen->add(rsp, static_cast(sizeof(void*) * registers.size())); jump_gen->push(0); jump_gen->push(rax); #else @@ -443,7 +460,7 @@ class kthook_simple { #ifdef KTHOOK_64_WIN // return the rsp to its initial state - jump_gen->sub(rsp, static_cast(sizeof(void*) * (registers.size()))); + jump_gen->sub(rsp, static_cast(sizeof(void*) * registers.size())); #else #endif @@ -467,7 +484,6 @@ class kthook_simple { jump_gen->add(rsp, 32); #endif // push original return address and return - jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); jump_gen->mov(rax, ptr[reinterpret_cast(&last_return_address)]); jump_gen->push(rax); jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); @@ -475,7 +491,6 @@ class kthook_simple { } else { using_ptr_to_return_address = true; - jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); jump_gen->mov(rax, rsp); jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); @@ -688,11 +703,26 @@ class kthook_signal { jump_gen->nop(3); detail::create_trampoline(hook_address, jump_gen); jump_gen->L(UserCode); - jump_gen->mov(rax, rsp); - jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); + if constexpr (create_context) { - jump_gen->mov(rsp, reinterpret_cast(&context.align)); jump_gen->pushfq(); + + jump_gen->push(rax); + jump_gen->push(rbx); + jump_gen->mov(rax, ptr[rsp + sizeof(std::uintptr_t)]); + jump_gen->xchg(rbx, rax); + jump_gen->mov(rax, reinterpret_cast(&context.flags)); + jump_gen->mov(ptr[rax], rbx); + jump_gen->pop(rbx); + jump_gen->pop(rax); + jump_gen->add(rsp, sizeof(cpu_ctx::eflags)); + + jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); + jump_gen->mov(rax, rsp); + jump_gen->mov(ptr[reinterpret_cast(&context.rsp)], rax); + jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); + + jump_gen->mov(rsp, reinterpret_cast(&context.flags)); jump_gen->push(r15); jump_gen->push(r14); jump_gen->push(r13); @@ -704,14 +734,18 @@ class kthook_signal { jump_gen->push(rdi); jump_gen->push(rsi); jump_gen->push(rbp); - jump_gen->push(rsp); + jump_gen->sub(rsp, sizeof(std::uintptr_t)); jump_gen->push(rdx); jump_gen->push(rcx); jump_gen->push(rbx); - jump_gen->mov(rax, ptr[reinterpret_cast(&last_return_address)]); + + jump_gen->mov(rax, ptr[reinterpret_cast(&context.rsp)]); jump_gen->mov(rsp, rax); - jump_gen->mov(ptr[reinterpret_cast(&context.rsp)], rax); + } else { + jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); } + jump_gen->mov(rax, rsp); + jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); #if defined(KTHOOK_64_WIN) constexpr std::array registers{rcx, rdx, r8, r9}; @@ -739,7 +773,7 @@ class kthook_signal { #ifdef KTHOOK_64_WIN jump_gen->mov(rax, reinterpret_cast(this)); // set rsp to next stack argument pointer - jump_gen->add(rsp, static_cast(sizeof(void*) * (registers.size()))); + jump_gen->add(rsp, static_cast(sizeof(void*) * registers.size())); jump_gen->push(0); jump_gen->push(rax); #else @@ -753,7 +787,7 @@ class kthook_signal { #ifdef KTHOOK_64_WIN // return the rsp to its initial state - jump_gen->sub(rsp, static_cast(sizeof(void*) * (registers.size()))); + jump_gen->sub(rsp, static_cast(sizeof(void*) * registers.size())); #else #endif @@ -777,7 +811,6 @@ class kthook_signal { jump_gen->add(rsp, 32); #endif // push original return address and return - jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); jump_gen->mov(rax, ptr[reinterpret_cast(&last_return_address)]); jump_gen->push(rax); jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); @@ -785,7 +818,6 @@ class kthook_signal { } else { using_ptr_to_return_address = true; - jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); jump_gen->mov(rax, rsp); jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); @@ -886,20 +918,30 @@ class kthook_naked { }; using cb_type = std::function; + friend std::uintptr_t detail::naked_relay(kthook_naked*); public: kthook_naked() : info(0, nullptr) { - } + }; - kthook_naked(std::uintptr_t destination, bool force_enable = true) - : info(destination, nullptr) { + kthook_naked(std::uintptr_t destination, cb_type callback_, bool force_enable = true) + : info(destination, nullptr), + callback(std::move(callback_)) { if (force_enable) { install(); } } - kthook_naked(void* destination, bool force_enable = true) - : kthook_naked(reinterpret_cast(destination), force_enable) { + kthook_naked(std::uintptr_t destination) + : info(destination, nullptr) { + } + + kthook_naked(void* destination) + : kthook_naked(reinterpret_cast(destination)) { + } + + kthook_naked(void* destination, cb_type callback, bool force_enable = true) + : kthook_naked(reinterpret_cast(destination), callback, force_enable) { } ~kthook_naked() { remove(); } @@ -951,18 +993,25 @@ class kthook_naked { auto hook_address = info.hook_address; - Xbyak::Label UserCode, ret_addr; + Xbyak::Label UserCode, ret_addr, jump_in_trampoline, skip_bytes, jump_out; jump_gen->jmp(UserCode, Xbyak::CodeGenerator::LabelType::T_NEAR); jump_gen->nop(3); detail::create_trampoline(hook_address, jump_gen); jump_gen->L(UserCode); - std::uintptr_t saved_rsp; + jump_gen->push(std::uintptr_t{}); + jump_gen->pushfq(); + jump_gen->mov(ptr[reinterpret_cast(&context.rax)], rax); jump_gen->mov(rax, rsp); - jump_gen->mov(ptr[reinterpret_cast(&saved_rsp)], rax); - jump_gen->mov(rsp, reinterpret_cast(&context.align)); - jump_gen->pushfq(); + jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); + jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); + + // [rsp - 0x00] == 0 + // [rsp - 0x08] == eflags + // last_return_address == rsp - 0x10 + + jump_gen->mov(rsp, reinterpret_cast(&context.flags)); jump_gen->push(r15); jump_gen->push(r14); jump_gen->push(r13); @@ -974,28 +1023,40 @@ class kthook_naked { jump_gen->push(rdi); jump_gen->push(rsi); jump_gen->push(rbp); - jump_gen->push(rsp); + jump_gen->sub(rsp, sizeof(std::uintptr_t)); jump_gen->push(rdx); jump_gen->push(rcx); jump_gen->push(rbx); - jump_gen->mov(rax, ptr[reinterpret_cast(&saved_rsp)]); + + jump_gen->mov(rax, ptr[reinterpret_cast(&last_return_address)]); jump_gen->mov(rsp, rax); + jump_gen->mov(ptr[reinterpret_cast(&context.flags)], rax); + jump_gen->add(rax, sizeof(cpu_ctx::eflags) + sizeof(std::uintptr_t)); jump_gen->mov(ptr[reinterpret_cast(&context.rsp)], rax); - jump_gen->mov(rax, info.hook_address + hook_size); + jump_gen->mov(rax, info.hook_address); jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); + jump_gen->mov(rax, ret_addr); #if defined(KTHOOK_64_WIN) jump_gen->mov(rcx, reinterpret_cast(this)); - jump_gen->sub(rsp, 32); + jump_gen->sub(rsp, sizeof(std::uintptr_t) * 8); #elif defined(KTHOOK_64_GCC) jump_gen->mov(rdi, reinterpret_cast(this)); #endif - jump_gen->mov(rax, ret_addr); + jump_gen->push(rax); jump_gen->jmp(ptr[rip]); - jump_gen->db(reinterpret_cast(&detail::naked_relay), 8); + jump_gen->db(reinterpret_cast(&detail::naked_relay), sizeof(std::uintptr_t)); jump_gen->L(ret_addr); + jump_gen->cmp(rax, -1); + jump_gen->je(jump_out); + + jump_gen->cmp(rax, hook_size); + jump_gen->jl(jump_in_trampoline); + + jump_gen->L(jump_out); + jump_gen->mov(rsp, reinterpret_cast(&context.rbx)); jump_gen->pop(rbx); jump_gen->pop(rcx); @@ -1012,16 +1073,55 @@ class kthook_naked { jump_gen->pop(r13); jump_gen->pop(r14); jump_gen->pop(r15); + + jump_gen->mov(rax, ptr[reinterpret_cast(&context.rsp)]); + jump_gen->mov(rsp, rax); + jump_gen->mov(rax, ptr[reinterpret_cast(&last_return_address)]); + jump_gen->push(rax); + jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); + jump_gen->sub(rsp, sizeof(cpu_ctx::eflags)); jump_gen->popfq(); - jump_gen->mov(rax, ptr[reinterpret_cast(&saved_rsp)]); + jump_gen->ret(); + + jump_gen->L(jump_in_trampoline); + + jump_gen->mov(rsp, rax); + jump_gen->mov(rax, skip_bytes); + jump_gen->add(rax, rsp); + + jump_gen->mov(ptr[reinterpret_cast(&last_return_address)], rax); + + jump_gen->mov(rsp, reinterpret_cast(&context.rbx)); + jump_gen->pop(rbx); + jump_gen->pop(rcx); + jump_gen->pop(rdx); + jump_gen->add(rsp, 8); + jump_gen->pop(rbp); + jump_gen->pop(rsi); + jump_gen->pop(rdi); + jump_gen->pop(r8); + jump_gen->pop(r9); + jump_gen->pop(r10); + jump_gen->pop(r11); + jump_gen->pop(r12); + jump_gen->pop(r13); + jump_gen->pop(r14); + jump_gen->pop(r15); + + jump_gen->mov(rax, ptr[reinterpret_cast(&context.rsp)]); jump_gen->mov(rsp, rax); + jump_gen->mov(rax, ptr[reinterpret_cast(&last_return_address)]); + jump_gen->push(rax); jump_gen->mov(rax, ptr[reinterpret_cast(&context.rax)]); + jump_gen->sub(rsp, sizeof(cpu_ctx::eflags)); + jump_gen->popfq(); - detail::create_trampoline(info.hook_address, jump_gen, true); + jump_gen->ret(); - jump_gen->jmp(ptr[rip]); - jump_gen->db(reinterpret_cast(&last_return_address), 8); + jump_gen->L(skip_bytes); + + detail::create_trampoline(info.hook_address, jump_gen); detail::flush_intruction_cache(jump_gen->getCode(), jump_gen->getSize()); return jump_gen->getCode(); @@ -1099,7 +1199,7 @@ class kthook_naked { std::unique_ptr trampoline_gen; const std::uint8_t* relay_jump{nullptr}; - bool installed; + bool installed{false}; }; } // namespace kthook diff --git a/include/kthook/x86/kthook_detail.hpp b/include/kthook/x86/kthook_detail.hpp index 72e92e5..9bd9844 100644 --- a/include/kthook/x86/kthook_detail.hpp +++ b/include/kthook/x86/kthook_detail.hpp @@ -71,7 +71,7 @@ struct function_traits { }; template -struct function_traits { +struct function_traits { static constexpr auto args_count = sizeof...(Args); static constexpr auto convention = cconv::ccdecl; using args = std::tuple; @@ -87,7 +87,7 @@ struct function_traits { }; template -struct function_traits { +struct function_traits { static constexpr auto args_count = sizeof...(Args); static constexpr auto convention = cconv::cthiscall; using args = std::tuple; @@ -95,7 +95,7 @@ struct function_traits { }; template -struct function_traits { +struct function_traits { static constexpr auto args_count = sizeof...(Args); static constexpr auto convention = cconv::cstdcall; using args = std::tuple; @@ -103,7 +103,7 @@ struct function_traits { }; template -struct function_traits { +struct function_traits { static constexpr auto args_count = sizeof...(Args); static constexpr auto convention = cconv::cfastcall; using args = std::tuple; diff --git a/include/kthook/x86/kthook_impl.hpp b/include/kthook/x86/kthook_impl.hpp index c8ce6cb..e82d34d 100644 --- a/include/kthook/x86/kthook_impl.hpp +++ b/include/kthook/x86/kthook_impl.hpp @@ -4,18 +4,7 @@ namespace kthook { #pragma pack(push, 1) struct cpu_ctx { - cpu_ctx() = default; - - std::uintptr_t edi; - std::uintptr_t esi; - std::uintptr_t ebp; - std::uintptr_t esp; - std::uintptr_t ebx; - std::uintptr_t edx; - std::uintptr_t ecx; - std::uintptr_t eax; - - struct EFLAGS { + struct eflags { public: std::uintptr_t CF : 1; @@ -57,15 +46,27 @@ struct cpu_ctx { private: std::uintptr_t reserved5 : 10; - } flags; + }; + + cpu_ctx() = default; + + std::uintptr_t edi; + std::uintptr_t esi; + std::uintptr_t ebp; + std::uintptr_t esp; + std::uintptr_t ebx; + std::uintptr_t edx; + std::uintptr_t ecx; + std::uintptr_t eax; - std::uint8_t align; + eflags* flags; }; #pragma pack(pop) namespace detail { struct cpu_ctx_empty { std::uintptr_t ecx; + void* flags; }; inline bool create_trampoline(std::uintptr_t hook_address, @@ -212,6 +213,7 @@ class kthook_simple { public: kthook_simple() : info(0, nullptr) { + context.flags = new cpu_ctx::eflags{}; }; kthook_simple(std::uintptr_t destination, cb_type callback_, bool force_enable = true) @@ -220,10 +222,12 @@ class kthook_simple { if (force_enable) { install(); } + context.flags = new cpu_ctx::eflags{}; } kthook_simple(std::uintptr_t destination) : info(destination, nullptr) { + context.flags = new cpu_ctx::eflags{}; } kthook_simple(void* destination) @@ -244,7 +248,10 @@ class kthook_simple { : kthook_simple(reinterpret_cast(destination), callback_, force_enable) { } - ~kthook_simple() { remove(); } + ~kthook_simple() { + remove(); + delete reinterpret_cast(context.flags); + } bool install() { if (installed) return false; @@ -282,7 +289,8 @@ class kthook_simple { template )> void set_cb_wrapped(C cb) { callback = [cb = std::forward(cb)](auto&&... args) { - std::apply(cb, detail::bind_values>(std::forward_as_tuple(std::forward(args)...))); + std::apply(cb, detail::bind_values>( + std::forward_as_tuple(std::forward(args)...))); }; } @@ -300,7 +308,7 @@ class kthook_simple { return reinterpret_cast(const_cast(trampoline_gen->getCode())); } - template + template Ret call_trampoline(Ts&&... args) const { return std::apply(get_trampoline(), detail::unpack(std::forward(args)...)); } @@ -323,19 +331,26 @@ class kthook_simple { jump_gen->L(UserCode); if constexpr (create_context) { - // save esp - jump_gen->mov(eax, esp); - jump_gen->mov(ptr[&last_return_address], eax); - jump_gen->mov(esp, reinterpret_cast(&context.align)); jump_gen->pushfd(); + jump_gen->mov(ptr[reinterpret_cast(&context.eax)], eax); + jump_gen->mov(ptr[reinterpret_cast(&context.ecx)], ecx); + jump_gen->mov(eax, ptr[reinterpret_cast(&context.flags)]); + jump_gen->mov(ecx, ptr[esp]); + jump_gen->mov(ptr[eax], ecx); + jump_gen->mov(eax, ptr[reinterpret_cast(&context.eax)]); + jump_gen->mov(ecx, ptr[reinterpret_cast(&context.ecx)]); + jump_gen->add(esp, sizeof(cpu_ctx::eflags)); + + jump_gen->mov(ptr[&last_return_address], esp); + jump_gen->mov(esp, reinterpret_cast(&context.flags)); jump_gen->pushad(); jump_gen->mov(esp, ptr[&last_return_address]); jump_gen->mov(ptr[reinterpret_cast(&context.esp)], esp); } - // save return address jump_gen->mov(eax, ptr[esp]); jump_gen->mov(ptr[&last_return_address], eax); + constexpr bool can_be_pushed = []() { if constexpr (function::args_count > 0) { using first = std::tuple_element_t<0, Args>; @@ -344,14 +359,22 @@ class kthook_simple { return false; }(); constexpr bool is_thiscall = (function::convention == detail::traits::cconv::cthiscall); + constexpr bool is_fastcall = (function::convention == detail::traits::cconv::cfastcall); #ifdef _WIN32 jump_gen->pop(eax); if constexpr (!std::is_void_v) { - if constexpr (sizeof(Ret) > 8 || !std::is_trivial_v) { - jump_gen->pop(eax); - jump_gen->push(reinterpret_cast(this)); - jump_gen->push(eax); - jump_gen->mov(eax, ptr[reinterpret_cast(&last_return_address)]); + constexpr bool is_fully_nontrivial = !std::is_trivial_v || !std::is_trivially_destructible_v; + if constexpr (sizeof(Ret) > 8 || is_fully_nontrivial) { + if constexpr ((is_thiscall || is_fastcall) && (sizeof(Ret) % 2 != 0 || is_fully_nontrivial)) { + jump_gen->push(reinterpret_cast(this)); + if constexpr (is_thiscall) + jump_gen->push(ecx); + } else { + jump_gen->pop(eax); + jump_gen->push(reinterpret_cast(this)); + jump_gen->push(eax); + jump_gen->mov(eax, ptr[reinterpret_cast(&last_return_address)]); + } } else { if constexpr (is_thiscall) { if constexpr (can_be_pushed) { @@ -373,8 +396,9 @@ class kthook_simple { // if Ret is class or union, memory for return value as first argument(hidden) // so we need to push our hook pointer after this hidden argument if constexpr (!std::is_void_v) { + constexpr bool is_fully_nontrivial = !std::is_trivial_v || !std::is_trivially_destructible_v; if constexpr (std::is_class_v || std::is_union_v || sizeof(Ret) > 8) { - if constexpr (is_thiscall) { + if constexpr (is_thiscall && (sizeof(Ret) % 2 != 0 || is_fully_nontrivial)) { jump_gen->push(reinterpret_cast(this)); jump_gen->push(ecx); } else { @@ -408,11 +432,7 @@ class kthook_simple { // call relay for restoring stack pointer after call jump_gen->call(relay_ptr); jump_gen->add(esp, 4); - jump_gen->mov(ptr[reinterpret_cast(&context.ecx)], ecx); - jump_gen->mov(ecx, ptr[&last_return_address]); - jump_gen->push(ecx); - jump_gen->mov(ecx, ptr[reinterpret_cast(&context.ecx)]); - jump_gen->ret(); + jump_gen->jmp(ptr[&last_return_address]); } else { jump_gen->push(eax); jump_gen->jmp(relay_ptr); @@ -519,6 +539,7 @@ class kthook_signal { public: kthook_signal() : info(0, nullptr) { + context.flags = new cpu_ctx::eflags{}; } kthook_signal(std::uintptr_t destination, bool force_enable = true) @@ -526,6 +547,7 @@ class kthook_signal { if (force_enable) { install(); } + context.flags = new cpu_ctx::eflags{}; } kthook_signal(void* destination, bool force_enable = true) @@ -537,7 +559,10 @@ class kthook_signal { : kthook_signal(reinterpret_cast(destination), force_enable) { } - ~kthook_signal() { remove(); } + ~kthook_signal() { + remove(); + delete reinterpret_cast(context.flags); + } bool install() { if (installed) return false; @@ -600,11 +625,18 @@ class kthook_signal { jump_gen->L(UserCode); if constexpr (create_context) { - // save esp - jump_gen->mov(eax, esp); - jump_gen->mov(ptr[&last_return_address], eax); - jump_gen->mov(esp, reinterpret_cast(&context.align)); jump_gen->pushfd(); + jump_gen->mov(ptr[reinterpret_cast(&context.eax)], eax); + jump_gen->mov(ptr[reinterpret_cast(&context.ecx)], ecx); + jump_gen->mov(eax, ptr[reinterpret_cast(&context.flags)]); + jump_gen->mov(ecx, ptr[esp]); + jump_gen->mov(ptr[eax], ecx); + jump_gen->mov(eax, ptr[reinterpret_cast(&context.eax)]); + jump_gen->mov(ecx, ptr[reinterpret_cast(&context.ecx)]); + jump_gen->add(esp, sizeof(cpu_ctx::eflags)); + + jump_gen->mov(ptr[&last_return_address], esp); + jump_gen->mov(esp, reinterpret_cast(&context.flags)); jump_gen->pushad(); jump_gen->mov(esp, ptr[&last_return_address]); jump_gen->mov(ptr[reinterpret_cast(&context.esp)], esp); @@ -613,7 +645,6 @@ class kthook_signal { jump_gen->mov(eax, ptr[esp]); jump_gen->mov(ptr[&last_return_address], eax); - // pop return address out constexpr bool can_be_pushed = []() { if constexpr (function::args_count > 0) { using first = std::tuple_element_t<0, Args>; @@ -622,14 +653,22 @@ class kthook_signal { return false; }(); constexpr bool is_thiscall = (function::convention == detail::traits::cconv::cthiscall); + constexpr bool is_fastcall = (function::convention == detail::traits::cconv::cfastcall); #ifdef _WIN32 jump_gen->pop(eax); if constexpr (!std::is_void_v) { - if constexpr (sizeof(Ret) > 8 || !std::is_trivial_v) { - jump_gen->pop(eax); - jump_gen->push(reinterpret_cast(this)); - jump_gen->push(eax); - jump_gen->mov(eax, ptr[reinterpret_cast(&last_return_address)]); + constexpr bool is_fully_nontrivial = !std::is_trivial_v || !std::is_trivially_destructible_v; + if constexpr (sizeof(Ret) > 8 || is_fully_nontrivial) { + if constexpr ((is_thiscall || is_fastcall) && (sizeof(Ret) % 2 != 0 || is_fully_nontrivial)) { + jump_gen->push(reinterpret_cast(this)); + if constexpr (is_thiscall) + jump_gen->push(ecx); + } else { + jump_gen->pop(eax); + jump_gen->push(reinterpret_cast(this)); + jump_gen->push(eax); + jump_gen->mov(eax, ptr[reinterpret_cast(&last_return_address)]); + } } else { if constexpr (is_thiscall) { if constexpr (can_be_pushed) { @@ -686,11 +725,7 @@ class kthook_signal { // call relay for restoring stack pointer after call jump_gen->call(relay_ptr); jump_gen->add(esp, 4); - jump_gen->mov(ptr[reinterpret_cast(&context.ecx)], ecx); - jump_gen->mov(ecx, ptr[&last_return_address]); - jump_gen->push(ecx); - jump_gen->mov(ecx, ptr[reinterpret_cast(&context.ecx)]); - jump_gen->ret(); + jump_gen->jmp(ptr[&last_return_address]); } else { jump_gen->push(eax); jump_gen->jmp(relay_ptr); @@ -781,6 +816,7 @@ class kthook_naked { } }; + friend std::uintptr_t detail::naked_relay(kthook_naked*); public: kthook_naked() : info(0, nullptr) { @@ -855,7 +891,7 @@ class kthook_naked { auto hook_address = info.hook_address; - Xbyak::Label UserCode, ret_addr; + Xbyak::Label UserCode, ret_addr, jump_in_trampoline, skip_bytes, jump_out; // this jump gets nopped when hook.remove() is called jump_gen->jmp(UserCode, Xbyak::CodeGenerator::LabelType::T_NEAR); jump_gen->nop(3); @@ -864,36 +900,85 @@ class kthook_naked { detail::create_trampoline(hook_address, jump_gen); jump_gen->L(UserCode); - // save esp - jump_gen->mov(ptr[reinterpret_cast(&context.eax)], eax); - jump_gen->mov(eax, esp); - jump_gen->mov(ptr[&last_return_address], eax); - jump_gen->mov(eax, ptr[reinterpret_cast(&context.eax)]); - - jump_gen->mov(esp, reinterpret_cast(&context.align)); jump_gen->pushfd(); + + jump_gen->mov(ptr[&last_return_address], esp); + + // &context.end -> esp + // pushad -> context.registers + // memory.esp -> esp + // esp -> context.esp + // &label ret_addr -> eax + // info.hook_address -> last_return_address + jump_gen->mov(esp, reinterpret_cast(&context.flags)); jump_gen->pushad(); jump_gen->mov(esp, ptr[&last_return_address]); + jump_gen->mov(ptr[reinterpret_cast(&context.flags)], esp); + jump_gen->sub(esp, sizeof(cpu_ctx::eflags)); jump_gen->mov(ptr[reinterpret_cast(&context.esp)], esp); + jump_gen->add(esp, sizeof(cpu_ctx::eflags)); jump_gen->mov(eax, ret_addr); - jump_gen->mov(dword[reinterpret_cast(&last_return_address)], info.hook_address + hook_size); + jump_gen->mov(dword[reinterpret_cast(&last_return_address)], info.hook_address); + // push this + // push eax(&label ret_addr) jump_gen->push(reinterpret_cast(this)); jump_gen->push(eax); + // GOTO callback(call) jump_gen->jmp(reinterpret_cast(&detail::naked_relay)); jump_gen->L(ret_addr); + + // restore stack jump_gen->add(esp, 0x04); + + // if need to skip trampoline + jump_gen->cmp(eax, ~0u); + jump_gen->je(jump_out); + + // if need to jump inside trampoline + jump_gen->cmp(eax, hook_size); + jump_gen->jl(jump_in_trampoline); + + jump_gen->L(jump_out); + + // esp -> &context.top + // context.registers -> popad + // context.esp -> esp + // stack -> popfd + // goto RETURN_ADDR jump_gen->mov(esp, reinterpret_cast(&context.edi)); jump_gen->popad(); - jump_gen->mov(esp, reinterpret_cast(&context.flags)); - jump_gen->popfd(); jump_gen->mov(esp, ptr[reinterpret_cast(&context.esp)]); + jump_gen->add(esp, sizeof(cpu_ctx::eflags)); + jump_gen->popfd(); + jump_gen->jmp(ptr[&last_return_address]); - detail::create_trampoline(info.hook_address, jump_gen); - + jump_gen->L(jump_in_trampoline); + // eax -> esp + // &label skip_bytes -> eax + // eax += esp + // so eax is address inside trampoline + jump_gen->mov(esp, eax); + jump_gen->mov(eax, skip_bytes); + jump_gen->add(eax, esp); + + // eax -> RETURN_ADDR + // esp -> &context.top + // popad + // context.esp -> esp + // stack -> popfd + // goto RETURN_ADDR + jump_gen->mov(dword[reinterpret_cast(&last_return_address)], eax); + jump_gen->mov(esp, reinterpret_cast(&context.edi)); + jump_gen->popad(); + jump_gen->mov(esp, ptr[reinterpret_cast(&context.esp)]); + jump_gen->add(esp, sizeof(cpu_ctx::eflags)); + jump_gen->popfd(); jump_gen->jmp(ptr[&last_return_address]); + jump_gen->L(skip_bytes); + detail::create_trampoline(info.hook_address, jump_gen); detail::flush_intruction_cache(jump_gen->getCode(), jump_gen->getSize()); return jump_gen->getCode(); diff --git a/include/kthook/x86_64/kthook_x86_64_detail.hpp b/include/kthook/x86_64/kthook_x86_64_detail.hpp index 4f91071..c62caaa 100644 --- a/include/kthook/x86_64/kthook_x86_64_detail.hpp +++ b/include/kthook/x86_64/kthook_x86_64_detail.hpp @@ -425,6 +425,45 @@ constexpr decltype(auto) unpack(Ts&&... args) { return unpack(std::forward_as_tuple(std::forward(args)...), std::make_index_sequence{}); } +// https://github.com/TsudaKageyu/minhook/blob/master/src/trampoline.h +#pragma pack(push, 1) +struct JCC_ABS { + std::uint8_t opcode; // 7* 0E: J** +16 + std::uint8_t dummy0; + std::uint8_t dummy1; // FF25 00000000: JMP [+6] + std::uint8_t dummy2; + std::uint32_t dummy3; + std::uint64_t address; // Absolute destination address +}; + +struct CALL_ABS { + std::uint8_t opcode0; // FF15 00000002: CALL [+6] + std::uint8_t opcode1; + std::uint32_t dummy0; + std::uint8_t dummy1; // EB 08: JMP +10 + std::uint8_t dummy2; + std::uint64_t address; // Absolute destination address +}; + +struct JMP_ABS { + std::uint8_t opcode0; // FF25 00000000: JMP [+6] + std::uint8_t opcode1; + std::uint32_t dummy; + std::uint64_t address; // Absolute destination address +}; + +typedef struct { + std::uint8_t opcode; // E9/E8 xxxxxxxx: JMP/CALL +5+xxxxxxxx + std::uint32_t operand; // Relative destination address +} JMP_REL, CALL_REL; + +struct JCC_REL { + std::uint8_t opcode0; // 0F8* xxxxxxxx: J** +6+xxxxxxxx + std::uint8_t opcode1; + std::uint32_t operand; // Relative destination address +}; +#pragma pack(pop) + template inline Ret signal_relay(HookPtrType* this_hook, Args&... args) { if constexpr (std::is_void_v) { @@ -465,63 +504,61 @@ inline Ret common_relay(CallbackT& cb, HookPtrType* this_hook, Args&... args) { return this_hook->get_trampoline()(args...); } - template #ifdef KTHOOK_32 -#ifdef __GNUC__ -__attribute__((cdecl)) inline void +#ifndef __GNUC__ +std::uintptr_t __cdecl #else -inline void __cdecl +__attribute__((cdecl)) std::uintptr_t #endif #else -inline void +std::uintptr_t #endif naked_relay(HookType* this_hook) { + auto ret_addr = this_hook->get_return_address(); + auto& cb = this_hook->get_callback(); if (cb) { cb(*this_hook); } + if (ret_addr != this_hook->get_return_address()) { + ret_addr = this_hook->get_return_address(); + auto hook_addr = this_hook->info.hook_address; + auto hook_size = this_hook->hook_size; + auto& original = this_hook->info.original_code; + if (ret_addr > hook_addr && ret_addr - hook_addr <= hook_size) { + std::uintptr_t result{}; + auto m = (ret_addr - hook_addr < hook_size) ? ret_addr - hook_addr : hook_size; + for (auto i = 0u; i < m;) { + hde hs; + std::size_t op_size = hde_disasm(original.get() + i, &hs); + if (hs.flags & F_ERROR) return 0; + if (((hs.opcode & 0xF0) == 0x70) || // one byte jump + ((hs.opcode & 0xFC) == 0xE0) || // LOOPNZ/LOOPZ/LOOP/JECXZ + ((hs.opcode2 & 0xF0) == 0x80)) { + result += sizeof(JCC_REL); + } + else if ((hs.opcode & 0xFD) == 0xE9) { + result += sizeof(JMP_REL); + } + else if (hs.opcode == 0xE8) { + result += sizeof(JMP_REL); + } + else { + result += op_size; + } + i += op_size; + } + // if need to get into trampoline + return result; + } else { + return ~0u; + } + } + // if need to get into trampoline + return 0; } -// https://github.com/TsudaKageyu/minhook/blob/master/src/trampoline.h -#pragma pack(push, 1) -struct JCC_ABS { - std::uint8_t opcode; // 7* 0E: J** +16 - std::uint8_t dummy0; - std::uint8_t dummy1; // FF25 00000000: JMP [+6] - std::uint8_t dummy2; - std::uint32_t dummy3; - std::uint64_t address; // Absolute destination address -}; - -struct CALL_ABS { - std::uint8_t opcode0; // FF15 00000002: CALL [+6] - std::uint8_t opcode1; - std::uint32_t dummy0; - std::uint8_t dummy1; // EB 08: JMP +10 - std::uint8_t dummy2; - std::uint64_t address; // Absolute destination address -}; - -struct JMP_ABS { - std::uint8_t opcode0; // FF25 00000000: JMP [+6] - std::uint8_t opcode1; - std::uint32_t dummy; - std::uint64_t address; // Absolute destination address -}; - -typedef struct { - std::uint8_t opcode; // E9/E8 xxxxxxxx: JMP/CALL +5+xxxxxxxx - std::uint32_t operand; // Relative destination address -} JMP_REL, CALL_REL; - -struct JCC_REL { - std::uint8_t opcode0; // 0F8* xxxxxxxx: J** +6+xxxxxxxx - std::uint8_t opcode1; - std::uint32_t operand; // Relative destination address -}; -#pragma pack(pop) - inline std::size_t detect_hook_size(std::uintptr_t addr) { size_t size = 0; while (size < 5) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fa056fe..7b4c1b8 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -44,6 +44,9 @@ foreach(_test_file ${TEST_SRC_FILES}) configure_file(${_test_file} "${CMAKE_BINARY_DIR}/${_test_name}-${conv}.cpp" COPYONLY) add_executable(${_test_name}-${conv} "${CMAKE_BINARY_DIR}/${_test_name}-${conv}.cpp") target_compile_definitions(${_test_name}-${conv} PRIVATE TEST_CCONV=${conv} CONV_NAME=${conv_name}) + if (conv STREQUAL CTHISCALL) + target_compile_definitions(${_test_name}-${conv} PRIVATE IS_THISCALL) + endif() target_link_libraries(${_test_name}-${conv} gtest gtest_main ${PROJECT_NAME}) add_test(${_test_name}-${conv} ${_test_name}-${conv}) set_tests_properties(${_test_name}-${conv} PROPERTIES TIMEOUT 5) diff --git a/tests/aggregates_test.cpp b/tests/aggregates_test.cpp index deffbae..08e8f4a 100644 --- a/tests/aggregates_test.cpp +++ b/tests/aggregates_test.cpp @@ -53,33 +53,24 @@ DECLARE_SIZE_ENLARGER(); class A { public: - NO_OPTIMIZE static BigAggregate -#ifdef KTHOOK_32 - TEST_CCONV -#endif + NO_OPTIMIZE static BigAggregate CCONV big_test_func(BigAggregate value) { SIZE_ENLARGER(); return value; } - NO_OPTIMIZE static MediumAggregate -#ifdef KTHOOK_32 - TEST_CCONV -#endif + NO_OPTIMIZE static MediumAggregate CCONV medium_test_func(MediumAggregate value) { SIZE_ENLARGER(); return value; } - NO_OPTIMIZE static SmallAggregate -#ifdef KTHOOK_32 - TEST_CCONV -#endif + NO_OPTIMIZE static SmallAggregate CCONV small_test_func(SmallAggregate value) { SIZE_ENLARGER(); return value; } }; -TEST(KthookSimpleTest, CREATE_NAME(BigAggregate)) { +TEST(kthook_simple, big_aggregate) { kthook::kthook_simple hook{&A::big_test_func}; hook.install(); @@ -96,7 +87,7 @@ TEST(KthookSimpleTest, CREATE_NAME(BigAggregate)) { EXPECT_EQ(counter, 2); } -TEST(KthookSignalTest, CREATE_NAME(BigAggregate)) { +TEST(kthook_signal, big_aggregate) { kthook::kthook_signal hook{&A::big_test_func}; { @@ -123,7 +114,7 @@ TEST(KthookSignalTest, CREATE_NAME(BigAggregate)) { } } -TEST(KthookSimpleTest, CREATE_NAME(MediumAggregate)) { +TEST(kthook_simple, medium_aggregate) { kthook::kthook_simple hook{&A::medium_test_func}; hook.install(); @@ -141,7 +132,7 @@ TEST(KthookSimpleTest, CREATE_NAME(MediumAggregate)) { EXPECT_EQ(counter, 2); } -TEST(KthookSignalTest, CREATE_NAME(MediumAggregate)) { +TEST(kthook_signal, medium_aggregate) { kthook::kthook_signal hook{&A::medium_test_func}; { @@ -169,7 +160,7 @@ TEST(KthookSignalTest, CREATE_NAME(MediumAggregate)) { } } -TEST(KthookSimpleTest, CREATE_NAME(SmallAggregate)) { +TEST(kthook_simple, small_aggregate) { kthook::kthook_simple hook{&A::small_test_func}; hook.install(); @@ -186,7 +177,7 @@ TEST(KthookSimpleTest, CREATE_NAME(SmallAggregate)) { EXPECT_EQ(counter, 2); } -TEST(KthookSignalTest, CREATE_NAME(SmallAggregate)) { +TEST(kthook_signal, small_aggregate) { kthook::kthook_signal hook{&A::small_test_func}; { diff --git a/tests/create_context_test.cpp b/tests/create_context_test.cpp new file mode 100644 index 0000000..2d13e97 --- /dev/null +++ b/tests/create_context_test.cpp @@ -0,0 +1,182 @@ +#include "gtest/gtest.h" +#include "kthook/kthook.hpp" +#include "test_common.hpp" + +#include "xbyak/xbyak.h" + +#define EQUALITY_CHECK(x) \ + if (lhs.x != rhs.x) { \ + return testing::AssertionFailure() << "lhs." << #x << "(" << lhs.x << ") != " << "rhs." << #x << "(" << rhs.x << ")"; \ + } + +#define IMM(X) reinterpret_cast(X) + +DECLARE_SIZE_ENLARGER(); + +class A { +public: + NO_OPTIMIZE static void THISCALL_REPLACEMENT + test_func() { + SIZE_ENLARGER(); + } +}; + +Xbyak::CodeGenerator gen; +kthook::cpu_ctx ctx; +bool inited = [] { + ctx.flags = new kthook::cpu_ctx::eflags; + return true; +}(); + +testing::AssertionResult operator==(const kthook::cpu_ctx::eflags& lhs, const kthook::cpu_ctx::eflags& rhs) { + EQUALITY_CHECK(CF) + EQUALITY_CHECK(PF) + EQUALITY_CHECK(AF) + EQUALITY_CHECK(ZF) + EQUALITY_CHECK(SF) + EQUALITY_CHECK(TF) + EQUALITY_CHECK(IF) + EQUALITY_CHECK(DF) + EQUALITY_CHECK(OF) + EQUALITY_CHECK(IOPL) + EQUALITY_CHECK(NT) + EQUALITY_CHECK(RF) + EQUALITY_CHECK(VM) + EQUALITY_CHECK(AC) + EQUALITY_CHECK(VIF) + EQUALITY_CHECK(VIP) + EQUALITY_CHECK(ID) + return testing::AssertionSuccess() << "args are equal"; +} + +#ifdef KTHOOK_32 +auto generate_code() { + using namespace Xbyak::util; + gen.pushfd(); + gen.push(eax); + gen.push(ebx); + gen.mov(eax, ptr[esp + 8]); + gen.mov(ebx, ptr[IMM(&ctx.flags)]); + gen.mov(ptr[ebx], eax); + gen.pop(ebx); + gen.pop(eax); + gen.popfd(); + gen.mov(ptr[IMM(&ctx.eax)], eax); + gen.mov(ptr[IMM(&ctx.edi)], edi); + gen.mov(ptr[IMM(&ctx.esi)], esi); + gen.mov(ptr[IMM(&ctx.ebp)], ebp); + gen.mov(ptr[IMM(&ctx.ebx)], ebx); + gen.mov(ptr[IMM(&ctx.edx)], edx); + gen.mov(ptr[IMM(&ctx.ecx)], ecx); + gen.mov(ptr[IMM(&ctx.eax)], eax); + gen.jmp(&A::test_func); + return gen.getCode(); +} +testing::AssertionResult operator==(const kthook::cpu_ctx& lhs, const kthook::cpu_ctx& rhs) { + EQUALITY_CHECK(edi) + EQUALITY_CHECK(esi) + EQUALITY_CHECK(ebp) + EQUALITY_CHECK(ebx) + EQUALITY_CHECK(edx) + EQUALITY_CHECK(ecx) + EQUALITY_CHECK(eax) + return testing::AssertionSuccess() << "args are equal"; +} +#else +auto generate_code() { +#define WRAP_MEM_REG(X) gen.mov(rax, X); gen.mov(ptr[IMM(&ctx.X)], rax) + + using namespace Xbyak::util; + gen.pushfq(); + gen.push(rax); + gen.mov(rax, ptr[rsp + 8]); + gen.mov(ptr[IMM(ctx.flags)], rax); + WRAP_MEM_REG(rbx); + WRAP_MEM_REG(rcx); + WRAP_MEM_REG(rdx); + gen.add(rsp, sizeof(std::uintptr_t) * 2); + WRAP_MEM_REG(rsp); + gen.sub(rsp, sizeof(std::uintptr_t) * 2); + WRAP_MEM_REG(rbp); + WRAP_MEM_REG(rsi); + WRAP_MEM_REG(rdi); + WRAP_MEM_REG(r8); + WRAP_MEM_REG(r9); + WRAP_MEM_REG(r10); + WRAP_MEM_REG(r11); + WRAP_MEM_REG(r12); + WRAP_MEM_REG(r13); + WRAP_MEM_REG(r14); + WRAP_MEM_REG(r15); + gen.pop(rax); + gen.mov(ptr[IMM(&ctx.rax)], rax); + gen.popfq(); + gen.jmp(ptr[rip]); + gen.db(IMM(&A::test_func), 8); + return gen.getCode(); + +#undef WRAP_MEM_REG +} + +testing::AssertionResult operator==(const kthook::cpu_ctx& lhs, const kthook::cpu_ctx& rhs) { + EQUALITY_CHECK(rax) + EQUALITY_CHECK(rbx) + EQUALITY_CHECK(rcx) + EQUALITY_CHECK(rdx) + EQUALITY_CHECK(rsp) + EQUALITY_CHECK(rbp) + EQUALITY_CHECK(rsi) + EQUALITY_CHECK(rdi) + EQUALITY_CHECK(r8) + EQUALITY_CHECK(r9) + EQUALITY_CHECK(r10) + EQUALITY_CHECK(r11) + EQUALITY_CHECK(r12) + EQUALITY_CHECK(r13) + EQUALITY_CHECK(r14) + EQUALITY_CHECK(r15) + return testing::AssertionSuccess() << "args are equal"; +} +#endif + +#undef EQUALITY_CHECK + +TEST(kthook_simple, function) { + kthook::kthook_simple hook{&A::test_func}; + EXPECT_TRUE(hook.install()); + + hook.set_cb([](const auto& hook, auto&&... args) { + EXPECT_TRUE(hook.get_context() == ctx); + }); + + std::memset(ctx.flags, 0, sizeof(*ctx.flags)); + std::memset(&ctx, 0, sizeof(ctx) - sizeof(ctx.flags)); + generate_code()(); +} + +TEST(kthook_signal, function) { + kthook::kthook_signal hook{&A::test_func, false}; + EXPECT_TRUE(hook.install()); + + hook.before.connect([](const auto& hook, auto&&... args) { + EXPECT_TRUE(hook.get_context() == ctx); + return false; + }); + + std::memset(ctx.flags, 0, sizeof(*ctx.flags)); + std::memset(&ctx, 0, sizeof(ctx) - sizeof(ctx.flags)); + generate_code()(); +} + +TEST(kthook_naked, function) { + kthook::kthook_naked hook{reinterpret_cast(&A::test_func)}; + EXPECT_TRUE(hook.install()); + + hook.set_cb([](const auto& hook, auto&&... args) { + EXPECT_TRUE(hook.get_context() == ctx); + }); + + std::memset(ctx.flags, 0, sizeof(*ctx.flags)); + std::memset(&ctx, 0, sizeof(ctx) - sizeof(ctx.flags)); + generate_code()(); +} diff --git a/tests/lots_of_arguments_test.cpp b/tests/lots_of_arguments_test.cpp index ebc0cb5..49058cf 100644 --- a/tests/lots_of_arguments_test.cpp +++ b/tests/lots_of_arguments_test.cpp @@ -37,17 +37,14 @@ DECLARE_SIZE_ENLARGER(); class A { public: - NO_OPTIMIZE static void -#ifdef KTHOOK_32 - TEST_CCONV -#endif + NO_OPTIMIZE static void CCONV test_func(int v1, float v2, long long v3, double v4, short v5, char v6, int v7, long double v8, float v9, int v10, int v11, long v12, long long v13, int v14, int v15, int v16) { SIZE_ENLARGER(); } }; -TEST(KthookSimpleLotsArgsTest, CREATE_NAME(HandlesKthookSimple)) { +TEST(kthook_simple, function) { kthook::kthook_simple hook{&A::test_func}; hook.install(); @@ -61,13 +58,15 @@ TEST(KthookSimpleLotsArgsTest, CREATE_NAME(HandlesKthookSimple)) { EXPECT_EQ(counter, 1); } -TEST(KthookSimpleLotsArgsTest, CREATE_NAME(Skip)) { +TEST(kthook_simple, arg_skip) { kthook::kthook_simple hook{&A::test_func}; hook.install(); int counter = 0; hook.set_cb_wrapped([&counter](const auto& h, kthook::take<8> v1_v8, float v9, kthook::take<7> v10_v16) { + ++counter; + return h.call_trampoline(v1_v8, v9, v10_v16); }); @@ -75,7 +74,7 @@ TEST(KthookSimpleLotsArgsTest, CREATE_NAME(Skip)) { EXPECT_EQ(counter, 1); } -TEST(KthookSignalLotsArgsTest, CREATE_NAME(HandlesKthookSignalBefore)) { +TEST(kthook_signal_before, function) { kthook::kthook_signal hook{&A::test_func}; int counter = 0; hook.before += [&counter](const auto& hook, auto&&... args) { @@ -87,7 +86,7 @@ TEST(KthookSignalLotsArgsTest, CREATE_NAME(HandlesKthookSignalBefore)) { EXPECT_EQ(counter, 1); } -TEST(KthookSignalLotsArgsTest, CREATE_NAME(HandlesKthookSignalAfter)) { +TEST(kthook_signal_after, function) { kthook::kthook_signal hook{&A::test_func}; int counter = 0; hook.after += [&counter](const auto& hook, auto&&... args) { diff --git a/tests/non_executable_test.cpp b/tests/non_executable_test.cpp index b136cc3..fd85d30 100644 --- a/tests/non_executable_test.cpp +++ b/tests/non_executable_test.cpp @@ -5,12 +5,17 @@ using test_signature = void (*)(); constexpr std::uintptr_t null_mem = 0x1000; -TEST(KthookSimpleTest, CREATE_NAME(HandlesKthookSimple)) { +TEST(kthook_simple, null_func) { kthook::kthook_simple hook{null_mem}; EXPECT_FALSE(hook.install()); } -TEST(KthookSiginalTest, CREATE_NAME(HandlesKthookSignal)) { +TEST(kthook_signal, null_func) { kthook::kthook_signal hook{null_mem}; EXPECT_FALSE(hook.install()); +} + +TEST(kthook_naked, null_func) { + kthook::kthook_naked hook{null_mem}; + EXPECT_FALSE(hook.install()); } \ No newline at end of file diff --git a/tests/return_value_test.cpp b/tests/return_value_test.cpp new file mode 100644 index 0000000..8482bd4 --- /dev/null +++ b/tests/return_value_test.cpp @@ -0,0 +1,443 @@ +#include "gtest/gtest.h" +#include "kthook/kthook.hpp" +#include "test_common.hpp" + +constexpr unsigned char byte_ret = 0x50; +constexpr unsigned short word_ret = 0x5050; +constexpr unsigned long dword_ret = 0x50505050ul; +constexpr unsigned long long qword_ret = 0x5050505050505050ull; + +DECLARE_SIZE_ENLARGER(); + + +struct Bool { + NO_OPTIMIZE static bool CCONV + test_func(int value) { + SIZE_ENLARGER() + return value; + } +}; + + +struct Byte { + NO_OPTIMIZE static unsigned char CCONV + test_func(int value) { + SIZE_ENLARGER() + return value; + } +}; + +struct Word { + NO_OPTIMIZE static unsigned short CCONV + test_func(int value) { + SIZE_ENLARGER() + return value; + } +}; + +struct Dword { + NO_OPTIMIZE static unsigned long CCONV + test_func(int value) { + SIZE_ENLARGER() + return value; + } +}; + +struct Qword { + NO_OPTIMIZE static unsigned long long CCONV + test_func(int value) { + SIZE_ENLARGER() + return value; + } +}; + +struct NonTriviallyConstructible { + NonTriviallyConstructible() { + v = 10; + } + + explicit NonTriviallyConstructible(int v) + : v(v) { + } + + NO_OPTIMIZE static NonTriviallyConstructible CCONV + test_func(int value) { + SIZE_ENLARGER() + return {}; + } + + friend bool operator==(const NonTriviallyConstructible& lhs, const NonTriviallyConstructible& rhs) { + return lhs.v == rhs.v; + } + + int v; +}; + +struct NonTriviallyCopyable { + NonTriviallyCopyable() = default; + NonTriviallyCopyable(const NonTriviallyCopyable& rhs) { + this->v = rhs.v; + } + NonTriviallyCopyable(int v) + : v(v) { + } + + NO_OPTIMIZE static NonTriviallyCopyable CCONV + test_func(int value) { + SIZE_ENLARGER() + return {}; + } + + friend bool operator==(const NonTriviallyCopyable& lhs, const NonTriviallyCopyable& rhs) { + return lhs.v == rhs.v; + } + + int v; +}; + +struct NonStandardLayout { + NonStandardLayout() = default; + explicit NonStandardLayout(int v) + : v(v) { + } + + NO_OPTIMIZE static NonStandardLayout CCONV + test_func(int value) { + SIZE_ENLARGER(); + return {}; + } + + friend bool operator==(const NonStandardLayout& lhs, const NonStandardLayout& rhs) { + return lhs.v == rhs.v; + } + +private: + int v{}; +}; + +TEST(kthook_simple, bool) { + kthook::kthook_simple hook{&Bool::test_func}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + + hook.set_cb([test_arg](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return true; + }); + + EXPECT_TRUE(Bool::test_func(test_arg)); + } + + { + const auto test_arg = GET_RANDOM_INT(); + + hook.set_cb([test_arg](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return false; + }); + + EXPECT_FALSE(Bool::test_func(test_arg)); + } +} + +TEST(kthook_simple, byte) { + kthook::kthook_simple hook{&Byte::test_func}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = GET_RANDOM_BYTE(); + + hook.set_cb([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return test_ret; + }); + + EXPECT_EQ(Byte::test_func(test_arg), test_ret); + } +} + +TEST(kthook_simple, word) { + kthook::kthook_simple hook{&Word::test_func}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = GET_RANDOM_WORD(); + + hook.set_cb([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return test_ret; + }); + + EXPECT_EQ(Word::test_func(test_arg), test_ret); + } +} + +TEST(kthook_simple, dword) { + kthook::kthook_simple hook{&Dword::test_func}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = GET_RANDOM_DWORD(); + + hook.set_cb([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return test_ret; + }); + + EXPECT_EQ(Dword::test_func(test_arg), test_ret); + } +} + +TEST(kthook_simple, qword) { + kthook::kthook_simple hook{&Qword::test_func}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = GET_RANDOM_QWORD(); + + hook.set_cb([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return test_ret; + }); + + EXPECT_EQ(Qword::test_func(test_arg), test_ret); + } +} + +TEST(kthook_simple, non_trivially_constructible) { + kthook::kthook_simple hook{&NonTriviallyConstructible::test_func}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = NonTriviallyConstructible{GET_RANDOM_INT()}; + + hook.set_cb([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return test_ret; + }); + + EXPECT_EQ(NonTriviallyConstructible::test_func(test_arg), test_ret); + } +} + +TEST(kthook_simple, non_trivially_copyable) { + kthook::kthook_simple hook{&NonTriviallyCopyable::test_func}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = NonTriviallyCopyable{GET_RANDOM_INT()}; + + hook.set_cb([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return test_ret; + }); + + EXPECT_EQ(NonTriviallyCopyable::test_func(test_arg), test_ret); + } +} + +TEST(kthook_simple, non_standard_copyable) { + kthook::kthook_simple hook{&NonStandardLayout::test_func}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = NonStandardLayout{GET_RANDOM_INT()}; + + hook.set_cb([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + + hook.get_trampoline()(value); + + EXPECT_EQ(value, test_arg); + return test_ret; + }); + + EXPECT_EQ(NonStandardLayout::test_func(test_arg), test_ret); + } +} + +TEST(kthook_signal, bool) { + kthook::kthook_signal hook{&Bool::test_func, false}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + + auto connection = hook.before.scoped_connect([test_arg](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(true); + }); + + EXPECT_TRUE(Bool::test_func(test_arg)); + } + + { + const auto test_arg = GET_RANDOM_INT(); + + auto connection = hook.before.scoped_connect([test_arg](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(false); + }); + + EXPECT_FALSE(Bool::test_func(test_arg)); + } +} + +TEST(kthook_signal, byte) { + kthook::kthook_signal hook{&Byte::test_func, false}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = GET_RANDOM_BYTE(); + + auto connection = hook.before.scoped_connect([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(test_ret); + }); + + EXPECT_EQ(Byte::test_func(test_arg), test_ret); + } +} + +TEST(kthook_signal, word) { + kthook::kthook_signal hook{&Word::test_func, false}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = GET_RANDOM_WORD(); + + auto connection = hook.before.scoped_connect([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(test_ret); + }); + + EXPECT_EQ(Word::test_func(test_arg), test_ret); + } +} + +TEST(kthook_signal, dword) { + kthook::kthook_signal hook{&Dword::test_func, false}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = GET_RANDOM_DWORD(); + + auto connection = hook.before.scoped_connect([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(test_ret); + }); + + EXPECT_EQ(Dword::test_func(test_arg), test_ret); + } +} + +TEST(kthook_signal, qword) { + kthook::kthook_signal hook{&Qword::test_func, false}; + EXPECT_TRUE(hook.install()); + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = GET_RANDOM_QWORD(); + + auto connection = hook.before.scoped_connect([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(test_ret); + }); + + EXPECT_EQ(Qword::test_func(test_arg), test_ret); + } +} + +TEST(kthook_signal, non_trivially_constructible) { + kthook::kthook_signal hook{&NonTriviallyConstructible::test_func, false}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = NonTriviallyConstructible{GET_RANDOM_INT()}; + + auto connection = hook.before.scoped_connect([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(test_ret); + }); + + EXPECT_EQ(NonTriviallyConstructible::test_func(test_arg), test_ret); + } +} + +TEST(kthook_signal, non_trivially_copyable) { + kthook::kthook_signal hook{&NonTriviallyCopyable::test_func, false}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = NonTriviallyCopyable{GET_RANDOM_INT()}; + + auto connection = hook.before.scoped_connect([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(test_ret); + }); + + EXPECT_EQ(NonTriviallyCopyable::test_func(test_arg), test_ret); + } +} + +TEST(kthook_signal, non_standard_layout) { + kthook::kthook_signal hook{&NonStandardLayout::test_func, false}; + EXPECT_TRUE(hook.install()); + + { + const auto test_arg = GET_RANDOM_INT(); + const auto test_ret = NonStandardLayout{GET_RANDOM_INT()}; + + auto connection = hook.before.scoped_connect([test_arg, test_ret](const auto& hook, int& value) { + EXPECT_EQ(value, test_arg); + return std::make_optional(test_ret); + }); + + EXPECT_EQ(NonStandardLayout::test_func(test_arg), test_ret); + } +} diff --git a/tests/simple_test.cpp b/tests/simple_test.cpp index 5e99090..0291c70 100644 --- a/tests/simple_test.cpp +++ b/tests/simple_test.cpp @@ -9,17 +9,32 @@ DECLARE_SIZE_ENLARGER(); class A { public: - NO_OPTIMIZE static int -#ifdef KTHOOK_32 - TEST_CCONV -#endif - test_func(int value) { + NO_OPTIMIZE static int CCONV + test_func(int value) { SIZE_ENLARGER(); return value; } }; -TEST(KthookSimpleTest, CREATE_NAME(HandlesSimpleUsage)) { +class AT { +public: + NO_OPTIMIZE static int TEST_THISCALL + test_func(int value) { + SIZE_ENLARGER(); + return value; + } +}; + +class AF { +public: + NO_OPTIMIZE static int TEST_FASTCALL + test_func(int value) { + SIZE_ENLARGER(); + return value; + } +}; + +TEST(kthook_simple, function) { kthook::kthook_simple hook{&A::test_func}; hook.install(); @@ -35,75 +50,35 @@ TEST(KthookSimpleTest, CREATE_NAME(HandlesSimpleUsage)) { EXPECT_EQ(A::test_func(test_val), return_default); } -#ifdef KTHOOK_64 -TEST(KthookNakedTest, CREATE_NAME(HandlesSimpleUsage)) { - kthook::kthook_naked hook{ reinterpret_cast(&A::test_func) }; - hook.install(); - - hook.set_cb([](const kthook::kthook_naked& hook) { - auto& ctx = hook.get_context(); - -#ifdef KTHOOK_WIN64 - EXPECT_EQ(ctx.rcx, test_val); -#else - EXPECT_EQ(ctx.rdi, test_val); -#endif - }); - - A::test_func(test_val); -} -#else -class AT { -public: - NO_OPTIMIZE static int -#ifdef KTHOOK_32 - CTHISCALL -#endif - test_func(int value) { - SIZE_ENLARGER(); - return value; - } -}; - -class AF { -public: - NO_OPTIMIZE static int -#ifdef KTHOOK_32 - CFASTCALL -#endif - test_func(int value) { - SIZE_ENLARGER(); - return value; - } -}; - -TEST(KthookNakedTest, CREATE_NAME(HandlesSimpleUsageThiscall)) { - kthook::kthook_naked hook{reinterpret_cast(&A::test_func)}; +TEST(kthook_naked, thiscall_function) { + kthook::kthook_naked hook{reinterpret_cast(&AT::test_func)}; hook.install(); hook.set_cb([](const kthook::kthook_naked& hook) { auto& ctx = hook.get_context(); - EXPECT_EQ(ctx.ecx, test_val); + auto arg1 = ctx.IARG1; + EXPECT_EQ(arg1, test_val); }); AT::test_func(test_val); } -TEST(KthookNakedTest, CREATE_NAME(HandlesSimpleUsageFastcall)) { - kthook::kthook_naked hook{reinterpret_cast(&A::test_func)}; + +TEST(kthook_naked, fastcall_function) { + kthook::kthook_naked hook{reinterpret_cast(&AF::test_func)}; hook.install(); hook.set_cb([](const kthook::kthook_naked& hook) { auto& ctx = hook.get_context(); - EXPECT_EQ(ctx.ecx, test_val); + auto arg1 = ctx.IARG1; + EXPECT_EQ(arg1, test_val); }); AF::test_func(test_val); } -#endif -TEST(KthookSiginalTest, CREATE_NAME(HandlesSimpleUsage)) { +TEST(kthook_signal, function) { kthook::kthook_signal hook{&A::test_func}; { @@ -117,7 +92,9 @@ TEST(KthookSiginalTest, CREATE_NAME(HandlesSimpleUsage)) { { auto connection = - hook.before.scoped_connect([&](const auto& hook, int& value) { return std::make_optional(return_default); }); + hook.before.scoped_connect([&](const auto& hook, int& value) { + return std::make_optional(return_default); + }); EXPECT_EQ(A::test_func(test_val), return_default); } @@ -128,4 +105,4 @@ TEST(KthookSiginalTest, CREATE_NAME(HandlesSimpleUsage)) { EXPECT_EQ(A::test_func(test_val), return_default); } -} \ No newline at end of file +} diff --git a/tests/test_common.hpp b/tests/test_common.hpp index 704e335..051b4d9 100644 --- a/tests/test_common.hpp +++ b/tests/test_common.hpp @@ -1,3 +1,6 @@ +#include +#include + #if defined(_MSC_VER) #define NO_OPTIMIZE __declspec(noinline) #elif defined(__clang__) @@ -8,8 +11,6 @@ #error Unknown compiler #endif -#define CREATE_NAME(X) X - #define DECLARE_SIZE_ENLARGER() static volatile unsigned long long a = 5; #define SIZE_ENLARGER() \ switch (a) { \ @@ -23,3 +24,63 @@ a = 5; \ break; \ } + +#ifdef KTHOOK_32 +#define CCONV TEST_CCONV +#define TEST_THISCALL CTHISCALL +#define TEST_FASTCALL CFASTCALL +#ifdef IS_THISCALL +#define THISCALL_REPLACEMENT CSTDCALL +#else +#define THISCALL_REPLACEMENT TEST_CCONV +#endif +#else +#define THISCALL_REPLACEMENT +#define CCONV +#define TEST_THISCALL +#define TEST_FASTCALL +#endif + +#define GET_RANDOM_BYTE() [] { std::random_device dev;\ + std::mt19937 rng(dev());\ + std::uniform_int_distribution dist(std::numeric_limits::min(),std::numeric_limits::max());\ + return dist(rng);\ + }() + +#define GET_RANDOM_WORD() [] { std::random_device dev;\ + std::mt19937 rng(dev());\ + std::uniform_int_distribution dist(std::numeric_limits::min(),std::numeric_limits::max());\ + return dist(rng);\ + }() + +#define GET_RANDOM_DWORD() [] { std::random_device dev;\ + std::mt19937 rng(dev());\ + std::uniform_int_distribution dist(std::numeric_limits::min(),std::numeric_limits::max());\ + return dist(rng);\ + }() + +#define GET_RANDOM_QWORD() [] { std::random_device dev;\ + std::mt19937 rng(dev());\ + std::uniform_int_distribution dist(std::numeric_limits::min(),std::numeric_limits::max());\ + return dist(rng);\ + }() + +#define GET_RANDOM_INT() [] { std::random_device dev;\ + std::mt19937 rng(dev());\ + std::uniform_int_distribution dist(std::numeric_limits::min(),std::numeric_limits::max());\ + return dist(rng);\ + }() + +#if defined(KTHOOK_64_WIN) +#define IARG1 rcx +#define IARG2 rdx +#define IARG3 r8 +#define IARG4 r9 +#elif defined(KTHOOK_64_GCC) +#define IARG1 rdi +#define IARG2 rsi +#define IARG3 rdx +#define IARG4 rcx +#else +#define IARG1 ecx +#endif