From e35aab16ea87b36077dceadbb2b4436304fe9e8d Mon Sep 17 00:00:00 2001
From: Stephen Miller <millerste004@gmail.com>
Date: Wed, 19 Feb 2025 21:54:46 -0600
Subject: [PATCH] Implement protecting multiple VMAs

A handful of games expect this to work, and updated versions of Grand Theft Auto V crash if it doesn't work.
---
 src/core/memory.cpp | 43 ++++++++++++++++++++++++++++---------------
 src/core/memory.h   |  4 +++-
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index 4717b3a74..ab2a3b66b 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -466,19 +466,14 @@ int MemoryManager::QueryProtection(VAddr addr, void** start, void** end, u32* pr
     return ORBIS_OK;
 }
 
-int MemoryManager::Protect(VAddr addr, size_t size, MemoryProt prot) {
-    std::scoped_lock lk{mutex};
-
-    // Find the virtual memory area that contains the specified address range.
-    auto it = FindVMA(addr);
-    if (it == vma_map.end() || !it->second.Contains(addr, size)) {
-        LOG_ERROR(Core, "Address range not mapped");
-        return ORBIS_KERNEL_ERROR_EINVAL;
-    }
-
-    VirtualMemoryArea& vma = it->second;
-    if (vma.type == VMAType::Free) {
-        LOG_ERROR(Core, "Cannot change protection on free memory region");
+s64 MemoryManager::ProtectBytes(VAddr addr, VirtualMemoryArea vma_base, size_t size,
+                                MemoryProt prot) {
+    const auto start_in_vma = addr - vma_base.base;
+    const auto adjusted_size =
+        vma_base.size - start_in_vma < size ? vma_base.size - start_in_vma : size;
+    
+    if (vma_base.type == VMAType::Free) {
+        LOG_ERROR(Kernel_Vmm, "Cannot change protection on free memory region");
         return ORBIS_KERNEL_ERROR_EINVAL;
     }
 
@@ -489,13 +484,13 @@ int MemoryManager::Protect(VAddr addr, size_t size, MemoryProt prot) {
 
     MemoryProt invalid_flags = prot & ~valid_flags;
     if (u32(invalid_flags) != 0 && u32(invalid_flags) != u32(MemoryProt::NoAccess)) {
-        LOG_ERROR(Core, "Invalid protection flags: prot = {:#x}, invalid flags = {:#x}", u32(prot),
+        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.prot = prot;
+    vma_base.prot = prot;
 
     // Set permissions
     Core::MemoryPermission perms{};
@@ -518,6 +513,24 @@ int MemoryManager::Protect(VAddr addr, size_t size, MemoryProt prot) {
 
     impl.Protect(addr, size, perms);
 
+    return adjusted_size;
+}
+
+s32 MemoryManager::Protect(VAddr addr, size_t size, MemoryProt prot) {
+    std::scoped_lock lk{mutex};
+    s64 protected_bytes = 0;
+    do {
+        auto it = FindVMA(addr + protected_bytes);
+        auto& vma_base = it->second;
+        auto result = 0;
+        result = ProtectBytes(addr + protected_bytes, vma_base, size - protected_bytes, prot);
+        if (result < 0) {
+            // ProtectBytes returned an error, return it
+            return result;
+        }
+        protected_bytes += result;
+    } while (protected_bytes < size);
+
     return ORBIS_OK;
 }
 
diff --git a/src/core/memory.h b/src/core/memory.h
index 59e48b248..7237b3aa4 100644
--- a/src/core/memory.h
+++ b/src/core/memory.h
@@ -196,7 +196,9 @@ public:
 
     int QueryProtection(VAddr addr, void** start, void** end, u32* prot);
 
-    int Protect(VAddr addr, size_t size, MemoryProt prot);
+    s32 Protect(VAddr addr, size_t size, MemoryProt prot);
+
+    s64 ProtectBytes(VAddr addr, VirtualMemoryArea vma_base, size_t size, MemoryProt prot);
 
     int VirtualQuery(VAddr addr, int flags, ::Libraries::Kernel::OrbisVirtualQueryInfo* info);