Fixup hooks impl, preserve %rax

2 files changed, 47 insertions, 23 deletions

diff --git a/example/examplePlugin/main.cpp b/example/examplePlugin/main.cpp
index cd8ddf45..050d403f 100644
--- a/example/examplePlugin/main.cpp
+++ b/example/examplePlugin/main.cpp
@@ -12,8 +12,10 @@
 
 // Methods
 inline std::unique_ptr<CHyprCustomLayout> g_pCustomLayout;
-inline CFunctionHook*                     g_pFocusHook = nullptr;
+inline CFunctionHook*                     g_pFocusHook  = nullptr;
+inline CFunctionHook*                     g_pMotionHook = nullptr;
 typedef void (*origFocusWindow)(void*, CWindow*, wlr_surface*);
+typedef void (*origMotion)(wlr_seat*, uint32_t, double, double);
 
 // Do NOT change this function.
 APICALL EXPORT std::string PLUGIN_API_VERSION() {
@@ -35,12 +37,15 @@ static void onNewWindow(void* self, std::any data) {
 }
 
 void hkFocusWindow(void* thisptr, CWindow* pWindow, wlr_surface* pSurface) {
-
-    // stuf
-
+    HyprlandAPI::addNotification(PHANDLE, getFormat("FocusWindow with %lx %lx", pWindow, pSurface), CColor{0.f, 1.f, 1.f, 1.f}, 5000);
     (*(origFocusWindow)g_pFocusHook->m_pOriginal)(thisptr, pWindow, pSurface);
 }
 
+void hkNotifyMotion(wlr_seat* wlr_seat, uint32_t time_msec, double sx, double sy) {
+    HyprlandAPI::addNotification(PHANDLE, getFormat("NotifyMotion with %lf %lf", sx, sy), CColor{0.f, 1.f, 1.f, 1.f}, 5000);
+    (*(origMotion)g_pMotionHook->m_pOriginal)(wlr_seat, time_msec, sx, sy);
+}
+
 APICALL EXPORT PLUGIN_DESCRIPTION_INFO PLUGIN_INIT(HANDLE handle) {
     PHANDLE = handle;
 
@@ -55,8 +60,10 @@ APICALL EXPORT PLUGIN_DESCRIPTION_INFO PLUGIN_INIT(HANDLE handle) {
 
     HyprlandAPI::addConfigValue(PHANDLE, "plugin:example:border_color", SConfigValue{.intValue = configStringToInt("rgb(44ee44)")});
 
-    g_pFocusHook = HyprlandAPI::createFunctionHook(PHANDLE, (void*)&CCompositor::focusWindow, (void*)&hkFocusWindow);
+    g_pFocusHook  = HyprlandAPI::createFunctionHook(PHANDLE, (void*)&CCompositor::focusWindow, (void*)&hkFocusWindow);
+    g_pMotionHook = HyprlandAPI::createFunctionHook(PHANDLE, (void*)&wlr_seat_pointer_notify_motion, (void*)&hkNotifyMotion);
     g_pFocusHook->hook();
+    g_pMotionHook->hook();
 
     HyprlandAPI::reloadConfig();
 
diff --git a/src/plugins/HookSystem.cpp b/src/plugins/HookSystem.cpp
index cbcea4a4..a8ddd90e 100644
--- a/src/plugins/HookSystem.cpp
+++ b/src/plugins/HookSystem.cpp
@@ -19,29 +19,34 @@ CFunctionHook::~CFunctionHook() {
     }
 }
 
-size_t probeMinimumJumpSize(void* start, size_t min) {
+size_t getInstructionLenAt(void* start) {
     ud_t udis;
 
     ud_init(&udis);
     ud_set_mode(&udis, 64);
+    ud_set_syntax(&udis, UD_SYN_INTEL);
 
-    size_t size = 0;
-
+    size_t curOffset = 1;
+    size_t insSize   = 0;
     while (true) {
-        size_t offset = 1;
-        while (true) {
+        ud_set_input_buffer(&udis, (uint8_t*)start, curOffset);
+        insSize = ud_disassemble(&udis);
+        if (insSize != curOffset)
+            break;
+        curOffset++;
+    }
 
-            ud_set_input_buffer(&udis, (uint8_t*)(start + size), offset);
-            if (offset != ud_disassemble(&udis))
-                break;
+    return insSize;
+}
 
-            offset++;
-        }
+size_t probeMinimumJumpSize(void* start, size_t min) {
 
-        size += offset;
+    size_t size = 0;
 
-        if (size > min)
-            break;
+    while (size <= min) {
+        // find info about this instruction
+        size_t insLen = getInstructionLenAt(start + size);
+        size += insLen;
     }
 
     return size;
@@ -56,24 +61,36 @@ bool CFunctionHook::hook() {
 
     // movabs $0,%rax | jmpq *%rax
     static constexpr uint8_t ABSOLUTE_JMP_ADDRESS[] = {0x48, 0xB8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xE0};
+    // pushq %rax
+    static constexpr uint8_t PUSH_RAX[] = {0x50};
+    // popq %rax
+    static constexpr uint8_t POP_RAX[] = {0x58};
+    // nop
+    static constexpr uint8_t NOP = 0x90;
 
     // get minimum size to overwrite
-    const auto HOOKSIZE = probeMinimumJumpSize(m_pSource, sizeof(ABSOLUTE_JMP_ADDRESS));
+    const auto HOOKSIZE = probeMinimumJumpSize(m_pSource, sizeof(ABSOLUTE_JMP_ADDRESS) + sizeof(PUSH_RAX) + sizeof(POP_RAX));
 
     // alloc trampoline
-    m_pTrampolineAddr = mmap(NULL, sizeof(ABSOLUTE_JMP_ADDRESS) + HOOKSIZE, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    m_pTrampolineAddr = mmap(NULL, sizeof(ABSOLUTE_JMP_ADDRESS) + HOOKSIZE + sizeof(PUSH_RAX), PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
 
     // populate trampoline
-    memcpy(m_pTrampolineAddr, m_pSource, HOOKSIZE);                                           // first, original func bytes
-    memcpy(m_pTrampolineAddr + HOOKSIZE, ABSOLUTE_JMP_ADDRESS, sizeof(ABSOLUTE_JMP_ADDRESS)); // then, our jump back
+    memcpy(m_pTrampolineAddr, m_pSource, HOOKSIZE);                                                              // first, original func bytes
+    memcpy(m_pTrampolineAddr + HOOKSIZE, PUSH_RAX, sizeof(PUSH_RAX));                                            // then, pushq %rax
+    memcpy(m_pTrampolineAddr + HOOKSIZE + sizeof(PUSH_RAX), ABSOLUTE_JMP_ADDRESS, sizeof(ABSOLUTE_JMP_ADDRESS)); // then, jump to source
 
     // fixup trampoline addr
-    *(uint64_t*)(m_pTrampolineAddr + HOOKSIZE + 2) = (uint64_t)(m_pSource + HOOKSIZE);
+    *(uint64_t*)(m_pTrampolineAddr + HOOKSIZE + 2 + sizeof(PUSH_RAX)) = (uint64_t)(m_pSource + sizeof(ABSOLUTE_JMP_ADDRESS));
 
     // make jump to hk
     mprotect(m_pSource - ((uint64_t)m_pSource) % sysconf(_SC_PAGE_SIZE), sysconf(_SC_PAGE_SIZE), PROT_READ | PROT_WRITE | PROT_EXEC);
     memcpy(m_pSource, ABSOLUTE_JMP_ADDRESS, sizeof(ABSOLUTE_JMP_ADDRESS));
 
+    // make popq %rax and NOP all remaining
+    memcpy(m_pSource + sizeof(ABSOLUTE_JMP_ADDRESS), POP_RAX, sizeof(POP_RAX));
+    size_t currentOp = sizeof(ABSOLUTE_JMP_ADDRESS) + sizeof(POP_RAX);
+    memset(m_pSource + currentOp, NOP, HOOKSIZE - currentOp);
+
     // fixup jump addr
     *(uint64_t*)(m_pSource + 2) = (uint64_t)(m_pDestination);