diff --git a/Source/Windows/ARM64EC/Module.cpp b/Source/Windows/ARM64EC/Module.cpp index 933e4383b..388597c85 100644 --- a/Source/Windows/ARM64EC/Module.cpp +++ b/Source/Windows/ARM64EC/Module.cpp @@ -792,12 +792,15 @@ void NotifyMemoryAlloc(void* Address, SIZE_T Size, ULONG Type, ULONG Prot, BOOL return; } - if (!After || Status) { - return; + if (!After) { + ThreadCreationMutex.lock(); + } else { + if (!Status) { + std::scoped_lock Lock(ThreadCreationMutex); + InvalidationTracker->HandleMemoryProtectionNotification(reinterpret_cast(Address), static_cast(Size), Prot); + } + ThreadCreationMutex.unlock(); } - - std::scoped_lock Lock(ThreadCreationMutex); - InvalidationTracker->HandleMemoryProtectionNotification(reinterpret_cast(Address), static_cast(Size), Prot); } void NotifyMemoryFree(void* Address, SIZE_T Size, ULONG FreeType, BOOL After, NTSTATUS Status) { @@ -805,15 +808,15 @@ void NotifyMemoryFree(void* Address, SIZE_T Size, ULONG FreeType, BOOL After, NT return; } - if (After) { - return; - } - - std::scoped_lock Lock(ThreadCreationMutex); - if (FreeType & MEM_DECOMMIT) { - InvalidationTracker->InvalidateAlignedInterval(reinterpret_cast(Address), static_cast(Size), true); - } else if (FreeType & MEM_RELEASE) { - InvalidationTracker->InvalidateContainingSection(reinterpret_cast(Address), true); + if (!After) { + ThreadCreationMutex.lock(); + if (FreeType & MEM_DECOMMIT) { + InvalidationTracker->InvalidateAlignedInterval(reinterpret_cast(Address), static_cast(Size), true); + } else if (FreeType & MEM_RELEASE) { + InvalidationTracker->InvalidateContainingSection(reinterpret_cast(Address), true); + } + } else { + ThreadCreationMutex.unlock(); } } @@ -822,12 +825,15 @@ void NotifyMemoryProtect(void* Address, SIZE_T Size, ULONG NewProt, BOOL After, return; } - if (!After || Status) { - return; + if (!After) { + ThreadCreationMutex.lock(); + } else { + if (!Status) { + std::scoped_lock Lock(ThreadCreationMutex); + InvalidationTracker->HandleMemoryProtectionNotification(reinterpret_cast(Address), static_cast(Size), NewProt); + } + ThreadCreationMutex.unlock(); } - - std::scoped_lock Lock(ThreadCreationMutex); - InvalidationTracker->HandleMemoryProtectionNotification(reinterpret_cast(Address), static_cast(Size), NewProt); } NTSTATUS NotifyMapViewOfSection(void* Unk1, void* Address, void* Unk2, SIZE_T Size, ULONG AllocType, ULONG Prot) { @@ -849,15 +855,16 @@ void NotifyUnmapViewOfSection(void* Address, BOOL After, NTSTATUS Status) { return; } - if (After) { - return; - } - - std::scoped_lock Lock(ThreadCreationMutex); - auto [Start, Size] = InvalidationTracker->InvalidateContainingSection(reinterpret_cast(Address), true); - if (Size) { - std::scoped_lock Lock(CTX->GetCodeInvalidationMutex()); - CTX->RemoveForceTSOInformation(Start, Size); + if (!After) { + ThreadCreationMutex.lock(); + std::scoped_lock Lock(ThreadCreationMutex); + auto [Start, Size] = InvalidationTracker->InvalidateContainingSection(reinterpret_cast(Address), true); + if (Size) { + std::scoped_lock Lock(CTX->GetCodeInvalidationMutex()); + CTX->RemoveForceTSOInformation(Start, Size); + } + } else { + ThreadCreationMutex.unlock(); } } diff --git a/Source/Windows/WOW64/Module.cpp b/Source/Windows/WOW64/Module.cpp index 6227f14a1..e2c5ba81f 100644 --- a/Source/Windows/WOW64/Module.cpp +++ b/Source/Windows/WOW64/Module.cpp @@ -847,30 +847,37 @@ void BTCpuNotifyMemoryDirty(void* Address, SIZE_T Size) { } void BTCpuNotifyMemoryAlloc(void* Address, SIZE_T Size, ULONG Type, ULONG Prot, BOOL After, ULONG Status) { - if (!After || Status) { - return; + if (!After) { + ThreadCreationMutex.lock(); + } else { + if (!Status) { + InvalidationTracker->HandleMemoryProtectionNotification(reinterpret_cast(Address), static_cast(Size), Prot); + } + ThreadCreationMutex.unlock(); } - std::scoped_lock Lock(ThreadCreationMutex); - InvalidationTracker->HandleMemoryProtectionNotification(reinterpret_cast(Address), static_cast(Size), Prot); } void BTCpuNotifyMemoryProtect(void* Address, SIZE_T Size, ULONG NewProt, BOOL After, ULONG Status) { - if (!After || Status) { - return; + if (!After) { + ThreadCreationMutex.lock(); + } else { + if (!Status) { + InvalidationTracker->HandleMemoryProtectionNotification(reinterpret_cast(Address), static_cast(Size), NewProt); + } + ThreadCreationMutex.unlock(); } - std::scoped_lock Lock(ThreadCreationMutex); - InvalidationTracker->HandleMemoryProtectionNotification(reinterpret_cast(Address), static_cast(Size), NewProt); } void BTCpuNotifyMemoryFree(void* Address, SIZE_T Size, ULONG FreeType, BOOL After, ULONG Status) { - if (After) { - return; - } - std::scoped_lock Lock(ThreadCreationMutex); - if (!Size) { - InvalidationTracker->InvalidateContainingSection(reinterpret_cast(Address), true); - } else if (FreeType & MEM_DECOMMIT) { - InvalidationTracker->InvalidateAlignedInterval(reinterpret_cast(Address), static_cast(Size), true); + if (!After) { + ThreadCreationMutex.lock(); + if (!Size) { + InvalidationTracker->InvalidateContainingSection(reinterpret_cast(Address), true); + } else if (FreeType & MEM_DECOMMIT) { + InvalidationTracker->InvalidateAlignedInterval(reinterpret_cast(Address), static_cast(Size), true); + } + } else { + ThreadCreationMutex.unlock(); } } @@ -881,12 +888,12 @@ NTSTATUS BTCpuNotifyMapViewOfSection(void* Unk1, void* Address, void* Unk2, SIZE } void BTCpuNotifyUnmapViewOfSection(void* Address, BOOL After, ULONG Status) { - if (After) { - return; + if (!After) { + ThreadCreationMutex.lock(); + InvalidationTracker->InvalidateContainingSection(reinterpret_cast(Address), true); + } else { + ThreadCreationMutex.unlock(); } - - std::scoped_lock Lock(ThreadCreationMutex); - InvalidationTracker->InvalidateContainingSection(reinterpret_cast(Address), true); } void BTCpuNotifyReadFile(HANDLE Handle, void* Address, SIZE_T Size, BOOL After, NTSTATUS Status) {}