diff --git a/Source/Windows/WOW64/Module.cpp b/Source/Windows/WOW64/Module.cpp index 677a16f00..a13a86bf3 100644 --- a/Source/Windows/WOW64/Module.cpp +++ b/Source/Windows/WOW64/Module.cpp @@ -35,6 +35,7 @@ $end_info$ #include #include #include +#include #include #include #include @@ -94,6 +95,7 @@ namespace { SYSTEM_CPU_INFORMATION CpuInfo{}; std::mutex ThreadSuspendLock; + std::unordered_set InitializedWOWThreads; // Set of TIDs, `ThreadSuspendLock` must be locked when accessing std::pair GetThreadTLS(HANDLE Thread) { THREAD_BASIC_INFORMATION Info; @@ -555,6 +557,8 @@ void BTCpuProcessInit() { NTSTATUS BTCpuThreadInit() { GetTLS().ThreadState() = CTX->CreateThread(nullptr, 0); + std::scoped_lock Lock(ThreadSuspendLock); + InitializedWOWThreads.emplace(GetCurrentThreadId()); return STATUS_SUCCESS; } @@ -564,6 +568,17 @@ NTSTATUS BTCpuThreadTerm(HANDLE Thread) { return Err; } + { + THREAD_BASIC_INFORMATION Info; + if (NTSTATUS Err = NtQueryInformationThread(Thread, ThreadBasicInformation, &Info, sizeof(Info), nullptr); Err) { + return Err; + } + + const auto ThreadTID = reinterpret_cast(Info.ClientId.UniqueThread); + std::scoped_lock Lock(ThreadSuspendLock); + InitializedWOWThreads.erase(ThreadTID); + } + CTX->DestroyThread(TLS.ThreadState()); return STATUS_SUCCESS; } @@ -666,6 +681,11 @@ NTSTATUS BTCpuSuspendLocalThread(HANDLE Thread, ULONG *Count) { } std::scoped_lock Lock(ThreadSuspendLock); + + // If the thread hasn't yet been initialized, suspend it without special handling as it wont yet have entered the JIT + if (!InitializedWOWThreads.contains(ThreadTID)) + return NtSuspendThread(Thread, Count); + // If CONTROL_IN_JIT is unset at this point, then it can never be set (and thus the JIT cannot be reentered) as // CONTROL_PAUSED has been set, as such, while this may redundantly request interrupts in rare cases it will never // miss them