feat: software breakpoint hook

std-microblock · std-microblock · commit 9ca4500cbe39 · 2025-07-26T20:23:03.000+08:00
diff --git a/include/blook/hook.h b/include/blook/hook.h
@@ -6,6 +6,8 @@
 #include "utils.h"
 #include "zasm/x86/assembler.hpp"
 #include <optional>
+#include <map>
+#include <Windows.h>
 
 namespace blook {
 
@@ -96,12 +98,10 @@ class VEHHookManager {
 
   struct SoftwareBreakpoint {
     void *address = nullptr;
-    std::vector<uint8_t> original_bytes;
   };
 
   struct PagefaultBreakpoint {
     void *address = nullptr;
-    int32_t origin_protection = 0;
   };
 
   struct HardwareBreakpointInformation {
@@ -110,6 +110,18 @@ class VEHHookManager {
     Trampoline trampoline;
   };
 
+  struct SoftwareBreakpointInformation {
+    SoftwareBreakpoint bp;
+    BreakpointCallback callback;
+    std::vector<uint8_t> original_bytes;
+  };
+
+  struct PagefaultBreakpointInformation {
+    PagefaultBreakpoint bp;
+    BreakpointCallback callback;
+    int32_t origin_protection = 0;
+  };
+
   static VEHHookManager &instance() {
     static VEHHookManager instance;
     return instance;
@@ -119,8 +131,16 @@ class VEHHookManager {
     short dr_index = -1;
   };
 
+  struct SoftwareBreakpointHandler {
+    void *address = nullptr;
+  };
+
+  struct PagefaultBreakpointHandler {
+    void *address = nullptr;
+  };
+
   using VEHHookHandler =
-      std::variant<std::monostate, HardwareBreakpointHandler>;
+      std::variant<std::monostate, HardwareBreakpointHandler, SoftwareBreakpointHandler, PagefaultBreakpointHandler>;
 
   VEHHookHandler add_breakpoint(HardwareBreakpoint bp,
                                 BreakpointCallback callback);
@@ -131,6 +151,10 @@ class VEHHookManager {
   void remove_breakpoint(const VEHHookHandler &handler);
   
   std::array<std::optional<HardwareBreakpointInformation>, 4> hw_breakpoints;
+  std::map<void *, SoftwareBreakpointInformation> sw_breakpoints;
+  std::map<void *, PagefaultBreakpointInformation> pf_breakpoints;
+  std::map<DWORD, void *> thread_bp_in_progress;
+
   void sync_hw_breakpoints();
 private:
 
diff --git a/src/hook.cpp b/src/hook.cpp
@@ -121,34 +121,77 @@ Trampoline Trampoline::make(Pointer pCode, size_t minByteSize,
 }
 
 static LONG blook_VectoredExceptionHandler(_EXCEPTION_POINTERS *ExceptionInfo) {
-  auto &manager = blook::VEHHookManager::instance();
-  auto code = ExceptionInfo->ExceptionRecord->ExceptionCode;
-  auto address = ExceptionInfo->ExceptionRecord->ExceptionAddress;
-  if (code == EXCEPTION_SINGLE_STEP) {
-    for (const auto &bp : manager.hw_breakpoints) {
-      if (bp.has_value() && bp->bp.address == address) {
-        if (bp->callback) {
-          size_t origRip = ExceptionInfo->ContextRecord->Rip;
-          VEHHookManager::VEHHookContext ctx(ExceptionInfo);
-          bp->callback(ctx);
-          if (ExceptionInfo->ContextRecord->Rip == origRip) {
-            // If the callback didn't change the RIP, jump to trampoline
-            ExceptionInfo->ContextRecord->Rip =
-                (size_t)bp->trampoline.pTrampoline.data();
-          }
+    auto &manager = blook::VEHHookManager::instance();
+    auto code = ExceptionInfo->ExceptionRecord->ExceptionCode;
+    auto address = ExceptionInfo->ExceptionRecord->ExceptionAddress;
+    auto thread_id = GetCurrentThreadId();
+
+    if (code == EXCEPTION_SINGLE_STEP) {
+        // Check if this is a re-enable step for a SW/Page-Guard BP
+        if (manager.thread_bp_in_progress.contains(thread_id)) {
+            void *bp_address = manager.thread_bp_in_progress.at(thread_id);
+            
+            if (manager.sw_breakpoints.contains(bp_address)) {
+                uint8_t int3 = 0xCC;
+                Pointer(bp_address).write(nullptr, std::span(&int3, 1));
+            } else if (manager.pf_breakpoints.contains(bp_address)) {
+                auto &bp = manager.pf_breakpoints.at(bp_address);
+                DWORD old_protect;
+                VirtualProtect(bp_address, 1, bp.origin_protection | PAGE_GUARD, &old_protect);
+            }
+            manager.thread_bp_in_progress.erase(thread_id);
+            return EXCEPTION_CONTINUE_EXECUTION;
         }
-      }
+
+        // If not a re-enable step, check for a hardware breakpoint
+        for (auto &bp_opt : manager.hw_breakpoints) {
+            if (bp_opt.has_value() && bp_opt->bp.address == address) {
+                if (bp_opt->callback) {
+                    size_t origRip = ExceptionInfo->ContextRecord->Rip;
+                    VEHHookManager::VEHHookContext ctx(ExceptionInfo);
+                    bp_opt->callback(ctx);
+                    if (ExceptionInfo->ContextRecord->Rip == origRip) {
+                        ExceptionInfo->ContextRecord->Rip =
+                            (size_t)bp_opt->trampoline.pTrampoline.data();
+                    }
+                }
+                return EXCEPTION_CONTINUE_EXECUTION;
+            }
+        }
+
+    } else if (code == EXCEPTION_BREAKPOINT) {
+        if (manager.sw_breakpoints.contains(address)) {
+            auto &bp = manager.sw_breakpoints.at(address);
+            Pointer(address).write(nullptr, bp.original_bytes);
+            ExceptionInfo->ContextRecord->Rip = (DWORD_PTR)address;
+
+            if (bp.callback) {
+                VEHHookManager::VEHHookContext ctx(ExceptionInfo);
+                bp.callback(ctx);
+            }
+            
+            ExceptionInfo->ContextRecord->EFlags |= (1 << 8);
+            manager.thread_bp_in_progress[thread_id] = address;
+            return EXCEPTION_CONTINUE_EXECUTION;
+        }
+    } else if (code == EXCEPTION_GUARD_PAGE) {
+        if (manager.pf_breakpoints.contains(address)) {
+            auto &bp = manager.pf_breakpoints.at(address);
+            if (bp.callback) {
+                VEHHookManager::VEHHookContext ctx(ExceptionInfo);
+                bp.callback(ctx);
+            }
+            manager.thread_bp_in_progress[thread_id] = address;
+            ExceptionInfo->ContextRecord->EFlags |= (1 << 8);
+            return EXCEPTION_CONTINUE_EXECUTION;
+        }
+        return EXCEPTION_CONTINUE_SEARCH;
     }
-    return EXCEPTION_CONTINUE_EXECUTION;
-  } else if (code == EXCEPTION_BREAKPOINT) {
-    // Handle software breakpoints here if needed
-  } else if (code == EXCEPTION_ACCESS_VIOLATION) {
-    // Handle pagefault breakpoints here if needed
-  }
 
-  return EXCEPTION_CONTINUE_SEARCH;
+    return EXCEPTION_CONTINUE_SEARCH;
 }
 
+
 void ensureVectoredExceptionHandler() {
   static bool handler_installed = false;
   if (!handler_installed) {
@@ -191,12 +234,54 @@ VEHHookManager::add_breakpoint(HardwareBreakpoint bp,
 VEHHookManager::VEHHookHandler
 VEHHookManager::add_breakpoint(SoftwareBreakpoint bp,
                                BreakpointCallback callback) {
-  throw std::runtime_error("Software breakpoints are not supported yet.");
+  ensureVectoredExceptionHandler();
+  if (sw_breakpoints.contains(bp.address)) {
+    throw std::runtime_error("Software breakpoint already exists at this address.");
+  }
+
+  SoftwareBreakpointInformation info;
+  info.bp = bp;
+  info.callback = callback;
+  
+  auto original_byte = Pointer(bp.address).try_read<uint8_t>();
+  if (!original_byte) {
+      throw std::runtime_error("Failed to read original byte for software breakpoint.");
+  }
+  info.original_bytes.push_back(*original_byte);
+
+  uint8_t int3 = 0xCC;
+  if (!Pointer(bp.address).write(nullptr, std::span(&int3, 1))) {
+      throw std::runtime_error("Failed to write INT3 for software breakpoint.");
+  }
+
+  sw_breakpoints[bp.address] = std::move(info);
+  return SoftwareBreakpointHandler{bp.address};
 }
 VEHHookManager::VEHHookHandler
 VEHHookManager::add_breakpoint(PagefaultBreakpoint bp,
                                BreakpointCallback callback) {
-  throw std::runtime_error("Pagefault breakpoints are not supported yet.");
+    ensureVectoredExceptionHandler();
+    if (pf_breakpoints.contains(bp.address)) {
+        throw std::runtime_error("Pagefault breakpoint already exists at this address.");
+    }
+
+    PagefaultBreakpointInformation info;
+    info.bp = bp;
+    info.callback = callback;
+
+    MEMORY_BASIC_INFORMATION mbi;
+    if (VirtualQuery(bp.address, &mbi, sizeof(mbi)) == 0) {
+        throw std::runtime_error("VirtualQuery failed for pagefault breakpoint.");
+    }
+    info.origin_protection = mbi.Protect;
+
+    DWORD old_protect;
+    if (!VirtualProtect(bp.address, 1, info.origin_protection | PAGE_GUARD, &old_protect)) {
+        throw std::runtime_error("VirtualProtect failed for pagefault breakpoint.");
+    }
+
+    pf_breakpoints[bp.address] = std::move(info);
+    return PagefaultBreakpointHandler{bp.address};
 }
 void VEHHookManager::remove_breakpoint(const VEHHookHandler &handler) {
   if (auto _hwbp = std::get_if<HardwareBreakpointHandler>(&handler)) {
@@ -206,7 +291,23 @@ void VEHHookManager::remove_breakpoint(const VEHHookHandler &handler) {
     }
     hw_breakpoints[_hwbp->dr_index].reset();
     sync_hw_breakpoints();
-  } else {
+  } else if (auto _swbp = std::get_if<SoftwareBreakpointHandler>(&handler)) {
+    if (!sw_breakpoints.contains(_swbp->address)) {
+        return;
+    }
+    auto& bp_info = sw_breakpoints.at(_swbp->address);
+    Pointer(_swbp->address).write(nullptr, bp_info.original_bytes);
+    sw_breakpoints.erase(_swbp->address);
+  } else if (auto _pfbp = std::get_if<PagefaultBreakpointHandler>(&handler)) {
+      if (!pf_breakpoints.contains(_pfbp->address)) {
+          return;
+      }
+      auto& bp_info = pf_breakpoints.at(_pfbp->address);
+      DWORD old_protect;
+      VirtualProtect(_pfbp->address, 1, bp_info.origin_protection, &old_protect);
+      pf_breakpoints.erase(_pfbp->address);
+  }
+  else {
     throw std::runtime_error("Unsupported breakpoint type.");
   }
 }
diff --git a/src/tests/test_windows.cpp b/src/tests/test_windows.cpp
@@ -377,6 +377,52 @@ TEST(VEHHookTest, HookMultipleFunctions) {
   ASSERT_FALSE(called4);
 }
 
+TEST(VEHHookTest, SoftwareBreakpoint) {
+    bool called = false;
+    auto handler = blook::VEHHookManager::instance().add_breakpoint(
+        blook::VEHHookManager::SoftwareBreakpoint{
+            .address = (void *)veh_test_target_func},
+        [&](blook::VEHHookManager::VEHHookContext &ctx) { called = true; });
+
+    ASSERT_TRUE(
+        std::holds_alternative<blook::VEHHookManager::SoftwareBreakpointHandler>(
+            handler));
+
+    veh_test_target_func(3, 5);
+
+    ASSERT_TRUE(called);
+
+    blook::VEHHookManager::instance().remove_breakpoint(handler);
+
+    called = false;
+    auto result = veh_test_target_func(2, 4);
+    EXPECT_EQ(result, 2 * 4 + 42);
+    ASSERT_FALSE(called);
+}
+
+// TEST(VEHHookTest, PagefaultBreakpoint) {
+//     bool called = false;
+//     auto handler = blook::VEHHookManager::instance().add_breakpoint(
+//         blook::VEHHookManager::PagefaultBreakpoint{
+//             .address = (void *)veh_test_target_func},
+//         [&](blook::VEHHookManager::VEHHookContext &ctx) { called = true; });
+
+//     ASSERT_TRUE(
+//         std::holds_alternative<blook::VEHHookManager::PagefaultBreakpointHandler>(
+//             handler));
+
+//     veh_test_target_func(3, 5);
+
+//     ASSERT_TRUE(called);
+
+//     blook::VEHHookManager::instance().remove_breakpoint(handler);
+
+//     called = false;
+//     auto result = veh_test_target_func(2, 4);
+//     EXPECT_EQ(result, 2 * 4 + 42);
+//     ASSERT_FALSE(called);
+// }
+
 int main(int argc, char **argv) {
   ::testing::InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();
