From 409d6918778652f3473eaa4d0ab099544b368ad3 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Sun, 18 Aug 2024 18:19:34 +0000
Subject: [PATCH] WOW64: Only handle unaligned atomics in JIT code

---
 Source/Windows/WOW64/Module.cpp | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/Source/Windows/WOW64/Module.cpp b/Source/Windows/WOW64/Module.cpp
index 301b4b9f7..4c0cb902f 100644
--- a/Source/Windows/WOW64/Module.cpp
+++ b/Source/Windows/WOW64/Module.cpp
@@ -125,7 +125,9 @@ uint64_t GetWowTEB(void* TEB) {
 
 bool IsAddressInJit(uint64_t Address) {
   const auto& Config = SignalDelegator->GetConfig();
-  if (Address >= Config.DispatcherBegin && Address < Config.DispatcherEnd) return true;
+  if (Address >= Config.DispatcherBegin && Address < Config.DispatcherEnd) {
+    return true;
+  }
 
   auto Thread = GetTLS().ThreadState();
   return Thread->CTX->IsAddressInCodeBuffer(Thread, Address);
@@ -650,22 +652,18 @@ NTSTATUS BTCpuSuspendLocalThread(HANDLE Thread, ULONG* Count) {
 NTSTATUS BTCpuResetToConsistentState(EXCEPTION_POINTERS* Ptrs) {
   auto* Context = Ptrs->ContextRecord;
   auto* Exception = Ptrs->ExceptionRecord;
-  if (Exception->ExceptionCode == EXCEPTION_DATATYPE_MISALIGNMENT && Context::HandleUnalignedAccess(Context)) {
-    LogMan::Msg::DFmt("Handled unaligned atomic: new pc: {:X}", Context->Pc);
-    NtContinue(Context, FALSE);
-  }
+  auto Thread = GetTLS().ThreadState();
 
   if (Exception->ExceptionCode == EXCEPTION_ACCESS_VIOLATION) {
     const auto FaultAddress = static_cast<uint64_t>(Exception->ExceptionInformation[1]);
 
-
     if (Context::HandleSuspendInterrupt(Context, FaultAddress)) {
       LogMan::Msg::DFmt("Resumed from suspend");
       NtContinue(Context, FALSE);
     }
 
     bool HandledRWX = false;
-    if (GetTLS().ThreadState()) {
+    if (Thread) {
       std::scoped_lock Lock(ThreadCreationMutex);
       HandledRWX = InvalidationTracker->HandleRWXAccessViolation(FaultAddress);
     }
@@ -676,10 +674,15 @@ NTSTATUS BTCpuResetToConsistentState(EXCEPTION_POINTERS* Ptrs) {
     }
   }
 
-  if (!IsAddressInJit(Context->Pc)) {
+  if (!Thread || !IsAddressInJit(Context->Pc)) {
     return STATUS_SUCCESS;
   }
 
+  if (Exception->ExceptionCode == EXCEPTION_DATATYPE_MISALIGNMENT && Context::HandleUnalignedAccess(Context)) {
+    LogMan::Msg::DFmt("Handled unaligned atomic: new pc: {:X}", Context->Pc);
+    NtContinue(Context, FALSE);
+  }
+
   LogMan::Msg::DFmt("Reconstructing context");
 
   WOW64_CONTEXT WowContext = Context::ReconstructWowContext(Context);