From e2b8151f6eca6c91b45e0d715b1533cde274eaa5 Mon Sep 17 00:00:00 2001 From: Willian Henrique Date: Sat, 12 Feb 2022 23:25:04 -0300 Subject: Fix hook_inline() breaking after a single call Committer note: I, mike, am a big dumb idiot. Thanks, Bill. --- src/hook.c | 6 ++++-- test/hook.test.c | 16 ++++++++++++++++ 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 + * Copyright © 2022 Willian Henrique * * 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 -- cgit v1.2.3