diff --git a/Source/Windows/ARM64EC/Module.cpp b/Source/Windows/ARM64EC/Module.cpp index 3aac5ff65..02ae43046 100644 --- a/Source/Windows/ARM64EC/Module.cpp +++ b/Source/Windows/ARM64EC/Module.cpp @@ -549,6 +549,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(MainModule)); + CPUFeatures.emplace(*CTX); X64ReturnInstr = ::VirtualAlloc(nullptr, FEXCore::Utils::FEX_PAGE_SIZE, MEM_COMMIT, PAGE_EXECUTE_READWRITE); @@ -696,6 +700,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(Address)); return STATUS_SUCCESS; } diff --git a/Source/Windows/Common/InvalidationTracker.cpp b/Source/Windows/Common/InvalidationTracker.cpp index b8ce08927..363609f19 100644 --- a/Source/Windows/Common/InvalidationTracker.cpp +++ b/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(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(Address), MemoryBasicInformation, &Info, sizeof(Info), nullptr)) { diff --git a/Source/Windows/Common/InvalidationTracker.h b/Source/Windows/Common/InvalidationTracker.h index f539c5b01..771b352e4 100644 --- a/Source/Windows/Common/InvalidationTracker.h +++ b/Source/Windows/Common/InvalidationTracker.h @@ -21,6 +21,7 @@ class InvalidationTracker { public: InvalidationTracker(FEXCore::Context::Context& CTX, const std::unordered_map& 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); diff --git a/Source/Windows/WOW64/Module.cpp b/Source/Windows/WOW64/Module.cpp index 5219ccdf7..ebf5f01b7 100644 --- a/Source/Windows/WOW64/Module.cpp +++ b/Source/Windows/WOW64/Module.cpp @@ -476,6 +476,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(MainModule)); + CPUFeatures.emplace(*CTX); // Allocate the syscall/unixcall trampolines in the lower 2GB of the address space diff --git a/Source/Windows/include/winternl.h b/Source/Windows/include/winternl.h index da486cfe6..d55d1e569 100644 --- a/Source/Windows/include/winternl.h +++ b/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*);