Merge pull request #5197 from bylaws/rwxp5

Windows: Improve handling of RWX memory

Source/Windows/ARM64EC/Module.cpp

@@ -126,6 +126,10 @@ struct ThreadCPUArea {
   }
 };
 
+struct FrontendThreadData {
+  bool InLockedRWXRead {};
+};
+
 namespace {
 fextl::unique_ptr<FEXCore::Context::Context> CTX;
 fextl::unique_ptr<FEX::DummyHandlers::DummySignalDelegator> SignalDelegator;
@@ -150,6 +154,10 @@ ThreadCPUArea GetCPUArea() {
   return ThreadCPUArea(NtCurrentTeb());
 }
 
+FrontendThreadData* GetFrontendThreadData(FEXCore::Core::InternalThreadState* Thread) {
+  return static_cast<FrontendThreadData*>(Thread->FrontendPtr);
+}
+
 bool IsEmulatorStackAddress(const ThreadCPUArea CPUArea, uint64_t Address) {
   return Address <= CPUArea.EmulatorStackBase() && Address >= CPUArea.EmulatorStackLimit();
 }
@@ -873,7 +881,30 @@ void BTCpu64NotifyMemoryDirty(void* Address, SIZE_T Size) {
   InvalidationTracker->InvalidateAlignedInterval(reinterpret_cast<uint64_t>(Address), static_cast<uint64_t>(Size), false);
 }
 
-void BTCpu64NotifyReadFile(HANDLE Handle, void* Address, SIZE_T Size, BOOL After, NTSTATUS Status) {}
+void BTCpu64NotifyReadFile(HANDLE Handle, void* Address, SIZE_T Size, BOOL After, NTSTATUS Status) {
+  auto* ThreadState = GetCPUArea().ThreadState();
+  if (!InvalidationTracker || !ThreadState) {
+    return;
+  }
+
+  auto& InLockedRWXRead = GetFrontendThreadData(ThreadState)->InLockedRWXRead;
+  if (!After) {
+    ThreadCreationMutex.lock();
+    CTX->GetCodeInvalidationMutex().lock();
+    if (InvalidationTracker->BeginUntrackedWriteLocked(reinterpret_cast<uint64_t>(Address), static_cast<uint64_t>(Size))) {
+      InLockedRWXRead = true;
+    } else {
+      CTX->GetCodeInvalidationMutex().unlock();
+      ThreadCreationMutex.unlock();
+    }
+  } else {
+    if (InLockedRWXRead) {
+      InLockedRWXRead = false;
+      CTX->GetCodeInvalidationMutex().unlock();
+      ThreadCreationMutex.unlock();
+    }
+  }
+}
 
 NTSTATUS ThreadInit() {
   std::scoped_lock Lock(ThreadCreationMutex);
@@ -928,6 +959,8 @@ NTSTATUS ThreadInit() {
                             .AMD64_ControlWord = 0x27f};
   Exception::LoadStateFromECContext(Thread, CPUArea.ContextAmd64().AMD64_Context);
 
+  Thread->FrontendPtr = new FrontendThreadData();
+
   {
     auto ThreadTID = GetCurrentThreadId();
     Threads.emplace(ThreadTID, Thread);
@@ -983,6 +1016,8 @@ NTSTATUS ThreadTerm(HANDLE Thread, LONG ExitCode) {
   }
   auto ThreadState = CPUArea.ThreadState();
 
+  delete GetFrontendThreadData(ThreadState);
+
   // GDT and LDT are mirrored, only free one.
   delete[] ThreadState->CurrentFrame->State.segment_arrays[FEXCore::Core::CPUState::SEGMENT_ARRAY_INDEX_GDT];
 

Source/Windows/Common/InvalidationTracker.cpp

@@ -142,54 +142,38 @@ void InvalidationTracker::InvalidateAlignedInterval(uint64_t Address, uint64_t S
 }
 
 void InvalidationTracker::ReprotectRWXIntervals(uint64_t Address, uint64_t Size) {
-  const auto End = Address + Size;
-  std::shared_lock Lock(IntervalsLock);
-
-  if (SMCDetectionDisabled) {
-    return;
-  }
-
-  do {
-    const auto Query = RWXIntervals.Query(Address);
-    if (Query.Enclosed) {
-      void* TmpAddress = reinterpret_cast<void*>(Address);
-      SIZE_T TmpSize = static_cast<SIZE_T>(std::min(End, Address + Query.Size) - Address);
-      ULONG TmpProt;
-      NtProtectVirtualMemory(NtCurrentProcess(), &TmpAddress, &TmpSize, PAGE_EXECUTE_READ, &TmpProt);
-    } else if (!Query.Size) {
-      // No more regions past `Address` in the interval list
-      break;
-    }
-
-    Address += Query.Size;
-  } while (Address < End);
+  ProtectRWXIntervalsInternal(Address, Size, false);
 }
 
 bool InvalidationTracker::HandleRWXAccessViolation(FEXCore::Core::InternalThreadState* Thread, uint64_t HostPc, uint64_t FaultAddress) {
   const bool NeedsInvalidate = [&](uint64_t Address) {
     std::shared_lock Lock(IntervalsLock);
-    const bool Enclosed = RWXIntervals.Query(Address).Enclosed;
-    // Invalidate just the single faulting page
-    if (!Enclosed) {
-      return false;
-    }
-
-    ULONG TmpProt;
-    void* TmpAddress = reinterpret_cast<void*>(Address);
-    SIZE_T TmpSize = 1;
-    NtProtectVirtualMemory(NtCurrentProcess(), &TmpAddress, &TmpSize, PAGE_EXECUTE_READWRITE, &TmpProt);
-    return true;
+    return RWXIntervals.Query(Address).Enclosed;
   }(FaultAddress);
 
   if (NeedsInvalidate) {
     // IntervalsLock cannot be held during invalidation
-    InvalidateIntervalInternal(FaultAddress & FEXCore::Utils::FEX_PAGE_MASK, FEXCore::Utils::FEX_PAGE_SIZE);
+    {
+      std::scoped_lock Lock(CTX.GetCodeInvalidationMutex());
+
+      InvalidateIntervalInternalLocked(FaultAddress & FEXCore::Utils::FEX_PAGE_MASK, FEXCore::Utils::FEX_PAGE_SIZE);
+
+      // Invalidate, then unprotect the faulting page with the compilation lock held to ensure that any racing invalidations are not dropped.
+      ULONG TmpProt;
+      void* TmpAddress = reinterpret_cast<void*>(FaultAddress);
+      SIZE_T TmpSize = 1;
+      NtProtectVirtualMemory(NtCurrentProcess(), &TmpAddress, &TmpSize, PAGE_EXECUTE_READWRITE, &TmpProt);
+    }
     DetectMonoBackpatcherBlock(Thread, HostPc);
     return true;
   }
   return false;
 }
 
+bool InvalidationTracker::BeginUntrackedWriteLocked(uint64_t Address, uint64_t Size) {
+  return ProtectRWXIntervalsInternal(Address, Size, true);
+}
+
 FEXCore::HLE::ExecutableRangeInfo InvalidationTracker::QueryExecutableRange(uint64_t Address) {
   std::shared_lock Lock(IntervalsLock);
   const auto XResult = XIntervals.Query(Address);
@@ -254,11 +238,53 @@ void InvalidationTracker::DisableSMCDetection() {
 }
 
 void InvalidationTracker::InvalidateIntervalInternal(uint64_t Address, uint64_t Size) {
-  std::scoped_lock Lock(CTX.GetCodeInvalidationMutex());
+  std::scoped_lock CodeLock(CTX.GetCodeInvalidationMutex());
+  InvalidateIntervalInternalLocked(Address, Size);
+}
+
+void InvalidationTracker::InvalidateIntervalInternalLocked(uint64_t Address, uint64_t Size) {
+  // NOTE: This assumes CodeInvalidationMutex is locked by the caller
   CTX.InvalidateCodeBuffersCodeRange(Address, Size);
   for (auto Thread : Threads) {
     CTX.InvalidateThreadCachedCodeRange(Thread.second, Address, Size);
   }
 }
 
+bool InvalidationTracker::ProtectRWXIntervalsInternal(uint64_t Address, uint64_t Size, bool ForWriteLocked) {
+  const auto End = Address + Size;
+  std::shared_lock Lock(IntervalsLock);
+
+  if (SMCDetectionDisabled) {
+    return false;
+  }
+
+  bool HitRWXInterval = false;
+  do {
+    const auto Query = RWXIntervals.Query(Address);
+    if (Query.Enclosed) {
+      if (!HitRWXInterval) {
+        if (ForWriteLocked) {
+          // If we are protecting as writable, then the entire range must be invalidated before any protections are
+          // applied and the invalidation mutex must be locked throughout.
+          // Do this lazily only when an RWX region is actually hit.
+          // NOTE: This assumes CodeInvalidationMutex is locked by the caller
+          InvalidateIntervalInternalLocked(Address, Size);
+        }
+        HitRWXInterval = true;
+      }
+      void* TmpAddress = reinterpret_cast<void*>(Address);
+      SIZE_T TmpSize = static_cast<SIZE_T>(std::min(End, Address + Query.Size) - Address);
+      ULONG TmpProt;
+      NtProtectVirtualMemory(NtCurrentProcess(), &TmpAddress, &TmpSize, ForWriteLocked ? PAGE_EXECUTE_READWRITE : PAGE_EXECUTE_READ, &TmpProt);
+    } else if (!Query.Size) {
+      // No more regions past `Address` in the interval list
+      break;
+    }
+
+    Address += Query.Size;
+  } while (Address < End);
+
+  return HitRWXInterval;
+}
+
 } // namespace FEX::Windows

Source/Windows/Common/InvalidationTracker.h

@@ -33,13 +33,23 @@ public:
   void InvalidateAlignedInterval(uint64_t Address, uint64_t Size, bool Free);
   void ReprotectRWXIntervals(uint64_t Address, uint64_t Size);
   bool HandleRWXAccessViolation(FEXCore::Core::InternalThreadState* Thread, uint64_t HostPC, uint64_t FaultAddress);
+
+  // Unprotects any RWX intervals in the input interval and invalidates code
+  // NOTE: CodeInvalidationMutex must be locked when calling this, and if true is returned, kept locked until the write ends.
+  bool BeginUntrackedWriteLocked(uint64_t Address, uint64_t Size);
+
   FEXCore::HLE::ExecutableRangeInfo QueryExecutableRange(uint64_t Address);
 
 private:
   void DetectMonoBackpatcherBlock(FEXCore::Core::InternalThreadState* Thread, uint64_t HostPC);
   void DisableSMCDetection();
   void InvalidateIntervalInternal(uint64_t Address, uint64_t Size);
+  // NOTE: This assumed CodeInvalidationMutex is locked by the caller
+  void InvalidateIntervalInternalLocked(uint64_t Address, uint64_t Size);
 
+  // NOTE: If ForWriteLocked is true then this assumes CodeInvalidationMutex is locked by the caller,
+  // and any code in the range will be invalidated before protection as RWX, otherwise protects as RX if false.
+  bool ProtectRWXIntervalsInternal(uint64_t Address, uint64_t Size, bool ForWriteLocked);
 
   FEXCore::IntervalList<uint64_t> XIntervals;
   FEXCore::IntervalList<uint64_t> RWXIntervals;

Source/Windows/WOW64/Module.cpp

@@ -107,6 +107,10 @@ struct TLS {
   }
 };
 
+struct FrontendThreadData {
+  bool InLockedRWXRead {};
+};
+
 class WowSyscallHandler;
 
 namespace {
@@ -142,6 +146,10 @@ TLS GetTLS() {
   return TLS {NtCurrentTeb()};
 }
 
+FrontendThreadData* GetFrontendThreadData(FEXCore::Core::InternalThreadState* Thread) {
+  return static_cast<FrontendThreadData*>(Thread->FrontendPtr);
+}
+
 uint64_t GetWowTEB(void* TEB) {
   static constexpr size_t WowTEBOffsetMemberOffset {0x180c};
   return static_cast<uint64_t>(
@@ -610,6 +618,8 @@ void BTCpuThreadInit() {
   TLS.ThreadState() = Thread;
   TLS.ControlWord().fetch_or(ControlBits::WOW_CPU_AREA_DIRTY, std::memory_order::relaxed);
 
+  Thread->FrontendPtr = new FrontendThreadData();
+
   auto ThreadTID = GetCurrentThreadId();
   Threads.emplace(ThreadTID, Thread);
   if (StatAllocHandler) {
@@ -656,6 +666,8 @@ void BTCpuThreadTerm(HANDLE Thread, LONG ExitCode) {
   }
   auto ThreadState = TLS.ThreadState();
 
+  delete GetFrontendThreadData(ThreadState);
+
   // GDT and LDT are mirrored, only free one.
   delete[] ThreadState->CurrentFrame->State.segment_arrays[FEXCore::Core::CPUState::SEGMENT_ARRAY_INDEX_GDT];
 
@@ -1004,7 +1016,25 @@ void BTCpuNotifyUnmapViewOfSection(void* Address, BOOL After, ULONG Status) {
   }
 }
 
-void BTCpuNotifyReadFile(HANDLE Handle, void* Address, SIZE_T Size, BOOL After, NTSTATUS Status) {}
+void BTCpuNotifyReadFile(HANDLE Handle, void* Address, SIZE_T Size, BOOL After, NTSTATUS Status) {
+  auto& InLockedRWXRead = GetFrontendThreadData(GetTLS().ThreadState())->InLockedRWXRead;
+  if (!After) {
+    ThreadCreationMutex.lock();
+    CTX->GetCodeInvalidationMutex().lock();
+    if (InvalidationTracker->BeginUntrackedWriteLocked(reinterpret_cast<uint64_t>(Address), static_cast<uint64_t>(Size))) {
+      InLockedRWXRead = true;
+    } else {
+      CTX->GetCodeInvalidationMutex().unlock();
+      ThreadCreationMutex.unlock();
+    }
+  } else {
+    if (InLockedRWXRead) {
+      InLockedRWXRead = false;
+      CTX->GetCodeInvalidationMutex().unlock();
+      ThreadCreationMutex.unlock();
+    }
+  }
+}
 
 BOOLEAN WINAPI BTCpuIsProcessorFeaturePresent(UINT Feature) {
   return CPUFeatures->IsFeaturePresent(Feature) ? TRUE : FALSE;