From ecc00216d640bb69015d7b683b33d8db87b58083 Mon Sep 17 00:00:00 2001
From: 66hh <49398720+66hh@users.noreply.github.com>
Date: Sat, 27 Dec 2025 13:51:10 +0800
Subject: [PATCH] Support ProcessInstrumentationCallback

---
 src/common/platform/process.hpp             |  6 ++++++
 src/windows-emulator/process_context.cpp    |  3 +++
 src/windows-emulator/process_context.hpp    |  1 +
 src/windows-emulator/syscall_dispatcher.cpp | 13 +++++++++++++
 src/windows-emulator/syscalls/process.cpp   | 17 +++++++++++++++++
 5 files changed, 40 insertions(+)

diff --git a/src/common/platform/process.hpp b/src/common/platform/process.hpp
index a77be9b9..f91e41cf 100644
--- a/src/common/platform/process.hpp
+++ b/src/common/platform/process.hpp
@@ -1124,4 +1124,10 @@ struct PROCESS_PRIORITY_CLASS
     UCHAR PriorityClass;
 };
 
+struct PROCESS_INSTRUMENTATION_CALLBACK_INFORMATION
+{
+    ULONG Version;
+    ULONG Reserved;
+    PVOID Callback;
+};
 // NOLINTEND(modernize-use-using,cppcoreguidelines-avoid-c-arrays,hicpp-avoid-c-arrays,modernize-avoid-c-arrays)
diff --git a/src/windows-emulator/process_context.cpp b/src/windows-emulator/process_context.cpp
index c48c155c..ea8cfee3 100644
--- a/src/windows-emulator/process_context.cpp
+++ b/src/windows-emulator/process_context.cpp
@@ -388,6 +388,7 @@ void process_context::setup(x86_64_emulator& emu, memory_manager& memory, regist
     this->rtl_user_thread_start = ntdll.find_export("RtlUserThreadStart");
     this->ki_user_apc_dispatcher = ntdll.find_export("KiUserApcDispatcher");
     this->ki_user_exception_dispatcher = ntdll.find_export("KiUserExceptionDispatcher");
+    this->instrumentation_callback = 0;
 
     this->default_register_set = emu.save_registers();
 }
@@ -413,6 +414,7 @@ void process_context::serialize(utils::buffer_serializer& buffer) const
     buffer.write_optional(this->rtl_user_thread_start32);
     buffer.write(this->ki_user_apc_dispatcher);
     buffer.write(this->ki_user_exception_dispatcher);
+    buffer.write(this->instrumentation_callback);
 
     buffer.write(this->events);
     buffer.write(this->files);
@@ -467,6 +469,7 @@ void process_context::deserialize(utils::buffer_deserializer& buffer)
     buffer.read_optional(this->rtl_user_thread_start32);
     buffer.read(this->ki_user_apc_dispatcher);
     buffer.read(this->ki_user_exception_dispatcher);
+    buffer.read(this->instrumentation_callback);
 
     buffer.read(this->events);
     buffer.read(this->files);
diff --git a/src/windows-emulator/process_context.hpp b/src/windows-emulator/process_context.hpp
index 634fe3c4..01070a06 100644
--- a/src/windows-emulator/process_context.hpp
+++ b/src/windows-emulator/process_context.hpp
@@ -110,6 +110,7 @@ struct process_context
     uint64_t rtl_user_thread_start{};
     uint64_t ki_user_apc_dispatcher{};
     uint64_t ki_user_exception_dispatcher{};
+    uint64_t instrumentation_callback{};
 
     // For WOW64 processes
     std::optional<emulator_object<PEB32>> peb32;
diff --git a/src/windows-emulator/syscall_dispatcher.cpp b/src/windows-emulator/syscall_dispatcher.cpp
index 370da9af..4aabed1c 100644
--- a/src/windows-emulator/syscall_dispatcher.cpp
+++ b/src/windows-emulator/syscall_dispatcher.cpp
@@ -101,6 +101,19 @@ void syscall_dispatcher::dispatch(windows_emulator& win_emu)
         }
 
         entry->second.handler(c);
+
+        if (context.instrumentation_callback != 0 && entry->second.name != "NtContinue")
+        {
+            uint64_t rip_old = emu.reg<uint64_t>(x86_register::rip);
+
+            // The increase in RIP caused by executing the syscall here has not yet occurred.
+            // If RIP is set directly, it will lead to an incorrect address, so the length of
+            // the syscall instruction needs to be subtracted.
+            emu.reg<uint64_t>(x86_register::rip, context.instrumentation_callback - 2);
+
+            emu.reg<uint64_t>(x86_register::r10, rip_old);
+        }
+
     }
     catch (std::exception& e)
     {
diff --git a/src/windows-emulator/syscalls/process.cpp b/src/windows-emulator/syscalls/process.cpp
index 0d9105e6..b3cd0fbc 100644
--- a/src/windows-emulator/syscalls/process.cpp
+++ b/src/windows-emulator/syscalls/process.cpp
@@ -281,6 +281,23 @@ namespace syscalls
             return STATUS_SUCCESS;
         }
 
+        if (info_class == ProcessInstrumentationCallback)
+        {
+
+            if (process_information_length != sizeof(PROCESS_INSTRUMENTATION_CALLBACK_INFORMATION))
+            {
+                return STATUS_BUFFER_OVERFLOW;
+            }
+
+            PROCESS_INSTRUMENTATION_CALLBACK_INFORMATION info;
+
+            c.emu.read_memory(process_information, &info, sizeof(PROCESS_INSTRUMENTATION_CALLBACK_INFORMATION));
+
+            c.proc.instrumentation_callback = (uint64_t)info.Callback;
+
+            return STATUS_SUCCESS;
+        }
+
         c.win_emu.log.error("Unsupported info process class: %X\n", info_class);
         c.emu.stop();