From 76f72852c8707987d39866084e6bf4e82a7aa304 Mon Sep 17 00:00:00 2001
From: Andrew Haberlandt <ahaberla@andrew.cmu.edu>
Date: Mon, 26 Aug 2024 21:48:38 +0000
Subject: [PATCH] Handle invalid hook addresses gracefully

Closes #37
---
 pyda_core/pyda_core_py.c |  9 ++++++++-
 tests/err_invalidhook.py | 18 ++++++++++++++++++
 tests/run_tests.py       |  8 ++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 tests/err_invalidhook.py

diff --git a/pyda_core/pyda_core_py.c b/pyda_core/pyda_core_py.c
index 9eda6e0..e43c9d8 100644
--- a/pyda_core/pyda_core_py.c
+++ b/pyda_core/pyda_core_py.c
@@ -559,9 +559,16 @@ PydaProcess_register_hook(PyObject *self, PyObject *args) {
 
 #ifdef PYDA_DYNAMORIO_CLIENT
     DEBUG_PRINTF("register_hook: %llx\n", addr);
-#endif // PYDA_DYNAMORIO_CLIENT
+    if (!dr_memory_is_readable((app_pc)addr, 1)) {
+        char buf[100];
+        snprintf(buf, sizeof(buf), "Hooked PC %" PRIxPTR " is invalid.", (uintptr_t)addr);
+        PyErr_SetString(PyExc_RuntimeError, buf);
+        return NULL;
+    }
+
     pyda_add_hook(p->main_thread->proc, addr, callback);
 
+#endif // PYDA_DYNAMORIO_CLIENT
     Py_INCREF(Py_None);
     return Py_None;
 }
diff --git a/tests/err_invalidhook.py b/tests/err_invalidhook.py
new file mode 100644
index 0000000..e110920
--- /dev/null
+++ b/tests/err_invalidhook.py
@@ -0,0 +1,18 @@
+from pyda import *
+from pwnlib.elf.elf import ELF
+from pwnlib.util.packing import u64
+import string
+import sys, time
+
+p = process()
+
+e = ELF(p.exe_path)
+e.address = p.maps[p.exe_path].base
+
+counter = 0
+def lib_hook(p):
+    global counter
+    counter += 1
+
+p.hook(0x1337133713371337, lib_hook)
+p.run()
\ No newline at end of file
diff --git a/tests/run_tests.py b/tests/run_tests.py
index 5fc2ede..b7a4d38 100644
--- a/tests/run_tests.py
+++ b/tests/run_tests.py
@@ -224,6 +224,14 @@ def no_warnings_or_errors(stdout: bytes, stderr: bytes) -> bool:
             lambda o, e: o.count(b"pass\n") == 1,
         ]
     )),
+
+    ("err_invalidhook", "simple.c", "err_invalidhook.py", RunOpts(), ExpectedResult(
+        retcode=0,
+        checkers=[
+            output_checker,
+            lambda o, e: e.count(b"RuntimeError: Hooked PC 1337133713371337 is invalid.") == 1,
+        ]
+    )),
 ]
 
 def main():