libstrat: implement mitm server registration

stratosphere/libstratosphere/include/stratosphere/sf/hipc/sf_hipc_server_manager.hpp

@@ -18,6 +18,12 @@
 #include "sf_hipc_server_domain_session_manager.hpp"
 #include "../../sm.hpp"
 
+namespace ams::ncm {
+
+    struct ProgramId;
+
+}
+
 namespace ams::sf::hipc {
 
     struct DefaultServerManagerOptions {
@@ -35,6 +41,8 @@ namespace ams::sf::hipc {
     class ServerManagerBase : public ServerDomainSessionManager {
         NON_COPYABLE(ServerManagerBase);
         NON_MOVEABLE(ServerManagerBase);
+        public:
+            using MitmQueryFunction = bool (*)(os::ProcessId, ncm::ProgramId);
         private:
             enum class UserDataTag : uintptr_t {
                 Server      = 1,
@@ -162,6 +170,8 @@ namespace ams::sf::hipc {
 
                 this->waitable_manager.LinkWaitableHolder(server);
             }
+
+            Result InstallMitmServerImpl(Handle *out_port_handle, sm::ServiceName service_name, MitmQueryFunction query_func);
         protected:
             virtual ServerBase *AllocateServer() = 0;
             virtual void DestroyServer(ServerBase *server)  = 0;
@@ -204,6 +214,18 @@ namespace ams::sf::hipc {
                 return ResultSuccess();
             }
 
+            template<typename ServiceImpl, auto MakeShared = std::make_shared<ServiceImpl>>
+            Result RegisterMitmServer(sm::ServiceName service_name, size_t max_sessions) {
+                static_assert(ServiceObjectTraits<ServiceImpl>::IsMitmServiceObject, "RegisterMitmServer requires mitm object. Use RegisterServer instead.");
+
+                /* Install mitm service. */
+                Handle port_handle;
+                R_TRY(this->InstallMitmServerImpl(&port_handle, service_name, max_sessions, &ServiceImpl::ShouldMitm));
+
+                this->RegisterServerImpl<ServiceImpl, MakeShared>(port_handle, service_name, true, cmif::ServiceObjectHolder());
+                return ResultSuccess();
+            }
+
             /* Processing. */
             os::WaitableHolder *WaitSignaled();
 
@@ -339,15 +361,15 @@ namespace ams::sf::hipc {
         public:
             ServerManager() : ServerManagerBase(this->domain_entry_storages, ManagerOptions::MaxDomainObjects) {
                 /* Clear storages. */
-                std::memset(this->server_storages, 0, sizeof(this->server_storages));
-                std::memset(this->server_allocated, 0, sizeof(this->server_allocated));
-                std::memset(this->session_storages, 0, sizeof(this->session_storages));
-                std::memset(this->session_allocated, 0, sizeof(this->session_allocated));
-                std::memset(this->pointer_buffer_storage, 0, sizeof(this->pointer_buffer_storage));
-                std::memset(this->saved_message_storage, 0, sizeof(this->saved_message_storage));
-                if constexpr (ManagerOptions::MaxDomains > 0) {
-                    std::memset(this->domain_allocated, 0, sizeof(this->domain_allocated));
-                }
+                #define SF_SM_MEMCLEAR(obj) if constexpr (sizeof(obj) > 0) { std::memset(obj, 0, sizeof(obj)); }
+                SF_SM_MEMCLEAR(this->server_storages);
+                SF_SM_MEMCLEAR(this->server_allocated);
+                SF_SM_MEMCLEAR(this->session_storages);
+                SF_SM_MEMCLEAR(this->session_allocated);
+                SF_SM_MEMCLEAR(this->pointer_buffer_storage);
+                SF_SM_MEMCLEAR(this->saved_message_storage);
+                SF_SM_MEMCLEAR(this->domain_allocated);
+                #undef SF_SM_MEMCLEAR
 
                 /* Set resource starts. */
                 this->pointer_buffers_start = util::AlignUp(reinterpret_cast<uintptr_t>(this->pointer_buffer_storage), 0x10);

stratosphere/libstratosphere/source/sf/hipc/sf_hipc_mitm_query_api.cpp

@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2018-2019 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+#include <stratosphere.hpp>
+#include "sf_hipc_mitm_query_api.hpp"
+
+namespace ams::sf::hipc::impl {
+
+    namespace {
+
+        class MitmQueryService : public IServiceObject {
+            private:
+                enum class CommandId {
+                    ShouldMitm = 65000,
+                };
+            private:
+                ServerManagerBase::MitmQueryFunction query_function;
+            public:
+                MitmQueryService(ServerManagerBase::MitmQueryFunction qf) : query_function(qf) { /* ... */ }
+
+                void ShouldMitm(sf::Out<bool> out, os::ProcessId process_id, ncm::ProgramId program_id) {
+                    out.SetValue(this->query_function(process_id, program_id));
+                }
+            public:
+                DEFINE_SERVICE_DISPATCH_TABLE {
+                    MAKE_SERVICE_COMMAND_META(ShouldMitm),
+                };
+        };
+
+        /* Globals. */
+        os::Mutex g_query_server_lock;
+        bool g_registered_any = false;
+
+        void QueryServerProcessThreadMain(void *query_server) {
+            reinterpret_cast<ServerManagerBase *>(query_server)->LoopProcess();
+        }
+
+        constexpr size_t QueryServerProcessThreadStackSize = 0x4000;
+        constexpr int    QueryServerProcessThreadPriority  = 27;
+        os::StaticThread<QueryServerProcessThreadStackSize> g_query_server_process_thread;
+
+    }
+
+    void RegisterMitmQueryHandle(Handle query_handle, ServerManagerBase::MitmQueryFunction query_func) {
+        std::scoped_lock lk(g_query_server_lock);
+
+        constexpr size_t MaxServers = 0;
+        sf::hipc::ServerManager<MaxServers> s_query_server;
+
+        R_ASSERT(s_query_server.RegisterSession(query_handle, cmif::ServiceObjectHolder(std::make_shared<MitmQueryService>(query_func))));
+
+        if (!g_registered_any) {
+            R_ASSERT(g_query_server_process_thread.Initialize(&QueryServerProcessThreadMain, &s_query_server, QueryServerProcessThreadPriority));
+            R_ASSERT(g_query_server_process_thread.Start());
+            g_registered_any = true;
+        }
+    }
+
+}

stratosphere/libstratosphere/source/sf/hipc/sf_hipc_mitm_query_api.hpp

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2018-2019 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+#pragma once
+#include <stratosphere.hpp>
+
+namespace ams::sf::hipc::impl {
+
+    void RegisterMitmQueryHandle(Handle query_handle, ServerManagerBase::MitmQueryFunction query_func);
+
+}

stratosphere/libstratosphere/source/sf/hipc/sf_hipc_server_manager.cpp

@@ -14,11 +14,22 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 #include <stratosphere.hpp>
+#include "sf_hipc_mitm_query_api.hpp"
 
 namespace ams::sf::hipc {
 
     ServerManagerBase::ServerBase::~ServerBase() { /* Pure virtual destructor, to prevent linker errors. */ }
 
+    Result ServerManagerBase::InstallMitmServerImpl(Handle *out_port_handle, sm::ServiceName service_name, ServerManagerBase::MitmQueryFunction query_func) {
+        /* Install the Mitm. */
+        Handle query_handle;
+        R_TRY(sm::mitm::InstallMitm(out_port_handle, &query_handle, service_name));
+
+        /* Register the query handle. */
+        impl::RegisterMitmQueryHandle(query_handle, query_func);
+        return ResultSuccess();
+    }
+
     void ServerManagerBase::RegisterSessionToWaitList(ServerSession *session) {
         session->has_received = false;
 

stratosphere/libstratosphere/source/sf/hipc/sf_hipc_server_session_manager.cpp

@@ -141,7 +141,7 @@ namespace ams::sf::hipc {
     Result ServerSessionManager::RegisterSession(Handle session_handle, cmif::ServiceObjectHolder &&obj) {
         /* We don't actually care about what happens to the session. It'll get linked. */
         ServerSession *session_ptr = nullptr;
-        return this->RegisterSession(&session_ptr, session_handle,std::forward<cmif::ServiceObjectHolder>(obj));
+        return this->RegisterSession(&session_ptr, session_handle, std::forward<cmif::ServiceObjectHolder>(obj));
     }
 
     Result ServerSessionManager::AcceptSession(Handle port_handle, cmif::ServiceObjectHolder &&obj) {