2 files changed, 41 insertions, 38 deletions

diff --git a/src/core/hle/service/sm/sm.cpp b/src/core/hle/service/sm/sm.cpp
index 15034abed..ae4dc4a75 100644
--- a/src/core/hle/service/sm/sm.cpp
+++ b/src/core/hle/service/sm/sm.cpp
@@ -4,6 +4,7 @@
 
 #include <tuple>
 #include "common/assert.h"
+#include "common/scope_exit.h"
 #include "core/core.h"
 #include "core/hle/ipc_helpers.h"
 #include "core/hle/kernel/k_client_port.h"
@@ -40,17 +41,13 @@ static ResultCode ValidateServiceName(const std::string& name) {
 }
 
 Kernel::KClientPort& ServiceManager::InterfaceFactory(ServiceManager& self, Core::System& system) {
-    ASSERT(self.sm_interface.expired());
-
-    auto sm = std::make_shared<SM>(self, system);
-    self.sm_interface = sm;
+    self.sm_interface = std::make_shared<SM>(self, system);
     self.controller_interface = std::make_unique<Controller>(system);
-
-    return sm->CreatePort();
+    return self.sm_interface->CreatePort();
 }
 
-ResultVal<Kernel::KServerPort*> ServiceManager::RegisterService(std::string name,
-                                                                u32 max_sessions) {
+ResultCode ServiceManager::RegisterService(std::string name, u32 max_sessions,
+                                           Kernel::SessionRequestHandlerPtr handler) {
 
     CASCADE_CODE(ValidateServiceName(name));
 
@@ -59,12 +56,9 @@ ResultVal<Kernel::KServerPort*> ServiceManager::RegisterService(std::string name
         return ERR_ALREADY_REGISTERED;
     }
 
-    auto* port = Kernel::KPort::Create(kernel);
-    port->Initialize(max_sessions, false, name);
+    registered_services.emplace(std::move(name), handler);
 
-    registered_services.emplace(std::move(name), port);
-
-    return MakeResult(&port->GetServerPort());
+    return ResultSuccess;
 }
 
 ResultCode ServiceManager::UnregisterService(const std::string& name) {
@@ -76,14 +70,11 @@ ResultCode ServiceManager::UnregisterService(const std::string& name) {
         return ERR_SERVICE_NOT_REGISTERED;
     }
 
-    iter->second->Close();
-
     registered_services.erase(iter);
     return ResultSuccess;
 }
 
 ResultVal<Kernel::KPort*> ServiceManager::GetServicePort(const std::string& name) {
-
     CASCADE_CODE(ValidateServiceName(name));
     auto it = registered_services.find(name);
     if (it == registered_services.end()) {
@@ -91,10 +82,13 @@ ResultVal<Kernel::KPort*> ServiceManager::GetServicePort(const std::string& name
         return ERR_SERVICE_NOT_REGISTERED;
     }
 
-    return MakeResult(it->second);
-}
+    auto* port = Kernel::KPort::Create(kernel);
+    port->Initialize(ServerSessionCountMax, false, name);
+    auto handler = it->second;
+    port->GetServerPort().SetSessionHandler(std::move(handler));
 
-SM::~SM() = default;
+    return MakeResult(port);
+}
 
 /**
  * SM::Initialize service function
@@ -156,11 +150,15 @@ ResultVal<Kernel::KClientSession*> SM::GetServiceImpl(Kernel::HLERequestContext&
         LOG_ERROR(Service_SM, "called service={} -> error 0x{:08X}", name, port_result.Code().raw);
         return port_result.Code();
     }
-    auto& port = port_result.Unwrap()->GetClientPort();
+    auto& port = port_result.Unwrap();
+    SCOPE_EXIT({ port->GetClientPort().Close(); });
+
+    server_ports.emplace_back(&port->GetServerPort());
 
     // Create a new session.
     Kernel::KClientSession* session{};
-    if (const auto result = port.CreateSession(std::addressof(session)); result.IsError()) {
+    if (const auto result = port->GetClientPort().CreateSession(std::addressof(session));
+        result.IsError()) {
         LOG_ERROR(Service_SM, "called service={} -> error 0x{:08X}", name, result.raw);
         return result;
     }
@@ -180,20 +178,21 @@ void SM::RegisterService(Kernel::HLERequestContext& ctx) {
     LOG_DEBUG(Service_SM, "called with name={}, max_session_count={}, is_light={}", name,
               max_session_count, is_light);
 
-    auto handle = service_manager.RegisterService(name, max_session_count);
-    if (handle.Failed()) {
-        LOG_ERROR(Service_SM, "failed to register service with error_code={:08X}",
-                  handle.Code().raw);
+    if (const auto result = service_manager.RegisterService(name, max_session_count, nullptr);
+        result.IsError()) {
+        LOG_ERROR(Service_SM, "failed to register service with error_code={:08X}", result.raw);
         IPC::ResponseBuilder rb{ctx, 2};
-        rb.Push(handle.Code());
+        rb.Push(result);
         return;
     }
 
-    IPC::ResponseBuilder rb{ctx, 2, 0, 1, IPC::ResponseBuilder::Flags::AlwaysMoveHandles};
-    rb.Push(handle.Code());
+    auto* port = Kernel::KPort::Create(kernel);
+    port->Initialize(ServerSessionCountMax, is_light, name);
+    SCOPE_EXIT({ port->GetClientPort().Close(); });
 
-    auto server_port = handle.Unwrap();
-    rb.PushMoveObjects(server_port);
+    IPC::ResponseBuilder rb{ctx, 2, 0, 1, IPC::ResponseBuilder::Flags::AlwaysMoveHandles};
+    rb.Push(ResultSuccess);
+    rb.PushMoveObjects(port->GetServerPort());
 }
 
 void SM::UnregisterService(Kernel::HLERequestContext& ctx) {
@@ -225,4 +224,10 @@ SM::SM(ServiceManager& service_manager_, Core::System& system_)
     });
 }
 
+SM::~SM() {
+    for (auto& server_port : server_ports) {
+        server_port->Close();
+    }
+}
+
 } // namespace Service::SM
diff --git a/src/core/hle/service/sm/sm.h b/src/core/hle/service/sm/sm.h
index ea37f11d4..068c78588 100644
--- a/src/core/hle/service/sm/sm.h
+++ b/src/core/hle/service/sm/sm.h
@@ -49,6 +49,7 @@ private:
     ServiceManager& service_manager;
     bool is_initialized{};
     Kernel::KernelCore& kernel;
+    std::vector<Kernel::KServerPort*> server_ports;
 };
 
 class ServiceManager {
@@ -58,7 +59,8 @@ public:
     explicit ServiceManager(Kernel::KernelCore& kernel_);
     ~ServiceManager();
 
-    ResultVal<Kernel::KServerPort*> RegisterService(std::string name, u32 max_sessions);
+    ResultCode RegisterService(std::string name, u32 max_sessions,
+                               Kernel::SessionRequestHandlerPtr handler);
     ResultCode UnregisterService(const std::string& name);
     ResultVal<Kernel::KPort*> GetServicePort(const std::string& name);
 
@@ -69,21 +71,17 @@ public:
             LOG_DEBUG(Service, "Can't find service: {}", service_name);
             return nullptr;
         }
-        auto* port = service->second;
-        if (port == nullptr) {
-            return nullptr;
-        }
-        return std::static_pointer_cast<T>(port->GetServerPort().GetSessionRequestHandler());
+        return std::static_pointer_cast<T>(service->second);
     }
 
     void InvokeControlRequest(Kernel::HLERequestContext& context);
 
 private:
-    std::weak_ptr<SM> sm_interface;
+    std::shared_ptr<SM> sm_interface;
     std::unique_ptr<Controller> controller_interface;
 
     /// Map of registered services, retrieved using GetServicePort.
-    std::unordered_map<std::string, Kernel::KPort*> registered_services;
+    std::unordered_map<std::string, Kernel::SessionRequestHandlerPtr> registered_services;
 
     /// Kernel context
     Kernel::KernelCore& kernel;