Windows: Track RWX regions in mapped images

Source/Windows/ARM64EC/Module.cpp

@@ -547,6 +547,10 @@ NTSTATUS ProcessInit() {
   CTX->SetSyscallHandler(SyscallHandler.get());
   CTX->InitCore();
   InvalidationTracker.emplace(*CTX, Threads);
+
+  auto MainModule = reinterpret_cast<__TEB*>(NtCurrentTeb())->Peb->ImageBaseAddress;
+  InvalidationTracker->HandleImageMap(reinterpret_cast<uint64_t>(MainModule));
+
   CPUFeatures.emplace(*CTX);
 
   X64ReturnInstr = ::VirtualAlloc(nullptr, FEXCore::Utils::FEX_PAGE_SIZE, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
@@ -694,6 +698,12 @@ void NotifyMemoryProtect(void* Address, SIZE_T Size, ULONG NewProt, BOOL After,
 }
 
 NTSTATUS NotifyMapViewOfSection(void* Unk1, void* Address, void* Unk2, SIZE_T Size, ULONG AllocType, ULONG Prot) {
+  if (!InvalidationTracker || !GetCPUArea().ThreadState()) {
+    return STATUS_SUCCESS;
+  }
+
+  std::scoped_lock Lock(ThreadCreationMutex);
+  InvalidationTracker->HandleImageMap(reinterpret_cast<uint64_t>(Address));
   return STATUS_SUCCESS;
 }
 

Source/Windows/Common/InvalidationTracker.cpp

@@ -18,14 +18,14 @@ void InvalidationTracker::HandleMemoryProtectionNotification(uint64_t Address, u
   const auto AlignedBase = Address & FEXCore::Utils::FEX_PAGE_MASK;
   const auto AlignedSize = (Address - AlignedBase + Size + FEXCore::Utils::FEX_PAGE_SIZE - 1) & FEXCore::Utils::FEX_PAGE_MASK;
 
-  if (Prot & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE)) {
+  if (Prot & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY)) {
     std::scoped_lock Lock(CTX.GetCodeInvalidationMutex());
     for (auto Thread : Threads) {
       CTX.InvalidateGuestCodeRange(Thread.second, AlignedBase, AlignedSize);
     }
   }
 
-  if (Prot & PAGE_EXECUTE_READWRITE) {
+  if (Prot & (PAGE_EXECUTE_WRITECOPY | PAGE_EXECUTE_READWRITE)) {
     LogMan::Msg::DFmt("Add SMC interval: {:X} - {:X}", AlignedBase, AlignedBase + AlignedSize);
     std::scoped_lock Lock(RWXIntervalsLock);
     RWXIntervals.Insert({AlignedBase, AlignedBase + AlignedSize});
@@ -35,6 +35,21 @@ void InvalidationTracker::HandleMemoryProtectionNotification(uint64_t Address, u
   }
 }
 
+void InvalidationTracker::HandleImageMap(uint64_t Address) {
+  auto* Nt = RtlImageNtHeader(reinterpret_cast<HMODULE>(Address));
+  auto* SectionsBegin = IMAGE_FIRST_SECTION(Nt);
+  auto* SectionsEnd = SectionsBegin + Nt->FileHeader.NumberOfSections;
+
+  for (auto* Section = SectionsBegin; Section != SectionsEnd; Section++) {
+    if ((Section->Characteristics & IMAGE_SCN_MEM_EXECUTE) && (Section->Characteristics & IMAGE_SCN_MEM_WRITE)) {
+      uint64_t SectionBase = Address + Section->VirtualAddress;
+      LogMan::Msg::DFmt("Add image SMC interval: {:X} - {:X}", SectionBase, SectionBase + Section->Misc.VirtualSize);
+      std::scoped_lock Lock(RWXIntervalsLock);
+      RWXIntervals.Insert({SectionBase, SectionBase + Section->Misc.VirtualSize});
+    }
+  }
+}
+
 void InvalidationTracker::InvalidateContainingSection(uint64_t Address, bool Free) {
   MEMORY_BASIC_INFORMATION Info;
   if (NtQueryVirtualMemory(NtCurrentProcess(), reinterpret_cast<void*>(Address), MemoryBasicInformation, &Info, sizeof(Info), nullptr)) {

Source/Windows/Common/InvalidationTracker.h

@@ -21,6 +21,7 @@ class InvalidationTracker {
 public:
   InvalidationTracker(FEXCore::Context::Context& CTX, const std::unordered_map<DWORD, FEXCore::Core::InternalThreadState*>& Threads);
   void HandleMemoryProtectionNotification(uint64_t Address, uint64_t Size, ULONG Prot);
+  void HandleImageMap(uint64_t Address);
   void InvalidateContainingSection(uint64_t Address, bool Free);
   void InvalidateAlignedInterval(uint64_t Address, uint64_t Size, bool Free);
   void ReprotectRWXIntervals(uint64_t Address, uint64_t Size);

Source/Windows/WOW64/Module.cpp

@@ -469,6 +469,10 @@ void BTCpuProcessInit() {
   CTX->SetSyscallHandler(SyscallHandler.get());
   CTX->InitCore();
   InvalidationTracker.emplace(*CTX, Threads);
+
+  auto MainModule = reinterpret_cast<__TEB*>(NtCurrentTeb())->Peb->ImageBaseAddress;
+  InvalidationTracker->HandleImageMap(reinterpret_cast<uint64_t>(MainModule));
+
   CPUFeatures.emplace(*CTX);
 
   // Allocate the syscall/unixcall trampolines in the lower 2GB of the address space

Source/Windows/include/winternl.h

@@ -477,6 +477,7 @@ NTSTATUS WINAPI RtlDeleteCriticalSection(RTL_CRITICAL_SECTION*);
 NTSTATUS WINAPI RtlEnterCriticalSection(RTL_CRITICAL_SECTION*);
 ULONG WINAPI RtlFindClearBitsAndSet(PRTL_BITMAP, ULONG, ULONG);
 ULONG WINAPI RtlGetCurrentDirectory_U(ULONG, LPWSTR);
+PIMAGE_NT_HEADERS WINAPI RtlImageNtHeader(HMODULE);
 PVOID WINAPI RtlImageDirectoryEntryToData(HMODULE, BOOL, WORD, ULONG*);
 void WINAPI RtlInitializeConditionVariable(RTL_CONDITION_VARIABLE*);
 NTSTATUS WINAPI RtlInitializeCriticalSection(RTL_CRITICAL_SECTION*);