From a78ffeeabaa0f2cd707f89f9a4b26800e6a34839 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Fri, 17 Nov 2023 23:52:13 +0000
Subject: [PATCH 1/2] WOW64: Call Wow64ProcessPendingCrossProcessItems on
 syscalls

This is needed to handle code invalidation requests sent by external
processes.
---
 Source/Windows/Defs/wow64.def     | 19 ++++++++++---------
 Source/Windows/WOW64/Module.cpp   |  1 +
 Source/Windows/include/winternl.h |  1 +
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/Source/Windows/Defs/wow64.def b/Source/Windows/Defs/wow64.def
index e5291402c..438ce7ebe 100644
--- a/Source/Windows/Defs/wow64.def
+++ b/Source/Windows/Defs/wow64.def
@@ -22,12 +22,13 @@ EXPORTS
   Wow64PassExceptionToGuest @16
   Wow64PrepareForDebuggerAttach @17 PRIVATE
   Wow64PrepareForException @18
-  Wow64RaiseException @19
-  Wow64ShallowThunkAllocObjectAttributes32TO64_FNC @20 PRIVATE
-  Wow64ShallowThunkAllocSecurityQualityOfService32TO64_FNC @21 PRIVATE
-  Wow64ShallowThunkSIZE_T32TO64 @22 PRIVATE
-  Wow64ShallowThunkSIZE_T64TO32 @23 PRIVATE
-  Wow64SuspendLocalThread @24
-  Wow64SystemServiceEx @25
-  Wow64ValidateUserCallTarget @26 PRIVATE
-  Wow64ValidateUserCallTargetFilter @27 PRIVATE
+  Wow64ProcessPendingCrossProcessItems @19
+  Wow64RaiseException @20
+  Wow64ShallowThunkAllocObjectAttributes32TO64_FNC @21 PRIVATE
+  Wow64ShallowThunkAllocSecurityQualityOfService32TO64_FNC @22 PRIVATE
+  Wow64ShallowThunkSIZE_T32TO64 @23 PRIVATE
+  Wow64ShallowThunkSIZE_T64TO32 @24 PRIVATE
+  Wow64SuspendLocalThread @25 PRIVATE
+  Wow64SystemServiceEx @26
+  Wow64ValidateUserCallTarget @27 PRIVATE
+  Wow64ValidateUserCallTargetFilter @28 PRIVATE
diff --git a/Source/Windows/WOW64/Module.cpp b/Source/Windows/WOW64/Module.cpp
index ee5dff680..677a16f00 100644
--- a/Source/Windows/WOW64/Module.cpp
+++ b/Source/Windows/WOW64/Module.cpp
@@ -460,6 +460,7 @@ public:
       const uint64_t EntryRAX = Frame->State.gregs[FEXCore::X86State::REG_RAX];
 
       Context::UnlockJITContext();
+      Wow64ProcessPendingCrossProcessItems();
       ReturnRAX = static_cast<uint64_t>(Wow64SystemServiceEx(static_cast<UINT>(EntryRAX),
                                                              reinterpret_cast<UINT *>(ReturnRSP + 4)));
       Context::LockJITContext();
diff --git a/Source/Windows/include/winternl.h b/Source/Windows/include/winternl.h
index dddcc6963..72fbb257a 100644
--- a/Source/Windows/include/winternl.h
+++ b/Source/Windows/include/winternl.h
@@ -90,6 +90,7 @@ typedef enum _MEMORY_INFORMATION_CLASS {
 } MEMORY_INFORMATION_CLASS;
 
 NTSTATUS WINAPI Wow64SystemServiceEx(UINT,UINT*);
+void WINAPI Wow64ProcessPendingCrossProcessItems(void);
 
 NTSTATUS WINAPI RtlWow64SetThreadContext(HANDLE,const WOW64_CONTEXT*);
 NTSTATUS WINAPI RtlWow64GetThreadContext(HANDLE,WOW64_CONTEXT*);

From 0806d4ec25415d32ffe903936b3378fdbaabac42 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Sat, 18 Nov 2023 00:51:54 +0000
Subject: [PATCH 2/2] WOW64: Handle target threads where the JIT is
 uninitialised

Suspend may be called on a thread before it has finished WOW64 initialisation,
keep track of all initialized threads and fallback to direct
NtSuspendThread when this is the case.
---
 Source/Windows/WOW64/Module.cpp | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

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 <atomic>
 #include <mutex>
 #include <utility>
+#include <unordered_set>
 #include <ntstatus.h>
 #include <windef.h>
 #include <winternl.h>
@@ -94,6 +95,7 @@ namespace {
   SYSTEM_CPU_INFORMATION CpuInfo{};
 
   std::mutex ThreadSuspendLock;
+  std::unordered_set<DWORD> InitializedWOWThreads; // Set of TIDs, `ThreadSuspendLock` must be locked when accessing
 
   std::pair<NTSTATUS, TLS> 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<uint64_t>(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