From 5833197b036d3821597f9d412a10aeccf60f43b3 Mon Sep 17 00:00:00 2001
From: icex2 <djh.icex2@gmail.com>
Date: Sun, 25 Feb 2024 09:07:54 +0100
Subject: [PATCH] fix(hook): Add missing hook_table_revert impl

Allow hooks to cleanup when they are shut down.
---
 src/main/hook/table.c | 82 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/src/main/hook/table.c b/src/main/hook/table.c
index ed67b3b..067ea75 100644
--- a/src/main/hook/table.c
+++ b/src/main/hook/table.c
@@ -15,12 +15,21 @@ static const size_t apiset_prefix_len = sizeof(apiset_prefix) - 1;
 static void hook_table_apply_to_all(
     const char *depname, const struct hook_symbol *syms, size_t nsyms);
 
+static void hook_table_revert_to_all(
+    const char *depname, const struct hook_symbol *syms, size_t nsyms);
+
 static void hook_table_apply_to_iid(
     HMODULE target,
     const pe_iid_t *iid,
     const struct hook_symbol *syms,
     size_t nsyms);
 
+static void hook_table_revert_to_iid(
+    HMODULE target,
+    const pe_iid_t *iid,
+    const struct hook_symbol *syms,
+    size_t nsyms);
+
 static bool hook_table_match_module(
     HMODULE target, const char *iid_name, const char *depname);
 
@@ -44,6 +53,23 @@ static void hook_table_apply_to_all(
     }
 }
 
+static void hook_table_revert_to_all(
+    const char *depname, const struct hook_symbol *syms, size_t nsyms)
+{
+    const peb_dll_t *dll;
+    HMODULE pe;
+
+    for (dll = peb_dll_get_first(); dll != NULL; dll = peb_dll_get_next(dll)) {
+        pe = peb_dll_get_base(dll);
+
+        if (pe == NULL) {
+            continue; /* ?? Happens sometimes. */
+        }
+
+        hook_table_revert(pe, depname, syms, nsyms);
+    }
+}
+
 void hook_table_apply(
     HMODULE target,
     const char *depname,
@@ -73,6 +99,35 @@ void hook_table_apply(
     }
 }
 
+void hook_table_revert(
+    HMODULE target,
+    const char *depname,
+    const struct hook_symbol *syms,
+    size_t nsyms)
+{
+    const pe_iid_t *iid;
+    const char *iid_name;
+
+    assert(depname != NULL);
+    assert(syms != NULL || nsyms == 0);
+
+    if (target == NULL) {
+        /*  Call out, which will then call us back repeatedly. Awkward, but
+            viewed from the outside it's good for usability. */
+
+        hook_table_revert_to_all(depname, syms, nsyms);
+    } else {
+        for (iid = pe_iid_get_first(target); iid != NULL;
+             iid = pe_iid_get_next(target, iid)) {
+            iid_name = pe_iid_get_name(target, iid);
+
+            if (hook_table_match_module(target, iid_name, depname)) {
+                hook_table_revert_to_iid(target, iid, syms, nsyms);
+            }
+        }
+    }
+}
+
 static void hook_table_apply_to_iid(
     HMODULE target,
     const pe_iid_t *iid,
@@ -101,6 +156,33 @@ static void hook_table_apply_to_iid(
     }
 }
 
+static void hook_table_revert_to_iid(
+    HMODULE target,
+    const pe_iid_t *iid,
+    const struct hook_symbol *syms,
+    size_t nsyms)
+{
+    struct pe_iat_entry iate;
+    size_t i;
+    size_t j;
+    const struct hook_symbol *sym;
+
+    i = 0;
+
+    while (pe_iid_get_iat_entry(target, iid, i++, &iate) == S_OK) {
+        for (j = 0; j < nsyms; j++) {
+            sym = &syms[j];
+
+            if (hook_table_match_proc(&iate, sym)) {
+                // Only revert-able if the original pointer was stored previously
+                if (sym->link != NULL && *sym->link != NULL) {
+                    pe_patch(iate.ppointer, sym->link, sizeof(*sym->link));
+                } 
+            }
+        }
+    }
+}
+
 static bool hook_table_match_module(
     HMODULE target, const char *iid_name, const char *depname)
 {