From 2e7b6de195419ffd4284e9357b65d246da7d09a2 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Thu, 26 Apr 2018 20:50:27 -0600
Subject: [PATCH] Loader: Automatically unload NRRs on service close.

---
 .../loader/source/ldr_registration.cpp        | 44 +++++++++++++------
 .../loader/source/ldr_registration.hpp        |  2 +
 stratosphere/loader/source/ldr_ro_service.hpp |  6 +++
 3 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/stratosphere/loader/source/ldr_registration.cpp b/stratosphere/loader/source/ldr_registration.cpp
index c7246d5cd..7df25314d 100644
--- a/stratosphere/loader/source/ldr_registration.cpp
+++ b/stratosphere/loader/source/ldr_registration.cpp
@@ -7,8 +7,7 @@ static Registration::List g_registration_list = {0};
 static u64 g_num_registered = 1;
 
 Registration::Process *Registration::GetFreeProcess() {
-    unsigned int i;
-    for (i = 0; i < REGISTRATION_LIST_MAX; i++) {
+    for (unsigned int i = 0; i < REGISTRATION_LIST_MAX; i++) {
         if (!g_registration_list.processes[i].in_use) {
             return &g_registration_list.processes[i];
         }
@@ -17,24 +16,30 @@ Registration::Process *Registration::GetFreeProcess() {
 }
 
 Registration::Process *Registration::GetProcess(u64 index) {
-    unsigned int i;
-    for (i = 0; i < REGISTRATION_LIST_MAX && (!g_registration_list.processes[i].in_use || g_registration_list.processes[i].index != index); i++) {
+    for (unsigned int i = 0; i < REGISTRATION_LIST_MAX; i++) {
+        if (g_registration_list.processes[i].in_use && g_registration_list.processes[i].index == index) {
+            return &g_registration_list.processes[i];
+        }
     }
-    if (i >= REGISTRATION_LIST_MAX) {
-        return NULL;
-    }
-    return &g_registration_list.processes[i];
+    return NULL;
 }
 
 Registration::Process *Registration::GetProcessByProcessId(u64 pid) {
-    unsigned int i;
-    for (i = 0; i < REGISTRATION_LIST_MAX && (!g_registration_list.processes[i].in_use || g_registration_list.processes[i].process_id != pid); i++) {
-        
+    for (unsigned int i = 0; i < REGISTRATION_LIST_MAX; i++) {
+        if (g_registration_list.processes[i].in_use && g_registration_list.processes[i].process_id == pid) {
+            return &g_registration_list.processes[i];
+        }
     }
-    if (i >= REGISTRATION_LIST_MAX) {
-        return NULL;
+    return NULL;
+}
+
+Registration::Process *Registration::GetProcessByRoService(void *service) {
+    for (unsigned int i = 0; i < REGISTRATION_LIST_MAX; i++) {
+        if (g_registration_list.processes[i].in_use && g_registration_list.processes[i].owner_ro_service == service) {
+            return &g_registration_list.processes[i];
+        }
     }
-    return &g_registration_list.processes[i];
+    return NULL;
 }
 
 bool Registration::RegisterTidSid(const TidSid *tid_sid, u64 *out_index) { 
@@ -103,6 +108,17 @@ void Registration::AddNsoInfo(u64 index, u64 base_address, u64 size, const unsig
     }
 }
 
+void Registration::CloseRoService(void *service, Handle process_h) {
+    Registration::Process *target_process = GetProcessByRoService(service);
+    if (target_process == NULL) {
+        return;
+    }
+    for (unsigned int i = 0; i < NRR_INFO_MAX; i++) {
+        if (target_process->nrr_infos[i].IsActive() && target_process->nrr_infos[i].process_handle == process_h) {
+            target_process->nrr_infos[i].Close();
+        }
+    }
+}
 
 Result Registration::AddNrrInfo(u64 index, MappedCodeMemory *nrr_info) {
     Registration::Process *target_process = GetProcess(index);
diff --git a/stratosphere/loader/source/ldr_registration.hpp b/stratosphere/loader/source/ldr_registration.hpp
index 8492248f0..f29d9e7b3 100644
--- a/stratosphere/loader/source/ldr_registration.hpp
+++ b/stratosphere/loader/source/ldr_registration.hpp
@@ -46,11 +46,13 @@ class Registration {
         static Registration::Process *GetFreeProcess();
         static Registration::Process *GetProcess(u64 index);
         static Registration::Process *GetProcessByProcessId(u64 pid);
+        static Registration::Process *GetProcessByRoService(void *service);
         static Result GetRegisteredTidSid(u64 index, Registration::TidSid *out);
         static bool RegisterTidSid(const TidSid *tid_sid, u64 *out_index);
         static bool UnregisterIndex(u64 index);
         static void SetProcessIdTidMinAndIs64BitAddressSpace(u64 index, u64 process_id, u64 tid_min, bool is_64_bit_addspace);
         static void AddNsoInfo(u64 index, u64 base_address, u64 size, const unsigned char *build_id);
+        static void CloseRoService(void *service, Handle process_h);
         static Result AddNrrInfo(u64 index, MappedCodeMemory *nrr_info);
         static Result RemoveNrrInfo(u64 index, u64 base_address);
         static Result GetNsoInfosForProcessId(NsoInfo *out, u32 max_out, u64 process_id, u32 *num_written);
diff --git a/stratosphere/loader/source/ldr_ro_service.hpp b/stratosphere/loader/source/ldr_ro_service.hpp
index f2a02f252..f47d9849b 100644
--- a/stratosphere/loader/source/ldr_ro_service.hpp
+++ b/stratosphere/loader/source/ldr_ro_service.hpp
@@ -18,6 +18,12 @@ class RelocatableObjectsService : IServiceObject {
     bool has_initialized;
     public:
         RelocatableObjectsService() : process_handle(0), process_id(U64_MAX), has_initialized(false) { }
+        ~RelocatableObjectsService() {
+            Registration::CloseRoService(this, this->process_handle);
+            if (this->has_initialized) {
+                svcCloseHandle(this->process_handle);
+            }
+        }
         virtual Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size);
         virtual Result handle_deferred() {
             /* This service will never defer. */