From bb3f8af81ad6b4c60cd8d1d33c18e7472759fe57 Mon Sep 17 00:00:00 2001
From: Stephen Miller <56742918+StevenMiller123@users.noreply.github.com>
Date: Tue, 3 Jun 2025 01:29:25 -0500
Subject: [PATCH] Core: Protect fixes (#3029)

* Swap do-while to while

If we use a do-while loop, we waste time if `aligned_size = 0`.  This is also still accurate to FreeBSD behavior, where it returns success if `start == end` during mprotect.
This also effectively prevents the memory assert seen in updated versions of RESIDENT EVIL 2 (CUSA09193)

* Move prot validation outside loop

The prot variable shouldn't change during a mprotect call, so we can check the flags before protecting instead.
Also cleans up the code for prot validation.
This should improve performance, and is more accurate to FreeBSD code.

* Add logging for protect calls

This will help in debugging future problems
---
 src/core/libraries/kernel/memory.cpp |  4 ++++
 src/core/memory.cpp                  | 32 +++++++++++++++-------------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/src/core/libraries/kernel/memory.cpp b/src/core/libraries/kernel/memory.cpp
index ce694dc1e..18676cbdf 100644
--- a/src/core/libraries/kernel/memory.cpp
+++ b/src/core/libraries/kernel/memory.cpp
@@ -264,6 +264,8 @@ int PS4_SYSV_ABI sceKernelQueryMemoryProtection(void* addr, void** start, void**
 }
 
 s32 PS4_SYSV_ABI sceKernelMprotect(const void* addr, u64 size, s32 prot) {
+    LOG_INFO(Kernel_Vmm, "called addr = {}, size = {:#x}, prot = {:#x}", fmt::ptr(addr), size,
+             prot);
     Core::MemoryManager* memory_manager = Core::Memory::Instance();
     Core::MemoryProt protection_flags = static_cast<Core::MemoryProt>(prot);
     return memory_manager->Protect(std::bit_cast<VAddr>(addr), size, protection_flags);
@@ -279,6 +281,8 @@ s32 PS4_SYSV_ABI posix_mprotect(const void* addr, u64 size, s32 prot) {
 }
 
 s32 PS4_SYSV_ABI sceKernelMtypeprotect(const void* addr, u64 size, s32 mtype, s32 prot) {
+    LOG_INFO(Kernel_Vmm, "called addr = {}, size = {:#x}, prot = {:#x}", fmt::ptr(addr), size,
+             prot);
     Core::MemoryManager* memory_manager = Core::Memory::Instance();
     Core::MemoryProt protection_flags = static_cast<Core::MemoryProt>(prot);
     return memory_manager->Protect(std::bit_cast<VAddr>(addr), size, protection_flags);
diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index bf9d1cabd..8550ece17 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -557,18 +557,6 @@ s64 MemoryManager::ProtectBytes(VAddr addr, VirtualMemoryArea vma_base, size_t s
         return adjusted_size;
     }
 
-    // Validate protection flags
-    constexpr static MemoryProt valid_flags = MemoryProt::NoAccess | MemoryProt::CpuRead |
-                                              MemoryProt::CpuReadWrite | MemoryProt::GpuRead |
-                                              MemoryProt::GpuWrite | MemoryProt::GpuReadWrite;
-
-    MemoryProt invalid_flags = prot & ~valid_flags;
-    if (u32(invalid_flags) != 0 && u32(invalid_flags) != u32(MemoryProt::NoAccess)) {
-        LOG_ERROR(Kernel_Vmm, "Invalid protection flags: prot = {:#x}, invalid flags = {:#x}",
-                  u32(prot), u32(invalid_flags));
-        return ORBIS_KERNEL_ERROR_EINVAL;
-    }
-
     // Change protection
     vma_base.prot = prot;
 
@@ -598,11 +586,25 @@ s64 MemoryManager::ProtectBytes(VAddr addr, VirtualMemoryArea vma_base, size_t s
 
 s32 MemoryManager::Protect(VAddr addr, size_t size, MemoryProt prot) {
     std::scoped_lock lk{mutex};
-    s64 protected_bytes = 0;
 
+    // Validate protection flags
+    constexpr static MemoryProt valid_flags = MemoryProt::NoAccess | MemoryProt::CpuRead |
+                                              MemoryProt::CpuReadWrite | MemoryProt::GpuRead |
+                                              MemoryProt::GpuWrite | MemoryProt::GpuReadWrite;
+
+    MemoryProt invalid_flags = prot & ~valid_flags;
+    if (invalid_flags != MemoryProt::NoAccess) {
+        LOG_ERROR(Kernel_Vmm, "Invalid protection flags");
+        return ORBIS_KERNEL_ERROR_EINVAL;
+    }
+
+    // Align addr and size to the nearest page boundary.
     auto aligned_addr = Common::AlignDown(addr, 16_KB);
     auto aligned_size = Common::AlignUp(size + addr - aligned_addr, 16_KB);
-    do {
+
+    // Protect all VMAs between aligned_addr and aligned_addr + aligned_size.
+    s64 protected_bytes = 0;
+    while (protected_bytes < aligned_size) {
         auto it = FindVMA(aligned_addr + protected_bytes);
         auto& vma_base = it->second;
         ASSERT_MSG(vma_base.Contains(addr + protected_bytes, 0), "Address {:#x} is out of bounds",
@@ -615,7 +617,7 @@ s32 MemoryManager::Protect(VAddr addr, size_t size, MemoryProt prot) {
             return result;
         }
         protected_bytes += result;
-    } while (protected_bytes < aligned_size);
+    }
 
     return ORBIS_OK;
 }