sm: Fix atmosphere-extension interaction with official JIT sysmodule usage

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. */

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);

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;

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);