From d219884bfa2a0c5bcb7935929bd17b528f24cc0e Mon Sep 17 00:00:00 2001
From: MasaGratoR <masagrator@gmail.com>
Date: Fri, 14 Apr 2023 18:33:16 +0200
Subject: [PATCH] tmem: Add tmemCloseHandle(), tmemWaitForPermission(), use
 those functions in nv.c to fix a race condition (#606)

---
 nx/include/switch/kernel/tmem.h | 15 +++++++++++
 nx/source/kernel/tmem.c         | 44 ++++++++++++++++++++++++++++++---
 nx/source/services/nv.c         |  6 +++++
 3 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/nx/include/switch/kernel/tmem.h b/nx/include/switch/kernel/tmem.h
index a5993ea3f1..ff39206a97 100644
--- a/nx/include/switch/kernel/tmem.h
+++ b/nx/include/switch/kernel/tmem.h
@@ -72,6 +72,21 @@ static inline void* tmemGetAddr(TransferMemory* t){
     return t->map_addr;
 }
 
+/**
+ * @brief Closes handle of a transfer memory object.
+ * @param t Transfer memory information structure.
+ * @return Result code.
+ */
+Result tmemCloseHandle(TransferMemory* t);
+
+/**
+ * @brief Waits until source backing memory permissions match perm.
+ * @param t Transfer memory information structure.
+ * @param perm Permissions which the source backing memory is expected to have before return.
+ * @return Result code.
+ */
+Result tmemWaitForPermission(TransferMemory* t, Permission perm);
+
 /**
  * @brief Frees up resources used by a transfer memory object, unmapping and closing handles, etc.
  * @param t Transfer memory information structure.
diff --git a/nx/source/kernel/tmem.c b/nx/source/kernel/tmem.c
index cbf1aeb1e7..e30dc75cb9 100644
--- a/nx/source/kernel/tmem.c
+++ b/nx/source/kernel/tmem.c
@@ -100,6 +100,45 @@ Result tmemUnmap(TransferMemory* t)
     return rc;
 }
 
+Result tmemCloseHandle(TransferMemory* t)
+{
+    Result rc = 0;
+
+    if (t->handle != INVALID_HANDLE) {
+        rc = svcCloseHandle(t->handle);
+        t->handle = INVALID_HANDLE;
+    }
+
+    return rc;
+}
+
+Result tmemWaitForPermission(TransferMemory* t, Permission perm)
+{
+    Result rc = 0;
+
+    if ((t->perm & perm) != perm) {
+        MemoryInfo m = {0};
+        u32 p = 0;
+        rc = svcQueryMemory(&m, &p, (u64)(t->src_addr));
+
+        if (R_FAILED(rc)) {
+            return rc;
+        }
+
+        while ((m.perm & perm) != perm) {
+            rc = svcQueryMemory(&m, &p, (u64)(t->src_addr));
+
+            if (R_FAILED(rc)) {
+                return rc;
+            }
+
+            svcSleepThread(100000);
+        }
+    }
+
+    return rc;
+}
+
 Result tmemClose(TransferMemory* t)
 {
     Result rc = 0;
@@ -109,16 +148,13 @@ Result tmemClose(TransferMemory* t)
     }
 
     if (R_SUCCEEDED(rc)) {
-        if (t->handle != INVALID_HANDLE) {
-            rc = svcCloseHandle(t->handle);
-        }
+        rc = tmemCloseHandle(t);
 
         if (t->src_addr != NULL) {
             __libnx_free(t->src_addr);
         }
 
         t->src_addr = NULL;
-        t->handle = INVALID_HANDLE;
     }
 
     return rc;
diff --git a/nx/source/services/nv.c b/nx/source/services/nv.c
index 037619abfe..cccfd355da 100644
--- a/nx/source/services/nv.c
+++ b/nx/source/services/nv.c
@@ -55,6 +55,11 @@ Result _nvInitialize(void) {
         if (R_SUCCEEDED(rc))
             rc = _nvCmdInitialize(CUR_PROCESS_HANDLE, g_nvTransfermem.handle, tmem_size);
 
+        Result rc2 = tmemCloseHandle(&g_nvTransfermem);
+        
+        if (R_SUCCEEDED(rc))
+            rc = rc2;
+
         // Clone the session handle - the cloned session is used to execute certain commands in parallel
         if (R_SUCCEEDED(rc))
             rc = serviceCloneEx(&g_nvSrv, 1, &g_nvSrvClone);
@@ -73,6 +78,7 @@ Result _nvInitialize(void) {
 void _nvCleanup(void) {
     serviceClose(&g_nvSrvClone);
     serviceClose(&g_nvSrv);
+    tmemWaitForPermission(&g_nvTransfermem, Perm_Rw);
     tmemClose(&g_nvTransfermem);
 }