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

* Fix session service disposal and improve transfer memory implementation

* Remove useless assignment

Ryujinx.HLE/HOS/Kernel/Ipc/KClientSession.cs

@@ -2,6 +2,7 @@ using Ryujinx.HLE.HOS.Kernel.Common;
 using Ryujinx.HLE.HOS.Kernel.Process;
 using Ryujinx.HLE.HOS.Kernel.Threading;
 using Ryujinx.HLE.HOS.Services;
+using System;
 
 namespace Ryujinx.HLE.HOS.Kernel.Ipc
 {
@@ -83,6 +84,11 @@ namespace Ryujinx.HLE.HOS.Kernel.Ipc
         {
             _parent.DisconnectClient();
             _parent.DecrementReferenceCount();
+
+            if (Service is IDisposable disposableObj)
+            {
+                disposableObj.Dispose();
+            }
         }
     }
 }

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,

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();
+            }
         }
     }
 }

Ryujinx.HLE/HOS/Kernel/SupervisorCall/Syscall.cs

@@ -991,16 +991,41 @@ namespace Ryujinx.HLE.HOS.Kernel.SupervisorCall
 
             KProcess process = _context.Scheduler.GetCurrentProcess();
 
-            KernelResult result = process.MemoryManager.ReserveTransferMemory(address, size, permission);
+            KResourceLimit resourceLimit = process.ResourceLimit;
+
+            if (resourceLimit != null && !resourceLimit.Reserve(LimitableResource.TransferMemory, 1))
+            {
+                return KernelResult.ResLimitExceeded;
+            }
+
+            void CleanUpForError()
+            {
+                resourceLimit?.Release(LimitableResource.TransferMemory, 1);
+            }
+
+            if (!process.MemoryManager.InsideAddrSpace(address, size))
+            {
+                CleanUpForError();
+
+                return KernelResult.InvalidMemState;
+            }
+
+            KTransferMemory transferMemory = new KTransferMemory(_context);
+
+            KernelResult result = transferMemory.Initialize(address, size, permission);
 
             if (result != KernelResult.Success)
             {
+                CleanUpForError();
+
                 return result;
             }
 
-            KTransferMemory transferMemory = new KTransferMemory(_context, address, size);
+            result = process.HandleTable.GenerateHandle(transferMemory, out handle);
+
+            transferMemory.DecrementReferenceCount();
 
-            return process.HandleTable.GenerateHandle(transferMemory, out handle);
+            return result;
         }
 
         public KernelResult MapPhysicalMemory(ulong address, ulong size)
@@ -1271,29 +1296,9 @@ namespace Ryujinx.HLE.HOS.Kernel.SupervisorCall
 
         public KernelResult CloseHandle(int handle)
         {
-            KProcess process = _context.Scheduler.GetCurrentProcess();
-
-            KAutoObject obj = process.HandleTable.GetObject<KAutoObject>(handle);
-
-            process.HandleTable.CloseHandle(handle);
-
-            if (obj == null)
-            {
-                return KernelResult.InvalidHandle;
-            }
-
-            if (obj is KSession session)
-            {
-                session.Dispose();
-            }
-            else if (obj is KTransferMemory transferMemory)
-            {
-                process.MemoryManager.ResetTransferMemory(
-                    transferMemory.Address,
-                    transferMemory.Size);
-            }
+            KProcess currentProcess = _context.Scheduler.GetCurrentProcess();
 
-            return KernelResult.Success;
+            return currentProcess.HandleTable.CloseHandle(handle) ? KernelResult.Success : KernelResult.InvalidHandle;
         }
 
         public KernelResult ResetSignal(int handle)

Ryujinx.HLE/HOS/Services/Sm/IUserInterface.cs

@@ -104,6 +104,9 @@ namespace Ryujinx.HLE.HOS.Services.Sm
                 throw new InvalidOperationException("Out of handles!");
             }
 
+            session.ServerSession.DecrementReferenceCount();
+            session.ClientSession.DecrementReferenceCount();
+
             context.Response.HandleDesc = IpcHandleDesc.MakeMove(handle);
 
             return ResultCode.Success;