8 files changed, 130 insertions, 73 deletions

diff --git a/src/core/hle/kernel/hle_ipc.cpp b/src/core/hle/kernel/hle_ipc.cpp
index 24700f7a5..689f58cf6 100644
--- a/src/core/hle/kernel/hle_ipc.cpp
+++ b/src/core/hle/kernel/hle_ipc.cpp
@@ -35,11 +35,11 @@ SessionRequestHandler::SessionRequestHandler() = default;
 SessionRequestHandler::~SessionRequestHandler() = default;
 
 void SessionRequestHandler::ClientConnected(KServerSession* session) {
-    session->SetHleHandler(shared_from_this());
+    session->SetSessionHandler(shared_from_this());
 }
 
 void SessionRequestHandler::ClientDisconnected(KServerSession* session) {
-    session->SetHleHandler(nullptr);
+    session->SetSessionHandler(nullptr);
 }
 
 HLERequestContext::HLERequestContext(KernelCore& kernel_, Core::Memory::Memory& memory_,
@@ -186,18 +186,6 @@ ResultCode HLERequestContext::WriteToOutgoingCommandBuffer(KThread& requesting_t
     auto& owner_process = *requesting_thread.GetOwnerProcess();
     auto& handle_table = owner_process.GetHandleTable();
 
-    // The data_size already includes the payload header, the padding and the domain header.
-    std::size_t size{};
-
-    if (IsTipc()) {
-        size = cmd_buf.size();
-    } else {
-        size = data_payload_offset + data_size - sizeof(IPC::DataPayloadHeader) / sizeof(u32) - 4;
-        if (Session()->IsDomain()) {
-            size -= sizeof(IPC::DomainMessageHeader) / sizeof(u32);
-        }
-    }
-
     for (auto& object : copy_objects) {
         Handle handle{};
         if (object) {
@@ -222,7 +210,7 @@ ResultCode HLERequestContext::WriteToOutgoingCommandBuffer(KThread& requesting_t
     if (Session()->IsDomain()) {
         current_offset = domain_offset - static_cast<u32>(domain_objects.size());
         for (const auto& object : domain_objects) {
-            server_session->AppendDomainRequestHandler(object);
+            server_session->AppendDomainHandler(object);
             cmd_buf[current_offset++] =
                 static_cast<u32_le>(server_session->NumDomainRequestHandlers());
         }
@@ -230,7 +218,7 @@ ResultCode HLERequestContext::WriteToOutgoingCommandBuffer(KThread& requesting_t
 
     // Copy the translated command buffer back into the thread's command buffer area.
     memory.WriteBlock(owner_process, requesting_thread.GetTLSAddress(), cmd_buf.data(),
-                      size * sizeof(u32));
+                      write_size * sizeof(u32));
 
     return RESULT_SUCCESS;
 }
diff --git a/src/core/hle/kernel/hle_ipc.h b/src/core/hle/kernel/hle_ipc.h
index e1b128281..51cd1a898 100644
--- a/src/core/hle/kernel/hle_ipc.h
+++ b/src/core/hle/kernel/hle_ipc.h
@@ -12,6 +12,8 @@
 #include <type_traits>
 #include <vector>
 #include <boost/container/small_vector.hpp>
+
+#include "common/assert.h"
 #include "common/common_types.h"
 #include "common/concepts.h"
 #include "common/swap.h"
@@ -84,6 +86,69 @@ public:
     void ClientDisconnected(KServerSession* session);
 };
 
+using SessionRequestHandlerPtr = std::shared_ptr<SessionRequestHandler>;
+
+/**
+ * Manages the underlying HLE requests for a session, and whether (or not) the session should be
+ * treated as a domain. This is managed separately from server sessions, as this state is shared
+ * when objects are cloned.
+ */
+class SessionRequestManager final {
+public:
+    SessionRequestManager() = default;
+
+    bool IsDomain() const {
+        return is_domain;
+    }
+
+    void ConvertToDomain() {
+        domain_handlers = {session_handler};
+        is_domain = true;
+    }
+
+    std::size_t DomainHandlerCount() const {
+        return domain_handlers.size();
+    }
+
+    bool HasSessionHandler() const {
+        return session_handler != nullptr;
+    }
+
+    SessionRequestHandler& SessionHandler() {
+        return *session_handler;
+    }
+
+    const SessionRequestHandler& SessionHandler() const {
+        return *session_handler;
+    }
+
+    void CloseDomainHandler(std::size_t index) {
+        if (index < DomainHandlerCount()) {
+            domain_handlers[index] = nullptr;
+        } else {
+            UNREACHABLE_MSG("Unexpected handler index {}", index);
+        }
+    }
+
+    SessionRequestHandlerPtr DomainHandler(std::size_t index) const {
+        ASSERT_MSG(index < DomainHandlerCount(), "Unexpected handler index {}", index);
+        return domain_handlers.at(index);
+    }
+
+    void AppendDomainHandler(SessionRequestHandlerPtr&& handler) {
+        domain_handlers.emplace_back(std::move(handler));
+    }
+
+    void SetSessionHandler(SessionRequestHandlerPtr&& handler) {
+        session_handler = std::move(handler);
+    }
+
+private:
+    bool is_domain{};
+    SessionRequestHandlerPtr session_handler;
+    std::vector<SessionRequestHandlerPtr> domain_handlers;
+};
+
 /**
  * Class containing information about an in-flight IPC request being handled by an HLE service
  * implementation. Services should avoid using old global APIs (e.g. Kernel::GetCommandBuffer()) and
@@ -239,18 +304,17 @@ public:
         copy_objects.emplace_back(object);
     }
 
-    void AddDomainObject(std::shared_ptr<SessionRequestHandler> object) {
+    void AddDomainObject(SessionRequestHandlerPtr object) {
         domain_objects.emplace_back(std::move(object));
     }
 
     template <typename T>
-    std::shared_ptr<T> GetDomainRequestHandler(std::size_t index) const {
-        return std::static_pointer_cast<T>(domain_request_handlers.at(index));
+    std::shared_ptr<T> GetDomainHandler(std::size_t index) const {
+        return std::static_pointer_cast<T>(manager->DomainHandler(index));
     }
 
-    void SetDomainRequestHandlers(
-        const std::vector<std::shared_ptr<SessionRequestHandler>>& handlers) {
-        domain_request_handlers = handlers;
+    void SetSessionRequestManager(std::shared_ptr<SessionRequestManager> manager_) {
+        manager = std::move(manager_);
     }
 
     /// Clears the list of objects so that no lingering objects are written accidentally to the
@@ -297,7 +361,7 @@ private:
     boost::container::small_vector<Handle, 8> copy_handles;
     boost::container::small_vector<KAutoObject*, 8> move_objects;
     boost::container::small_vector<KAutoObject*, 8> copy_objects;
-    boost::container::small_vector<std::shared_ptr<SessionRequestHandler>, 8> domain_objects;
+    boost::container::small_vector<SessionRequestHandlerPtr, 8> domain_objects;
 
     std::optional<IPC::CommandHeader> command_header;
     std::optional<IPC::HandleDescriptorHeader> handle_descriptor_header;
@@ -311,12 +375,12 @@ private:
 
     u32_le command{};
     u64 pid{};
+    u32 write_size{};
     u32 data_payload_offset{};
     u32 handles_offset{};
     u32 domain_offset{};
-    u32 data_size{};
 
-    std::vector<std::shared_ptr<SessionRequestHandler>> domain_request_handlers;
+    std::shared_ptr<SessionRequestManager> manager;
     bool is_thread_waiting{};
 
     KernelCore& kernel;
diff --git a/src/core/hle/kernel/k_client_port.h b/src/core/hle/kernel/k_client_port.h
index d00ce3ddd..8501156e8 100644
--- a/src/core/hle/kernel/k_client_port.h
+++ b/src/core/hle/kernel/k_client_port.h
@@ -31,6 +31,9 @@ public:
     const KPort* GetParent() const {
         return parent;
     }
+    KPort* GetParent() {
+        return parent;
+    }
 
     s32 GetNumSessions() const {
         return num_sessions;
diff --git a/src/core/hle/kernel/k_port.cpp b/src/core/hle/kernel/k_port.cpp
index feb2bb11f..723a1b571 100644
--- a/src/core/hle/kernel/k_port.cpp
+++ b/src/core/hle/kernel/k_port.cpp
@@ -56,11 +56,8 @@ ResultCode KPort::EnqueueSession(KServerSession* session) {
 
     R_UNLESS(state == State::Normal, ResultPortClosed);
 
-    if (server.HasHLEHandler()) {
-        server.GetHLEHandler()->ClientConnected(session);
-    } else {
-        server.EnqueueSession(session);
-    }
+    server.GetSessionRequestHandler()->ClientConnected(session);
+    server.EnqueueSession(session);
 
     return RESULT_SUCCESS;
 }
diff --git a/src/core/hle/kernel/k_server_port.h b/src/core/hle/kernel/k_server_port.h
index e76792253..d1a757ec3 100644
--- a/src/core/hle/kernel/k_server_port.h
+++ b/src/core/hle/kernel/k_server_port.h
@@ -32,26 +32,24 @@ public:
     explicit KServerPort(KernelCore& kernel_);
     virtual ~KServerPort() override;
 
-    using HLEHandler = std::shared_ptr<SessionRequestHandler>;
-
     void Initialize(KPort* parent_, std::string&& name_);
 
     /// Whether or not this server port has an HLE handler available.
-    bool HasHLEHandler() const {
-        return hle_handler != nullptr;
+    bool HasSessionRequestHandler() const {
+        return session_handler != nullptr;
     }
 
     /// Gets the HLE handler for this port.
-    HLEHandler GetHLEHandler() const {
-        return hle_handler;
+    SessionRequestHandlerPtr GetSessionRequestHandler() const {
+        return session_handler;
     }
 
     /**
      * Sets the HLE handler template for the port. ServerSessions crated by connecting to this port
      * will inherit a reference to this handler.
      */
-    void SetHleHandler(HLEHandler hle_handler_) {
-        hle_handler = std::move(hle_handler_);
+    void SetSessionHandler(SessionRequestHandlerPtr&& handler) {
+        session_handler = std::move(handler);
     }
 
     void EnqueueSession(KServerSession* pending_session);
@@ -73,7 +71,7 @@ private:
 
 private:
     SessionList session_list;
-    HLEHandler hle_handler;
+    SessionRequestHandlerPtr session_handler;
     KPort* parent{};
 };
 
diff --git a/src/core/hle/kernel/k_server_session.cpp b/src/core/hle/kernel/k_server_session.cpp
index 8850d9af5..457fdfd60 100644
--- a/src/core/hle/kernel/k_server_session.cpp
+++ b/src/core/hle/kernel/k_server_session.cpp
@@ -23,7 +23,8 @@
 
 namespace Kernel {
 
-KServerSession::KServerSession(KernelCore& kernel_) : KSynchronizationObject{kernel_} {}
+KServerSession::KServerSession(KernelCore& kernel_)
+    : KSynchronizationObject{kernel_}, manager{std::make_shared<SessionRequestManager>()} {}
 
 KServerSession::~KServerSession() {
     kernel.ReleaseServiceThread(service_thread);
@@ -43,14 +44,8 @@ void KServerSession::Destroy() {
 }
 
 void KServerSession::OnClientClosed() {
-    // We keep a shared pointer to the hle handler to keep it alive throughout
-    // the call to ClientDisconnected, as ClientDisconnected invalidates the
-    // hle_handler member itself during the course of the function executing.
-    std::shared_ptr<SessionRequestHandler> handler = hle_handler;
-    if (handler) {
-        // Note that after this returns, this server session's hle_handler is
-        // invalidated (set to null).
-        handler->ClientDisconnected(this);
+    if (manager->HasSessionHandler()) {
+        manager->SessionHandler().ClientDisconnected(this);
     }
 }
 
@@ -66,12 +61,12 @@ bool KServerSession::IsSignaled() const {
     return false;
 }
 
-void KServerSession::AppendDomainRequestHandler(std::shared_ptr<SessionRequestHandler> handler) {
-    domain_request_handlers.push_back(std::move(handler));
+void KServerSession::AppendDomainHandler(SessionRequestHandlerPtr handler) {
+    manager->AppendDomainHandler(std::move(handler));
 }
 
 std::size_t KServerSession::NumDomainRequestHandlers() const {
-    return domain_request_handlers.size();
+    return manager->DomainHandlerCount();
 }
 
 ResultCode KServerSession::HandleDomainSyncRequest(Kernel::HLERequestContext& context) {
@@ -80,14 +75,14 @@ ResultCode KServerSession::HandleDomainSyncRequest(Kernel::HLERequestContext& co
     }
 
     // Set domain handlers in HLE context, used for domain objects (IPC interfaces) as inputs
-    context.SetDomainRequestHandlers(domain_request_handlers);
+    context.SetSessionRequestManager(manager);
 
     // If there is a DomainMessageHeader, then this is CommandType "Request"
     const auto& domain_message_header = context.GetDomainMessageHeader();
     const u32 object_id{domain_message_header.object_id};
     switch (domain_message_header.command) {
     case IPC::DomainMessageHeader::CommandType::SendMessage:
-        if (object_id > domain_request_handlers.size()) {
+        if (object_id > manager->DomainHandlerCount()) {
             LOG_CRITICAL(IPC,
                          "object_id {} is too big! This probably means a recent service call "
                          "to {} needed to return a new interface!",
@@ -95,12 +90,12 @@ ResultCode KServerSession::HandleDomainSyncRequest(Kernel::HLERequestContext& co
             UNREACHABLE();
             return RESULT_SUCCESS; // Ignore error if asserts are off
         }
-        return domain_request_handlers[object_id - 1]->HandleSyncRequest(*this, context);
+        return manager->DomainHandler(object_id - 1)->HandleSyncRequest(*this, context);
 
     case IPC::DomainMessageHeader::CommandType::CloseVirtualHandle: {
         LOG_DEBUG(IPC, "CloseVirtualHandle, object_id=0x{:08X}", object_id);
 
-        domain_request_handlers[object_id - 1] = nullptr;
+        manager->CloseDomainHandler(object_id - 1);
 
         IPC::ResponseBuilder rb{context, 2};
         rb.Push(RESULT_SUCCESS);
@@ -133,14 +128,14 @@ ResultCode KServerSession::CompleteSyncRequest(HLERequestContext& context) {
     if (IsDomain() && context.HasDomainMessageHeader()) {
         result = HandleDomainSyncRequest(context);
         // If there is no domain header, the regular session handler is used
-    } else if (hle_handler != nullptr) {
+    } else if (manager->HasSessionHandler()) {
         // If this ServerSession has an associated HLE handler, forward the request to it.
-        result = hle_handler->HandleSyncRequest(*this, context);
+        result = manager->SessionHandler().HandleSyncRequest(*this, context);
     }
 
     if (convert_to_domain) {
-        ASSERT_MSG(IsSession(), "ServerSession is already a domain instance.");
-        domain_request_handlers = {hle_handler};
+        ASSERT_MSG(!IsDomain(), "ServerSession is already a domain instance.");
+        manager->ConvertToDomain();
         convert_to_domain = false;
     }
 
diff --git a/src/core/hle/kernel/k_server_session.h b/src/core/hle/kernel/k_server_session.h
index 597d76d38..dd4de2904 100644
--- a/src/core/hle/kernel/k_server_session.h
+++ b/src/core/hle/kernel/k_server_session.h
@@ -12,6 +12,7 @@
 #include <boost/intrusive/list.hpp>
 
 #include "common/threadsafe_queue.h"
+#include "core/hle/kernel/hle_ipc.h"
 #include "core/hle/kernel/k_synchronization_object.h"
 #include "core/hle/kernel/service_thread.h"
 #include "core/hle/result.h"
@@ -64,8 +65,8 @@ public:
      * instead of the regular IPC machinery. (The regular IPC machinery is currently not
      * implemented.)
      */
-    void SetHleHandler(std::shared_ptr<SessionRequestHandler> hle_handler_) {
-        hle_handler = std::move(hle_handler_);
+    void SetSessionHandler(SessionRequestHandlerPtr handler) {
+        manager->SetSessionHandler(std::move(handler));
     }
 
     /**
@@ -82,7 +83,7 @@ public:
 
     /// Adds a new domain request handler to the collection of request handlers within
     /// this ServerSession instance.
-    void AppendDomainRequestHandler(std::shared_ptr<SessionRequestHandler> handler);
+    void AppendDomainHandler(SessionRequestHandlerPtr handler);
 
     /// Retrieves the total number of domain request handlers that have been
     /// appended to this ServerSession instance.
@@ -90,12 +91,7 @@ public:
 
     /// Returns true if the session has been converted to a domain, otherwise False
     bool IsDomain() const {
-        return !IsSession();
-    }
-
-    /// Returns true if this session has not been converted to a domain, otherwise false.
-    bool IsSession() const {
-        return domain_request_handlers.empty();
+        return manager->IsDomain();
     }
 
     /// Converts the session to a domain at the end of the current command
@@ -103,6 +99,21 @@ public:
         convert_to_domain = true;
     }
 
+    /// Gets the session request manager, which forwards requests to the underlying service
+    std::shared_ptr<SessionRequestManager>& GetSessionRequestManager() {
+        return manager;
+    }
+
+    /// Gets the session request manager, which forwards requests to the underlying service
+    const std::shared_ptr<SessionRequestManager>& GetSessionRequestManager() const {
+        return manager;
+    }
+
+    /// Sets the session request manager, which forwards requests to the underlying service
+    void SetSessionRequestManager(std::shared_ptr<SessionRequestManager> manager_) {
+        manager = std::move(manager_);
+    }
+
 private:
     /// Queues a sync request from the emulated application.
     ResultCode QueueSyncRequest(KThread* thread, Core::Memory::Memory& memory);
@@ -114,11 +125,8 @@ private:
     /// object handle.
     ResultCode HandleDomainSyncRequest(Kernel::HLERequestContext& context);
 
-    /// This session's HLE request handler (applicable when not a domain)
-    std::shared_ptr<SessionRequestHandler> hle_handler;
-
-    /// This is the list of domain request handlers (after conversion to a domain)
-    std::vector<std::shared_ptr<SessionRequestHandler>> domain_request_handlers;
+    /// This session's HLE request handlers
+    std::shared_ptr<SessionRequestManager> manager;
 
     /// When set to True, converts the session to a domain at the end of the command
     bool convert_to_domain{};
diff --git a/src/core/hle/kernel/k_session.h b/src/core/hle/kernel/k_session.h
index 16901e19c..a981fd1f6 100644
--- a/src/core/hle/kernel/k_session.h
+++ b/src/core/hle/kernel/k_session.h
@@ -66,6 +66,10 @@ public:
         return port;
     }
 
+    KClientPort* GetParent() {
+        return port;
+    }
+
 private:
     enum class State : u8 {
         Invalid = 0,