From 909397233c6abe04b72737a0e12678c107176321 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Thu, 17 Sep 2020 08:24:47 -0700
Subject: [PATCH] sm: Fix atmosphere-extension interaction with official JIT
 sysmodule usage

---
 .../sm/source/impl/sm_service_manager.cpp     | 28 +++++++++++++++++++
 .../sm/source/impl/sm_service_manager.hpp     |  3 ++
 stratosphere/sm/source/sm_user_service.cpp    |  6 ++++
 stratosphere/sm/source/sm_user_service.hpp    |  2 ++
 4 files changed, 39 insertions(+)

diff --git a/stratosphere/sm/source/impl/sm_service_manager.cpp b/stratosphere/sm/source/impl/sm_service_manager.cpp
index 96be1cbac..b773ddab4 100644
--- a/stratosphere/sm/source/impl/sm_service_manager.cpp
+++ b/stratosphere/sm/source/impl/sm_service_manager.cpp
@@ -350,6 +350,12 @@ namespace ams::sm::impl {
             return service == ServiceName::Encode("fsp-srv");
         }
 
+        bool ShouldCloseOnClientDisconnect(ServiceName service) {
+            /* jit sysmodule is closed and relaunched by am for each application which uses it. */
+            constexpr auto JitU = ServiceName::Encode("jit:u");
+            return service == JitU;
+        }
+
         Result GetMitmServiceHandleImpl(Handle *out, ServiceInfo *service_info, const MitmProcessInfo &client_info) {
             /* Send command to query if we should mitm. */
             bool should_mitm;
@@ -419,6 +425,28 @@ namespace ams::sm::impl {
 
             return ResultSuccess();
         }
+
+    }
+
+    /* Client disconnection callback. */
+    void OnClientDisconnected(os::ProcessId process_id) {
+        /* Ensure that the process id is valid. */
+        if (process_id == os::InvalidProcessId) {
+            return;
+        }
+
+        /* NOTE: Nintendo unregisters all services a process hosted on client close. */
+        /* We do not do this as an atmosphere extension, in order to reduce the number */
+        /* of sessions open at any given time. */
+        /* However, certain system behavior (jit) relies on this occurring. */
+        /* As such, we will special case the system components which rely on the behavior. */
+        for (size_t i = 0; i < ServiceCountMax; i++) {
+            if (g_service_list[i].name != InvalidServiceName && g_service_list[i].owner_process_id == process_id) {
+                if (ShouldCloseOnClientDisconnect(g_service_list[i].name)) {
+                    g_service_list[i].Free();
+                }
+            }
+        }
     }
 
     /* Process management. */
diff --git a/stratosphere/sm/source/impl/sm_service_manager.hpp b/stratosphere/sm/source/impl/sm_service_manager.hpp
index e27a18c86..cc628eec6 100644
--- a/stratosphere/sm/source/impl/sm_service_manager.hpp
+++ b/stratosphere/sm/source/impl/sm_service_manager.hpp
@@ -19,6 +19,9 @@
 
 namespace ams::sm::impl {
 
+    /* Client disconnection callback. */
+    void OnClientDisconnected(os::ProcessId process_id);
+
     /* Process management. */
     Result RegisterProcess(os::ProcessId process_id, ncm::ProgramId program_id, cfg::OverrideStatus, const void *acid_sac, size_t acid_sac_size, const void *aci_sac, size_t aci_sac_size);
     Result UnregisterProcess(os::ProcessId process_id);
diff --git a/stratosphere/sm/source/sm_user_service.cpp b/stratosphere/sm/source/sm_user_service.cpp
index bf5eb4162..de9c8b6e2 100644
--- a/stratosphere/sm/source/sm_user_service.cpp
+++ b/stratosphere/sm/source/sm_user_service.cpp
@@ -19,6 +19,12 @@
 
 namespace ams::sm {
 
+    UserService::~UserService() {
+        if (this->has_initialized) {
+            impl::OnClientDisconnected(this->process_id);
+        }
+    }
+
     Result UserService::RegisterClient(const sf::ClientProcessId &client_process_id) {
         this->process_id = client_process_id.GetValue();
         this->has_initialized = true;
diff --git a/stratosphere/sm/source/sm_user_service.hpp b/stratosphere/sm/source/sm_user_service.hpp
index 1bf82caad..dfb64b530 100644
--- a/stratosphere/sm/source/sm_user_service.hpp
+++ b/stratosphere/sm/source/sm_user_service.hpp
@@ -26,6 +26,8 @@ namespace ams::sm {
             bool has_initialized = false;
         private:
             Result EnsureInitialized();
+        public:
+            virtual ~UserService();
         public:
             /* Official commands. */
             Result RegisterClient(const sf::ClientProcessId &client_process_id);