From 899aa0546cfd8cd44ce917d1b0067f0c60822d46 Mon Sep 17 00:00:00 2001
From: Gabriel A <gab.dark.100@gmail.com>
Date: Tue, 26 Mar 2024 22:50:20 -0300
Subject: [PATCH] Reprotection for bridges

---
 src/Ryujinx.Cpu/Jit/AddressSpacePartition.cs  | 22 ++++++++++++-
 .../Jit/AddressSpacePartitioned.cs            | 31 ++++++++++++++++---
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/src/Ryujinx.Cpu/Jit/AddressSpacePartition.cs b/src/Ryujinx.Cpu/Jit/AddressSpacePartition.cs
index 80460fa4fd..74c5a18acd 100644
--- a/src/Ryujinx.Cpu/Jit/AddressSpacePartition.cs
+++ b/src/Ryujinx.Cpu/Jit/AddressSpacePartition.cs
@@ -209,6 +209,7 @@ namespace Ryujinx.Cpu.Jit
         private ulong _cachedLastPagePa;
         private MemoryBlock _firstPageMemoryForUnmap;
         private ulong _firstPageOffsetForLateMap;
+        private MemoryPermission _firstPageMemoryProtection;
 
         public ulong Address { get; }
         public ulong Size { get; }
@@ -294,6 +295,11 @@ namespace Ryujinx.Cpu.Jit
             Debug.Assert(va + size <= EndAddress);
 
             _baseMemory.Reprotect(va - Address, size, protection, false);
+
+            if (va == Address)
+            {
+                _firstPageMemoryProtection = protection;
+            }
         }
 
         public void Reprotect(
@@ -319,7 +325,7 @@ namespace Ryujinx.Cpu.Jit
             return _baseMemory.GetPointer(va - Address, size);
         }
 
-        public void InsertBridgeAtEnd(AddressSpacePartition partitionAfter)
+        public void InsertBridgeAtEnd(AddressSpacePartition partitionAfter, bool useProtectionMirrors)
         {
             ulong firstPagePa = partitionAfter?._firstPagePa ?? ulong.MaxValue;
             ulong lastPagePa = _lastPagePa ?? ulong.MaxValue;
@@ -331,6 +337,12 @@ namespace Ryujinx.Cpu.Jit
                     (MemoryBlock firstPageMemory, ulong firstPageOffset) = partitionAfter.GetFirstPageMemoryAndOffset();
 
                     _baseMemory.MapView(firstPageMemory, firstPageOffset, Size, _hostPageSize);
+
+                    if (!useProtectionMirrors)
+                    {
+                        _baseMemory.Reprotect(Size, _hostPageSize, partitionAfter._firstPageMemoryProtection, throwOnFail: false);
+                    }
+
                     _firstPageMemoryForUnmap = firstPageMemory;
                     _firstPageOffsetForLateMap = firstPageOffset;
                 }
@@ -350,6 +362,14 @@ namespace Ryujinx.Cpu.Jit
             }
         }
 
+        public void ReprotectBridge(MemoryPermission protection)
+        {
+            if (_firstPageMemoryForUnmap != null)
+            {
+                _baseMemory.Reprotect(Size, _hostPageSize, protection, throwOnFail: false);
+            }
+        }
+
         private (MemoryBlock, ulong) GetFirstPageMemoryAndOffset()
         {
             _treeLock.EnterReadLock();
diff --git a/src/Ryujinx.Cpu/Jit/AddressSpacePartitioned.cs b/src/Ryujinx.Cpu/Jit/AddressSpacePartitioned.cs
index 244dd0e75b..444076752a 100644
--- a/src/Ryujinx.Cpu/Jit/AddressSpacePartitioned.cs
+++ b/src/Ryujinx.Cpu/Jit/AddressSpacePartitioned.cs
@@ -99,7 +99,7 @@ namespace Ryujinx.Cpu.Jit
             {
                 while (va < endVa)
                 {
-                    AddressSpacePartition partition = FindPartition(va);
+                    AddressSpacePartition partition = FindPartitionWithIndex(va, out int partitionIndex);
 
                     if (partition == null)
                     {
@@ -117,6 +117,13 @@ namespace Ryujinx.Cpu.Jit
                     else
                     {
                         partition.ReprotectAligned(clampedVa, clampedEndVa - clampedVa, protection);
+
+                        if (clampedVa == partition.Address &&
+                            partitionIndex > 0 &&
+                            _partitions[partitionIndex - 1].EndAddress == partition.Address)
+                        {
+                            _partitions[partitionIndex - 1].ReprotectBridge(protection);
+                        }
                     }
 
                     va += clampedEndVa - clampedVa;
@@ -187,21 +194,21 @@ namespace Ryujinx.Cpu.Jit
             {
                 if (_partitions[partitionIndex - 1].EndAddress == _partitions[partitionIndex].Address)
                 {
-                    _partitions[partitionIndex - 1].InsertBridgeAtEnd(_partitions[partitionIndex]);
+                    _partitions[partitionIndex - 1].InsertBridgeAtEnd(_partitions[partitionIndex], _useProtectionMirrors);
                 }
                 else
                 {
-                    _partitions[partitionIndex - 1].InsertBridgeAtEnd(null);
+                    _partitions[partitionIndex - 1].InsertBridgeAtEnd(null, _useProtectionMirrors);
                 }
             }
 
             if (partitionIndex + 1 < _partitions.Count && _partitions[partitionIndex].EndAddress == _partitions[partitionIndex + 1].Address)
             {
-                _partitions[partitionIndex].InsertBridgeAtEnd(_partitions[partitionIndex + 1]);
+                _partitions[partitionIndex].InsertBridgeAtEnd(_partitions[partitionIndex + 1], _useProtectionMirrors);
             }
             else
             {
-                _partitions[partitionIndex].InsertBridgeAtEnd(null);
+                _partitions[partitionIndex].InsertBridgeAtEnd(null, _useProtectionMirrors);
             }
         }
 
@@ -241,6 +248,20 @@ namespace Ryujinx.Cpu.Jit
             return null;
         }
 
+        private AddressSpacePartition FindPartitionWithIndex(ulong va, out int index)
+        {
+            lock (_partitions)
+            {
+                index = FindPartitionIndexLocked(va);
+                if (index >= 0)
+                {
+                    return _partitions[index];
+                }
+            }
+
+            return null;
+        }
+
         private int FindPartitionIndexLocked(ulong va)
         {
             int left = 0;