From fe8d031708038612f89a02b797fd30cafe2d3b47 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Tue, 27 Jul 2021 23:55:53 -0700
Subject: [PATCH] cs: fix allocator aborts

---
 .../stratosphere/pgl/pgl_event_observer.hpp   | 10 ++++-
 .../source/cs/cs_command_processor.cpp        |  1 +
 .../source/pgl/pgl_shell_api.cpp              | 45 +++++++++++++++++--
 stratosphere/cs/source/cs_main.cpp            |  2 +-
 4 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/libraries/libstratosphere/include/stratosphere/pgl/pgl_event_observer.hpp b/libraries/libstratosphere/include/stratosphere/pgl/pgl_event_observer.hpp
index ecfcbdcbc..692340c39 100644
--- a/libraries/libstratosphere/include/stratosphere/pgl/pgl_event_observer.hpp
+++ b/libraries/libstratosphere/include/stratosphere/pgl/pgl_event_observer.hpp
@@ -84,11 +84,17 @@ namespace ams::pgl {
     class EventObserver {
         NON_COPYABLE(EventObserver);
         private:
-            std::unique_ptr<impl::EventObserverInterface> m_impl;
+            struct Deleter {
+                void operator()(impl::EventObserverInterface *);
+            };
+        public:
+            using UniquePtr = std::unique_ptr<impl::EventObserverInterface, Deleter>;
+        private:
+            UniquePtr m_impl;
         public:
             EventObserver() { /* ... */ }
 
-            explicit EventObserver(std::unique_ptr<impl::EventObserverInterface> impl) : m_impl(std::move(impl)) { /* ... */ }
+            explicit EventObserver(UniquePtr impl) : m_impl(std::move(impl)) { /* ... */ }
 
             EventObserver(EventObserver &&rhs) {
                 m_impl = std::move(rhs.m_impl);
diff --git a/libraries/libstratosphere/source/cs/cs_command_processor.cpp b/libraries/libstratosphere/source/cs/cs_command_processor.cpp
index d00ae2a72..434433c9e 100644
--- a/libraries/libstratosphere/source/cs/cs_command_processor.cpp
+++ b/libraries/libstratosphere/source/cs/cs_command_processor.cpp
@@ -33,6 +33,7 @@ namespace ams::cs {
                 SendFirmwareVersion(socket, header);
                 break;
             /* TODO: Command support. */
+            /* TODO: Command support. */
             default:
                 scs::CommandProcessor::ProcessCommand(header, body, socket);
                 break;
diff --git a/libraries/libstratosphere/source/pgl/pgl_shell_api.cpp b/libraries/libstratosphere/source/pgl/pgl_shell_api.cpp
index 83e1f03fb..46c56d09d 100644
--- a/libraries/libstratosphere/source/pgl/pgl_shell_api.cpp
+++ b/libraries/libstratosphere/source/pgl/pgl_shell_api.cpp
@@ -18,6 +18,44 @@
 
 namespace ams::pgl {
 
+    namespace {
+
+        struct PglEventObserverAllocator;
+        using RemoteAllocator     = ams::sf::ExpHeapStaticAllocator<1_KB, PglEventObserverAllocator>;
+        using RemoteObjectFactory = ams::sf::ObjectFactory<typename RemoteAllocator::Policy>;
+
+        class StaticAllocatorInitializer {
+            public:
+                StaticAllocatorInitializer() {
+                    RemoteAllocator::Initialize(lmem::CreateOption_None);
+                }
+        } g_static_allocator_initializer;
+
+        template<typename T, typename... Args>
+        T *AllocateFromStaticExpHeap(Args &&... args) {
+            T * const object = static_cast<T *>(RemoteAllocator::Allocate(sizeof(T)));
+            if (AMS_LIKELY(object != nullptr)) {
+                std::construct_at(object, std::forward<Args>(args)...);
+            }
+            return object;
+        }
+
+        template<typename T>
+        void FreeToStaticExpHeap(T *object) {
+            return RemoteAllocator::Deallocate(object, sizeof(T));
+        }
+
+        template<typename T, typename... Args> requires std::derived_from<T, impl::EventObserverInterface>
+        EventObserver::UniquePtr MakeUniqueFromStaticExpHeap(Args &&... args) {
+            return EventObserver::UniquePtr{AllocateFromStaticExpHeap<T>(std::forward<Args>(args)...)};
+        }
+
+    }
+
+    void EventObserver::Deleter::operator()(impl::EventObserverInterface *obj) {
+        FreeToStaticExpHeap(obj);
+    }
+
     Result Initialize() {
         return ::pglInitialize();
     }
@@ -79,17 +117,16 @@ namespace ams::pgl {
         ::PglEventObserver obs;
         R_TRY(::pglGetEventObserver(std::addressof(obs)));
 
-        /* TODO: Real allocator */
         if (hos::GetVersion() >= hos::Version_12_0_0) {
-            auto observer_holder = std::make_unique<impl::EventObserverByTipc<RemoteEventObserver>>(obs);
+            auto observer_holder = MakeUniqueFromStaticExpHeap<impl::EventObserverByTipc<RemoteEventObserver>>(obs);
             R_UNLESS(observer_holder != nullptr, pgl::ResultOutOfMemory());
 
             *out = pgl::EventObserver(std::move(observer_holder));
         } else {
-            auto remote_observer = ams::sf::CreateSharedObjectEmplaced<pgl::sf::IEventObserver, RemoteEventObserver>(obs);
+            auto remote_observer = RemoteObjectFactory::CreateSharedEmplaced<pgl::sf::IEventObserver, RemoteEventObserver>(obs);
             R_UNLESS(remote_observer != nullptr, pgl::ResultOutOfMemory());
 
-            auto observer_holder = std::make_unique<impl::EventObserverByCmif>(std::move(remote_observer));
+            auto observer_holder = MakeUniqueFromStaticExpHeap<impl::EventObserverByCmif>(std::move(remote_observer));
             R_UNLESS(observer_holder != nullptr, pgl::ResultOutOfMemory());
 
             *out = pgl::EventObserver(std::move(observer_holder));
diff --git a/stratosphere/cs/source/cs_main.cpp b/stratosphere/cs/source/cs_main.cpp
index b72898285..86d882cb8 100644
--- a/stratosphere/cs/source/cs_main.cpp
+++ b/stratosphere/cs/source/cs_main.cpp
@@ -70,7 +70,7 @@ namespace ams::cs {
 
         alignas(os::MemoryPageSize) constinit u8 g_heap_memory[32_KB];
 
-        alignas(0x40) constinit u8 g_htcs_buffer[1_KB];
+        alignas(0x40) constinit u8 g_htcs_buffer[2_KB];
 
         constinit os::SdkMutex g_heap_mutex;
         constinit lmem::HeapHandle g_heap_handle;