Fix session service disposal and improve transfer memory implementation (#1397)

* Fix session service disposal and improve transfer memory implementation

* Remove useless assignment

2 files changed, 83 insertions, 87 deletions

diff --git a/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryManager.cs b/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryManager.cs
index b13e2841..43e7ad69 100644
--- a/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryManager.cs
+++ b/Ryujinx.HLE/HOS/Kernel/Memory/KMemoryManager.cs
@@ -1129,82 +1129,6 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
             }
         }
 
-        public KernelResult ReserveTransferMemory(ulong address, ulong size, MemoryPermission permission)
-        {
-            lock (_blocks)
-            {
-                if (CheckRange(
-                    address,
-                    size,
-                    MemoryState.TransferMemoryAllowed | MemoryState.IsPoolAllocated,
-                    MemoryState.TransferMemoryAllowed | MemoryState.IsPoolAllocated,
-                    MemoryPermission.Mask,
-                    MemoryPermission.ReadAndWrite,
-                    MemoryAttribute.Mask,
-                    MemoryAttribute.None,
-                    MemoryAttribute.IpcAndDeviceMapped,
-                    out MemoryState state,
-                    out _,
-                    out MemoryAttribute attribute))
-                {
-                    // TODO: Missing checks.
-
-                    if (!_blockAllocator.CanAllocate(MaxBlocksNeededForInsertion))
-                    {
-                        return KernelResult.OutOfResource;
-                    }
-
-                    ulong pagesCount = size / PageSize;
-
-                    attribute |= MemoryAttribute.Borrowed;
-
-                    InsertBlock(address, pagesCount, state, permission, attribute);
-
-                    return KernelResult.Success;
-                }
-                else
-                {
-                    return KernelResult.InvalidMemState;
-                }
-            }
-        }
-
-        public KernelResult ResetTransferMemory(ulong address, ulong size)
-        {
-            lock (_blocks)
-            {
-                if (CheckRange(
-                    address,
-                    size,
-                    MemoryState.TransferMemoryAllowed | MemoryState.IsPoolAllocated,
-                    MemoryState.TransferMemoryAllowed | MemoryState.IsPoolAllocated,
-                    MemoryPermission.None,
-                    MemoryPermission.None,
-                    MemoryAttribute.Mask,
-                    MemoryAttribute.Borrowed,
-                    MemoryAttribute.IpcAndDeviceMapped,
-                    out MemoryState state,
-                    out _,
-                    out _))
-                {
-                    if (!_blockAllocator.CanAllocate(MaxBlocksNeededForInsertion))
-                    {
-                        return KernelResult.OutOfResource;
-                    }
-
-                    ulong pagesCount = size / PageSize;
-
-                    InsertBlock(address, pagesCount, state, MemoryPermission.ReadAndWrite);
-
-                    return KernelResult.Success;
-                }
-                else
-                {
-                    return KernelResult.InvalidMemState;
-                }
-            }
-        }
-
         public KernelResult SetProcessMemoryPermission(ulong address, ulong size, MemoryPermission permission)
         {
             lock (_blocks)
@@ -2195,6 +2119,22 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                 MemoryAttribute.Borrowed);
         }
 
+        public KernelResult BorrowTransferMemory(KPageList pageList, ulong address, ulong size, MemoryPermission permission)
+        {
+            return SetAttributesAndChangePermission(
+                address,
+                size,
+                MemoryState.TransferMemoryAllowed,
+                MemoryState.TransferMemoryAllowed,
+                MemoryPermission.Mask,
+                MemoryPermission.ReadAndWrite,
+                MemoryAttribute.Mask,
+                MemoryAttribute.None,
+                permission,
+                MemoryAttribute.Borrowed,
+                pageList);
+        }
+
         private KernelResult SetAttributesAndChangePermission(
             ulong            address,
             ulong            size,
@@ -2233,14 +2173,7 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
 
                     if (pageList != null)
                     {
-                        KPageList currPageList = new KPageList();
-
-                        AddVaRangeToPageList(currPageList, address, pagesCount);
-
-                        if (!currPageList.IsEqual(pageList))
-                        {
-                            return KernelResult.InvalidMemRange;
-                        }
+                        AddVaRangeToPageList(pageList, address, pagesCount);
                     }
 
                     if (!_blockAllocator.CanAllocate(MaxBlocksNeededForInsertion))
@@ -2297,6 +2230,22 @@ namespace Ryujinx.HLE.HOS.Kernel.Memory
                 MemoryAttribute.Borrowed);
         }
 
+        public KernelResult UnborrowTransferMemory(ulong address, ulong size, KPageList pageList)
+        {
+            return ClearAttributesAndChangePermission(
+                address,
+                size,
+                MemoryState.TransferMemoryAllowed,
+                MemoryState.TransferMemoryAllowed,
+                MemoryPermission.None,
+                MemoryPermission.None,
+                MemoryAttribute.Mask,
+                MemoryAttribute.Borrowed,
+                MemoryPermission.ReadAndWrite,
+                MemoryAttribute.Borrowed,
+                pageList);
+        }
+
         private KernelResult ClearAttributesAndChangePermission(
             ulong            address,
             ulong            size,
diff --git a/Ryujinx.HLE/HOS/Kernel/Memory/KTransferMemory.cs b/Ryujinx.HLE/HOS/Kernel/Memory/KTransferMemory.cs
index 6da0c405..d3e6208e 100644
--- a/Ryujinx.HLE/HOS/Kernel/Memory/KTransferMemory.cs
+++ b/Ryujinx.HLE/HOS/Kernel/Memory/KTransferMemory.cs
@@ -1,16 +1,63 @@
 using Ryujinx.HLE.HOS.Kernel.Common;
+using Ryujinx.HLE.HOS.Kernel.Process;
+using System;
 
 namespace Ryujinx.HLE.HOS.Kernel.Memory
 {
     class KTransferMemory : KAutoObject
     {
+        private KProcess _creator;
+
+        private readonly KPageList _pageList;
+
         public ulong Address { get; private set; }
-        public ulong Size    { get; private set; }
+        public ulong Size => _pageList.GetPagesCount() * KMemoryManager.PageSize;
+
+        public MemoryPermission Permission { get; private set; }
 
-        public KTransferMemory(KernelContext context, ulong address, ulong size) : base(context)
+        private bool _hasBeenInitialized;
+        private bool _isMapped;
+
+        public KTransferMemory(KernelContext context) : base(context)
         {
+            _pageList = new KPageList();
+        }
+
+        public KernelResult Initialize(ulong address, ulong size, MemoryPermission permission)
+        {
+            KProcess creator = KernelContext.Scheduler.GetCurrentProcess();
+
+            _creator = creator;
+
+            KernelResult result = creator.MemoryManager.BorrowTransferMemory(_pageList, address, size, permission);
+
+            if (result != KernelResult.Success)
+            {
+                return result;
+            }
+
+            creator.IncrementReferenceCount();
+
+            Permission = permission;
             Address = address;
-            Size    = size;
+            _hasBeenInitialized = true;
+            _isMapped = false;
+
+            return result;
+        }
+
+        protected override void Destroy()
+        {
+            if (_hasBeenInitialized)
+            {
+                if (!_isMapped && _creator.MemoryManager.UnborrowTransferMemory(Address, Size, _pageList) != KernelResult.Success)
+                {
+                    throw new InvalidOperationException("Unexpected failure restoring transfer memory attributes.");
+                }
+
+                _creator.ResourceLimit?.Release(LimitableResource.TransferMemory, 1);
+                _creator.DecrementReferenceCount();
+            }
         }
     }
 }
\ No newline at end of file