diff --git a/include/stratosphere/mitm/mitm_server.hpp b/include/stratosphere/mitm/mitm_server.hpp index 8a15533f..dbb597b7 100644 --- a/include/stratosphere/mitm/mitm_server.hpp +++ b/include/stratosphere/mitm/mitm_server.hpp @@ -21,6 +21,8 @@ #include "sm_mitm.h" #include "mitm_session.hpp" +void RegisterMitmServerQueryHandle(Handle query_h, ServiceObjectHolder &&service); + template class MitmServer : public IWaitable { static_assert(std::is_base_of::value, "MitM Service Objects must derive from IMitmServiceObject"); @@ -30,7 +32,8 @@ class MitmServer : public IWaitable { char mitm_name[9]; public: - MitmServer(Handle *out_query_h, const char *service_name, unsigned int max_s) : port_handle(0), max_sessions(max_s) { + MitmServer(const char *service_name, unsigned int max_s) : port_handle(0), max_sessions(max_s) { + Handle query_h = 0; Result rc = smMitMInitialize(); if (R_FAILED(rc)) { fatalSimple(rc); @@ -38,9 +41,10 @@ class MitmServer : public IWaitable { strncpy(mitm_name, service_name, 8); mitm_name[8] = '\x00'; - if (R_FAILED((rc = smMitMInstall(&this->port_handle, out_query_h, mitm_name)))) { + if (R_FAILED((rc = smMitMInstall(&this->port_handle, &query_h, mitm_name)))) { fatalSimple(rc); } + RegisterMitmServerQueryHandle(query_h, std::move(ServiceObjectHolder(std::move(std::make_shared>())))); smMitMExit(); } @@ -99,10 +103,6 @@ class MitmServer : public IWaitable { template static void AddMitmServerToManager(SessionManagerBase *manager, const char *srv_name, unsigned int max_sessions) { - Handle query_h; - auto *srv = new MitmServer(&query_h, srv_name, max_sessions); - manager->AddSession(query_h, std::move(ServiceObjectHolder(std::move(std::make_shared>())))); + auto *srv = new MitmServer(srv_name, max_sessions); manager->AddWaitable(srv); } - - diff --git a/source/mitm_server.cpp b/source/mitm_server.cpp new file mode 100644 index 00000000..b1a3129b --- /dev/null +++ b/source/mitm_server.cpp @@ -0,0 +1,50 @@ +/* + * Copyright (c) 2018 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 . + */ + +#include +#include +#include + +static HosMutex g_server_query_mutex; +static HosThread g_server_query_manager_thread; +static SessionManagerBase *g_server_query_manager = nullptr; + +static void ServerQueryManagerThreadFunc(void *arg) { + g_server_query_manager->Process(); +} + +void RegisterMitmServerQueryHandle(Handle query_h, ServiceObjectHolder &&service) { + std::scoped_lock lock(g_server_query_mutex); + + const bool exists = g_server_query_manager != nullptr; + if (!exists) { + /* Create a new waitable manager if it doesn't exist already. */ + g_server_query_manager = new WaitableManager(1); + } + + /* Add session to the manager. */ + g_server_query_manager->AddSession(query_h, std::move(service)); + + /* If this is our first time, launch thread. */ + if (!exists) { + if (R_FAILED(g_server_query_manager_thread.Initialize(&ServerQueryManagerThreadFunc, nullptr, 0x4000, 27))) { + std::abort(); + } + if (R_FAILED(g_server_query_manager_thread.Start())) { + std::abort(); + } + } +} \ No newline at end of file