Fix hook_inline() breaking after a single call

Committer note: I, mike, am a big dumb idiot. Thanks, Bill.

2 files changed, 20 insertions, 2 deletions

diff --git a/src/hook.c b/src/hook.c
index f625966..09dc403 100644
--- a/src/hook.c
+++ b/src/hook.c
@@ -1,5 +1,6 @@
 /*
  * Copyright © 2021 Michael Smith <mikesmiffy128@gmail.com>
+ * Copyright © 2022 Willian Henrique <wsimanbrazil@yahoo.com.br>
  *
  * Permission to use, copy, modify, and/or distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -58,10 +59,11 @@ void *hook_inline(void *func_, void *target) {
 		len += ud_insn_len(&udis);
 	}
 	// for simplicity, just bump alloc the trampoline. no need to free anyway
-	if (nexttrampoline - trampolines > len + 6) goto nospc;
+	if (nexttrampoline - trampolines > sizeof(trampolines) - len - 6) goto nospc;
 	uchar *trampoline = (uchar *)InterlockedExchangeAdd(
 			(volatile long *)&nexttrampoline, len + 6);
-	if (trampoline - trampolines > len + 6) { // avoid TOCTOU
+	// avoid TOCTOU
+	if (trampoline - trampolines > sizeof(trampolines) - len - 6) {
 nospc:	con_warn("hook_inline: out of trampoline space\n");
 		return 0;
 	}
diff --git a/test/hook.test.c b/test/hook.test.c
index 831fbf6..50e07a8 100644
--- a/test/hook.test.c
+++ b/test/hook.test.c
@@ -25,6 +25,12 @@ static int (*orig_some_function)(int, int);
 static int some_hook(int a, int b) {
 	return orig_some_function(a, b) + 5;
 }
+__attribute__((noinline))
+static int other_function(int a, int b) { return a - b; }
+static int (*orig_other_function)(int, int);
+static int other_hook(int a, int b) {
+	return orig_other_function(a, b) + 5;
+}
 
 TEST("Inline hooks should be able to wrap the original function", 0) {
 	orig_some_function = hook_inline(&some_function, &some_hook);
@@ -39,6 +45,16 @@ TEST("Inline hooks should be removable again", 0) {
 	return some_function(5, 5) == 10;
 }
 
+TEST("Multiple functions should be able to be inline hooked at once", 0) {
+	orig_some_function = hook_inline(&some_function, &some_hook);
+	if (!orig_some_function) return false;
+
+	orig_other_function = hook_inline(&other_function, &other_hook);
+	if (!orig_other_function) return false;
+
+	return other_function(5, 5) == 5;
+}
+
 #endif
 
 // vi: sw=4 ts=4 noet tw=80 cc=80