From 7e536f74ae824307376f7671d989424b4b5e7047 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Thu, 14 Oct 2021 00:40:43 -0700
Subject: [PATCH] strat: fix some mitm server management logic

---
 .../stratosphere/sf/hipc/sf_hipc_server_manager.hpp        | 4 ++--
 .../sf/hipc/sf_hipc_server_session_manager.hpp             | 6 ++++++
 .../source/sf/hipc/sf_hipc_server_manager.cpp              | 7 +++++++
 stratosphere/sm/source/impl/sm_service_manager.cpp         | 6 ++++--
 4 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/libraries/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_manager.hpp b/libraries/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_manager.hpp
index 6a36e9201..2a9879a72 100644
--- a/libraries/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_manager.hpp
+++ b/libraries/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_manager.hpp
@@ -289,7 +289,7 @@ namespace ams::sf::hipc {
             util::TypedStorage<ServerSession> m_session_storages[MaxSessions];
             bool m_session_allocated[MaxSessions];
             u8 m_pointer_buffer_storage[0x10 + (MaxSessions * ManagerOptions::PointerBufferSize)];
-            u8 m_saved_message_storage[0x10 + (MaxSessions * (ManagerOptions::CanDeferInvokeRequest ? hipc::TlsMessageBufferSize : 0))];
+            u8 m_saved_message_storage[0x10 + (MaxSessions * ((ManagerOptions::CanDeferInvokeRequest || ManagerOptions::CanManageMitmServers) ? hipc::TlsMessageBufferSize : 0))];
             uintptr_t m_pointer_buffers_start;
             uintptr_t m_saved_messages_start;
 
@@ -403,7 +403,7 @@ namespace ams::sf::hipc {
             }
 
             virtual cmif::PointerAndSize GetSessionSavedMessageBuffer(const ServerSession *session) const override final {
-                if constexpr (ManagerOptions::CanDeferInvokeRequest) {
+                if constexpr (ManagerOptions::CanDeferInvokeRequest || ManagerOptions::CanManageMitmServers) {
                     return this->GetObjectBySessionIndex(session, m_saved_messages_start, hipc::TlsMessageBufferSize);
                 } else {
                     return cmif::PointerAndSize();
diff --git a/libraries/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_session_manager.hpp b/libraries/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_session_manager.hpp
index bbcf0468b..af10c89c1 100644
--- a/libraries/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_session_manager.hpp
+++ b/libraries/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_session_manager.hpp
@@ -69,6 +69,12 @@ namespace ams::sf::hipc {
                 AMS_ABORT_UNLESS(util::GetReference(m_forward_service) != nullptr);
             }
 
+            ~ServerSession() {
+                if (m_has_forward_service) {
+                    util::DestroyAt(m_forward_service);
+                }
+            }
+
             ALWAYS_INLINE bool IsMitmSession() const {
                 return m_has_forward_service;
             }
diff --git a/libraries/libstratosphere/source/sf/hipc/sf_hipc_server_manager.cpp b/libraries/libstratosphere/source/sf/hipc/sf_hipc_server_manager.cpp
index 6717f4f6e..7e8f3cbb9 100644
--- a/libraries/libstratosphere/source/sf/hipc/sf_hipc_server_manager.cpp
+++ b/libraries/libstratosphere/source/sf/hipc/sf_hipc_server_manager.cpp
@@ -135,6 +135,13 @@ namespace ams::sf::hipc {
             if (!session->m_has_received) {
                 R_TRY(this->ReceiveRequest(session, tls_message));
                 session->m_has_received = true;
+
+                if (this->CanManageMitmServers()) {
+                    const cmif::PointerAndSize &saved_message = session->m_saved_message;
+                    AMS_ABORT_UNLESS(tls_message.GetSize() == saved_message.GetSize());
+
+                    std::memcpy(saved_message.GetPointer(), tls_message.GetPointer(), tls_message.GetSize());
+                }
             }
 
             R_TRY_CATCH(this->ProcessRequest(session, tls_message)) {
diff --git a/stratosphere/sm/source/impl/sm_service_manager.cpp b/stratosphere/sm/source/impl/sm_service_manager.cpp
index 2ae192cac..ddf219f76 100644
--- a/stratosphere/sm/source/impl/sm_service_manager.cpp
+++ b/stratosphere/sm/source/impl/sm_service_manager.cpp
@@ -624,9 +624,11 @@ namespace ams::sm::impl {
             R_TRY(ValidateAccessControl(AccessControlEntry(proc->access_control, proc->access_control_size), service, false, false));
         }
 
-        /* Get service info. Check to see if we need to defer this until later. */
+        /* Get service info/mitm info. */
         ServiceInfo *service_info = GetServiceInfo(service);
-        MitmInfo *mitm_info = GetMitmInfo(service_info);
+        MitmInfo *mitm_info = service_info != nullptr ? GetMitmInfo(service_info) : nullptr;
+
+        /* Check to see if we need to defer until later. */
         R_UNLESS(service_info != nullptr,                           tipc::ResultRequestDeferred());
         R_UNLESS(!ShouldDeferForInit(service),                      tipc::ResultRequestDeferred());
         R_UNLESS(!HasFutureMitmDeclaration(service),                tipc::ResultRequestDeferred());