Implement handler removing when no hooks are active

geode-sdk · Jan 19, 2025 · dc14d4c · dc14d4c
1 parent 315bf46
commit dc14d4c
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/loader/src/loader/HookImpl.cpp b/loader/src/loader/HookImpl.cpp
@@ -82,7 +82,7 @@ Result<> Hook::Impl::enable() {
 Result<> Hook::Impl::disable() {
     if (!m_enabled)
         return Ok();
-    GEODE_UNWRAP_INTO(auto handler, LoaderImpl::get()->getHandler(m_address));
+    GEODE_UNWRAP_INTO(auto handler, LoaderImpl::get()->getOrRemoveHandler(m_address));
     tulip::hook::removeHook(handler, m_handle);
     m_enabled = false;
     log::debug("Disabled {} hook", m_displayName);

diff --git a/loader/src/loader/LoaderImpl.cpp b/loader/src/loader/LoaderImpl.cpp
@@ -948,15 +948,29 @@ Result<tulip::hook::HandlerHandle> Loader::Impl::getHandler(void* address) {
     if (!m_handlerHandles.count(address)) {
         return Err("Handler does not exist at address");
     }
-    return Ok(m_handlerHandles[address]);
+    return Ok(m_handlerHandles[address].first);
 }
 
 Result<tulip::hook::HandlerHandle> Loader::Impl::getOrCreateHandler(void* address, tulip::hook::HandlerMetadata const& metadata) {
     if (m_handlerHandles.count(address)) {
-        return Ok(m_handlerHandles[address]);
+        m_handlerHandles[address].second++;
+        return Ok(m_handlerHandles[address].first);
     }
     GEODE_UNWRAP_INTO(auto handle, tulip::hook::createHandler(address, metadata));
-    m_handlerHandles[address] = handle;
+    m_handlerHandles[address].first = handle;
+    m_handlerHandles[address].second = 1;
+    return Ok(handle);
+}
+
+Result<tulip::hook::HandlerHandle> Loader::Impl::getOrRemoveHandler(void* address) {
+    if (!m_handlerHandles.count(address)) {
+        return Err("Handler does not exist at address");
+    }
+    auto handle = m_handlerHandles[address].first;
+    if (m_handlerHandles[address].second == 1) {
+        GEODE_UNWRAP(tulip::hook::removeHandler(handle));
+    }
+    m_handlerHandles[address].second--;
     return Ok(handle);
 }
 

diff --git a/loader/src/loader/LoaderImpl.hpp b/loader/src/loader/LoaderImpl.hpp
@@ -68,10 +68,11 @@ namespace geode {
         Mod* takeNextMod();
         void releaseNextMod();
 
-        std::unordered_map<void*, tulip::hook::HandlerHandle> m_handlerHandles;
+        std::unordered_map<void*, std::pair<tulip::hook::HandlerHandle, size_t>> m_handlerHandles;
 
         Result<tulip::hook::HandlerHandle> getHandler(void* address);
         Result<tulip::hook::HandlerHandle> getOrCreateHandler(void* address, tulip::hook::HandlerMetadata const& metadata);
+        Result<tulip::hook::HandlerHandle> getOrRemoveHandler(void* address);
 
         bool loadHooks();