Skip to content

Commit 58aab31

Browse files
std-microblockstd-microblock
committed
feat: veh hook
1 parent c680d76 commit 58aab31

File tree

4 files changed

+182
-33
lines changed

4 files changed

+182
-33
lines changed

‎include/blook/hook.h‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -82,13 +82,15 @@ class InlineHook {
8282

8383
class VEHHookManager {
8484
public:
85-
struct VEHHookContext {};
85+
struct VEHHookContext {
86+
_EXCEPTION_POINTERS *exception_info;
87+
};
8688
using BreakpointCallback = std::function<void(VEHHookContext &ctx)>;
8789

8890
struct HardwareBreakpoint {
8991
void *address = nullptr;
9092
short dr_index = -1;
91-
int size = 0;
93+
int size = 1;
9294
HwBp::When when = HwBp::When::Executed;
9395
};
9496

@@ -105,8 +107,7 @@ class VEHHookManager {
105107
struct HardwareBreakpointInformation {
106108
HardwareBreakpoint bp;
107109
BreakpointCallback callback;
108-
void *trampoline_address = nullptr;
109-
size_t trampoline_size = 0;
110+
Trampoline trampoline;
110111
};
111112

112113
static VEHHookManager &instance() {
@@ -128,11 +129,11 @@ class VEHHookManager {
128129
VEHHookHandler add_breakpoint(PagefaultBreakpoint bp,
129130
BreakpointCallback callback);
130131
void remove_breakpoint(const VEHHookHandler &handler);
131-
132-
private:
132+
133133
std::array<std::optional<HardwareBreakpointInformation>, 4> hw_breakpoints;
134-
135134
void sync_hw_breakpoints();
135+
private:
136+
136137
VEHHookManager() {}
137138
};
138139
} // namespace blook

‎src/hook.cpp‎

Lines changed: 54 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -120,9 +120,47 @@ Trampoline Trampoline::make(Pointer pCode, size_t minByteSize,
120120
return t;
121121
}
122122

123+
static LONG blook_VectoredExceptionHandler(_EXCEPTION_POINTERS *ExceptionInfo) {
124+
auto &manager = blook::VEHHookManager::instance();
125+
auto code = ExceptionInfo->ExceptionRecord->ExceptionCode;
126+
auto address = ExceptionInfo->ExceptionRecord->ExceptionAddress;
127+
if (code == EXCEPTION_SINGLE_STEP) {
128+
for (const auto &bp : manager.hw_breakpoints) {
129+
if (bp.has_value() && bp->bp.address == address) {
130+
if (bp->callback) {
131+
size_t origRip = ExceptionInfo->ContextRecord->Rip;
132+
VEHHookManager::VEHHookContext ctx(ExceptionInfo);
133+
bp->callback(ctx);
134+
if (ExceptionInfo->ContextRecord->Rip == origRip) {
135+
// If the callback didn't change the RIP, jump to trampoline
136+
ExceptionInfo->ContextRecord->Rip =
137+
(size_t)bp->trampoline.pTrampoline.data();
138+
}
139+
}
140+
}
141+
}
142+
return EXCEPTION_CONTINUE_EXECUTION;
143+
} else if (code == EXCEPTION_BREAKPOINT) {
144+
// Handle software breakpoints here if needed
145+
} else if (code == EXCEPTION_ACCESS_VIOLATION) {
146+
// Handle pagefault breakpoints here if needed
147+
}
148+
149+
return EXCEPTION_CONTINUE_SEARCH;
150+
}
151+
152+
void ensureVectoredExceptionHandler() {
153+
static bool handler_installed = false;
154+
if (!handler_installed) {
155+
AddVectoredExceptionHandler(1, blook_VectoredExceptionHandler);
156+
handler_installed = true;
157+
}
158+
}
159+
123160
VEHHookManager::VEHHookHandler
124161
VEHHookManager::add_breakpoint(HardwareBreakpoint bp,
125162
BreakpointCallback callback) {
163+
ensureVectoredExceptionHandler();
126164
if (bp.dr_index == -1) {
127165
for (int i = 0; i < 4; ++i) {
128166
if (!hw_breakpoints[i].has_value()) {
@@ -144,12 +182,9 @@ VEHHookManager::add_breakpoint(HardwareBreakpoint bp,
144182
hw_breakpoints[bp.dr_index] = {
145183
.bp = bp,
146184
.callback = callback,
147-
.trampoline_address = nullptr,
148-
.trampoline_size = 0,
185+
.trampoline = Trampoline::make(bp.address, bp.size, true),
149186
};
150187

151-
// allocate trampoline
152-
153188
sync_hw_breakpoints();
154189
return HardwareBreakpointHandler{bp.dr_index};
155190
}
@@ -178,17 +213,23 @@ void VEHHookManager::remove_breakpoint(const VEHHookHandler &handler) {
178213
void VEHHookManager::sync_hw_breakpoints() {
179214
auto threads = Process::self()->threads();
180215
for (const auto &thread : threads) {
181-
for (auto &bp : hw_breakpoints) {
216+
for (int i = 0; i < 4; ++i) {
217+
auto &bp = hw_breakpoints[i];
182218
if (bp.has_value()) {
183-
HwBp::Set(bp->bp.address, bp->bp.size, bp->bp.when, bp->bp.dr_index,
184-
OpenThread(THREAD_GET_CONTEXT | THREAD_SET_CONTEXT |
185-
THREAD_SUSPEND_RESUME,
186-
false, thread.id));
219+
if (auto res = HwBp::Set(
220+
bp->bp.address, bp->bp.size, bp->bp.when, bp->bp.dr_index,
221+
OpenThread(THREAD_GET_CONTEXT | THREAD_SET_CONTEXT |
222+
THREAD_SUSPEND_RESUME,
223+
false, thread.id));
224+
res.m_error != HwBp::Result::Success) {
225+
std::println(
226+
"Failed to set hardware breakpoint on DR{} for thread ID {}: {}",
227+
bp->bp.dr_index, thread.id, (int)res.m_error);
228+
}
187229
} else {
188-
HwBp::Remove(bp->bp.dr_index,
189-
OpenThread(THREAD_GET_CONTEXT | THREAD_SET_CONTEXT |
190-
THREAD_SUSPEND_RESUME,
191-
false, thread.id));
230+
HwBp::Remove(i, OpenThread(THREAD_GET_CONTEXT | THREAD_SET_CONTEXT |
231+
THREAD_SUSPEND_RESUME,
232+
false, thread.id));
192233
}
193234
}
194235
}

‎src/tests/test_windows.cpp‎

Lines changed: 119 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -254,21 +254,128 @@ TEST(BlookModuleTests, ExportParsing) {
254254
EXPECT_EQ(exported_func->data<void *>(), (void *)&f_test_exports);
255255
}
256256

257-
class VEHHookTest : public ::testing::Test {
258-
protected:
259-
std::shared_ptr<blook::InlineHook> hook_AplusB;
260-
std::shared_ptr<blook::InlineHook> hook_GetTickCount;
257+
#pragma optimize("", off)
258+
int veh_test_target_func(int a, int b) { return a * b + 42; }
259+
int veh_test_target_func2(int a, int b) { return a * b + 423; }
260+
int veh_test_target_func3(int a, int b) { return a * b + 412; }
261+
int veh_test_target_func4(int a, int b) { return a * b + 426; }
262+
#pragma optimize("", on)
261263

262-
void TearDown() override {
264+
TEST(VEHHookTest, HookFunction) {
265+
bool called = false;
266+
auto handler = blook::VEHHookManager::instance().add_breakpoint(
267+
blook::VEHHookManager::HardwareBreakpoint{
268+
.address = (void *)veh_test_target_func},
269+
[&](blook::VEHHookManager::VEHHookContext &ctx) { called = true; });
263270

264-
if (hook_AplusB && hook_AplusB->is_installed()) {
265-
hook_AplusB->uninstall();
266-
}
267-
if (hook_GetTickCount && hook_GetTickCount->is_installed()) {
268-
hook_GetTickCount->uninstall();
271+
ASSERT_TRUE(
272+
std::holds_alternative<blook::VEHHookManager::HardwareBreakpointHandler>(
273+
handler));
274+
275+
veh_test_target_func(3, 5);
276+
277+
ASSERT_TRUE(called);
278+
279+
blook::VEHHookManager::instance().remove_breakpoint(handler);
280+
281+
auto result = veh_test_target_func(2, 4);
282+
EXPECT_EQ(result, 2 * 4 + 42);
283+
}
284+
285+
TEST(VEHHookTest, HookFunctionInAnotherThread) {
286+
bool called = false;
287+
auto handler = blook::VEHHookManager::instance().add_breakpoint(
288+
blook::VEHHookManager::HardwareBreakpoint{
289+
.address = (void *)veh_test_target_func},
290+
[&](blook::VEHHookManager::VEHHookContext &ctx) {
291+
called = true;
292+
ExitThread(0);
293+
});
294+
295+
ASSERT_TRUE(
296+
std::holds_alternative<blook::VEHHookManager::HardwareBreakpointHandler>(
297+
handler));
298+
299+
std::thread t([]() {
300+
while (true) {
301+
veh_test_target_func(7, 8);
269302
}
270-
}
271-
};
303+
});
304+
blook::VEHHookManager::instance().sync_hw_breakpoints();
305+
t.join();
306+
307+
ASSERT_TRUE(called);
308+
309+
blook::VEHHookManager::instance().remove_breakpoint(handler);
310+
}
311+
312+
TEST(VEHHookTest, HookMultipleFunctions) {
313+
bool called1 = false;
314+
bool called2 = false;
315+
bool called3 = false;
316+
bool called4 = false;
317+
318+
auto handler1 = blook::VEHHookManager::instance().add_breakpoint(
319+
blook::VEHHookManager::HardwareBreakpoint{
320+
.address = (void *)veh_test_target_func},
321+
[&](blook::VEHHookManager::VEHHookContext &ctx) { called1 = true; });
322+
323+
auto handler2 = blook::VEHHookManager::instance().add_breakpoint(
324+
blook::VEHHookManager::HardwareBreakpoint{
325+
.address = (void *)veh_test_target_func2},
326+
[&](blook::VEHHookManager::VEHHookContext &ctx) { called2 = true; });
327+
328+
auto handler3 = blook::VEHHookManager::instance().add_breakpoint(
329+
blook::VEHHookManager::HardwareBreakpoint{
330+
.address = (void *)veh_test_target_func3},
331+
[&](blook::VEHHookManager::VEHHookContext &ctx) { called3 = true; });
332+
333+
auto handler4 = blook::VEHHookManager::instance().add_breakpoint(
334+
blook::VEHHookManager::HardwareBreakpoint{
335+
.address = (void *)veh_test_target_func4},
336+
[&](blook::VEHHookManager::VEHHookContext &ctx) { called4 = true; });
337+
338+
ASSERT_TRUE(
339+
std::holds_alternative<blook::VEHHookManager::HardwareBreakpointHandler>(
340+
handler1));
341+
ASSERT_TRUE(
342+
std::holds_alternative<blook::VEHHookManager::HardwareBreakpointHandler>(
343+
handler2));
344+
ASSERT_TRUE(
345+
std::holds_alternative<blook::VEHHookManager::HardwareBreakpointHandler>(
346+
handler3));
347+
ASSERT_TRUE(
348+
std::holds_alternative<blook::VEHHookManager::HardwareBreakpointHandler>(
349+
handler4));
350+
351+
veh_test_target_func(3, 5);
352+
veh_test_target_func2(6, 7);
353+
veh_test_target_func3(8, 9);
354+
veh_test_target_func4(10, 11);
355+
356+
ASSERT_TRUE(called1);
357+
ASSERT_TRUE(called2);
358+
ASSERT_TRUE(called3);
359+
ASSERT_TRUE(called4);
360+
361+
blook::VEHHookManager::instance().remove_breakpoint(handler1);
362+
blook::VEHHookManager::instance().remove_breakpoint(handler2);
363+
blook::VEHHookManager::instance().remove_breakpoint(handler3);
364+
blook::VEHHookManager::instance().remove_breakpoint(handler4);
365+
called1 = called2 = called3 = called4 = false;
366+
veh_test_target_func(2, 4);
367+
veh_test_target_func2(3, 5);
368+
veh_test_target_func3(4, 6);
369+
veh_test_target_func4(5, 7);
370+
EXPECT_EQ(veh_test_target_func(2, 4), 2 * 4 + 42);
371+
EXPECT_EQ(veh_test_target_func2(3, 5), 3 * 5 + 423);
372+
EXPECT_EQ(veh_test_target_func3(4, 6), 4 * 6 + 412);
373+
EXPECT_EQ(veh_test_target_func4(5, 7), 5 * 7 + 426);
374+
ASSERT_FALSE(called1);
375+
ASSERT_FALSE(called2);
376+
ASSERT_FALSE(called3);
377+
ASSERT_FALSE(called4);
378+
}
272379

273380
int main(int argc, char **argv) {
274381
::testing::InitGoogleTest(&argc, argv);

‎xmake.lua‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ add_requires("zasm 2025.03.02", "gtest")
1212

1313
target("blook")
1414
set_kind("static")
15-
add_files("src/*.cpp")
15+
add_files("src/*.cpp", "src/**/*.cc")
1616
add_defines("WIN32_LEAN_AND_MEAN", "NOMINMAX")
1717
if is_plat("windows") then
1818
add_files("src/platform/windows/*.cpp")

0 commit comments

Comments
 (0)