From 6fca901aa2e4f19c200dcfac1938d52fa904fbcf Mon Sep 17 00:00:00 2001 From: Lion Kortlepel Date: Wed, 8 Apr 2026 18:44:37 +0200 Subject: [PATCH] implement connection limiter to replace manual limiting code --- CMakeLists.txt | 2 ++ include/TConnectionLimiter.h | 62 ++++++++++++++++++++++++++++++++++ include/TNetwork.h | 6 ++-- src/TConnectionLimiter.cpp | 61 ++++++++++++++++++++++++++++++++++ src/TNetwork.cpp | 64 ++++++++---------------------------- 5 files changed, 142 insertions(+), 53 deletions(-) create mode 100644 include/TConnectionLimiter.h create mode 100644 src/TConnectionLimiter.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 67e758d..cdec565 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -52,6 +52,7 @@ set(PRJ_HEADERS include/Settings.h include/Profiling.h include/ChronoWrapper.h + include/TConnectionLimiter.h ) # add all source files (.cpp) to this, except the one with main() set(PRJ_SOURCES @@ -78,6 +79,7 @@ set(PRJ_SOURCES src/Settings.cpp src/Profiling.cpp src/ChronoWrapper.cpp + src/TConnectionLimiter.cpp ) find_package(Lua REQUIRED) diff --git a/include/TConnectionLimiter.h b/include/TConnectionLimiter.h new file mode 100644 index 0000000..ab70fa5 --- /dev/null +++ b/include/TConnectionLimiter.h @@ -0,0 +1,62 @@ +#pragma once + +#include +#include +#include +#include +#include + +class TConnectionLimiter { +public: + class TGuard { + public: + TGuard() = default; + TGuard(TConnectionLimiter* owner, std::string ip); + + TGuard(const TGuard&) = delete; + TGuard& operator=(const TGuard&) = delete; + + ~TGuard() { Release(); } + + // not threadsafe + TGuard(TGuard&& other) noexcept { *this = std::move(other); } + // not threadsafe + TGuard& operator=(TGuard&& other) noexcept; + + private: + friend class TConnectionLimiter; + + // not threadsafe + void Release(); + + TConnectionLimiter* mOwner { nullptr }; + std::string mIp; + }; + + TConnectionLimiter(size_t maxPerIp, size_t maxGlobal); + + [[nodiscard]] std::optional TryAcquire(const std::string& ip); + +private: + void Release(const std::string& ip) { + std::unique_lock Lock { mMutex }; + auto It = mPerIp.find(ip); + if (It != mPerIp.end()) { + // this guard exists to avoid underflow in case something goes wrong and this gets called too many times + if (It->second > 0) + --It->second; + if (It->second == 0) + mPerIp.erase(It); + } + // this guard exists to avoid underflow in case something goes wrong and this gets called too many times + if (mGlobal > 0) + --mGlobal; + } + + const size_t mMaxPerIp; + const size_t mMaxGlobal; + + std::mutex mMutex { }; + std::unordered_map mPerIp { }; + size_t mGlobal = 0; +}; diff --git a/include/TNetwork.h b/include/TNetwork.h index d5e5473..d1391b4 100644 --- a/include/TNetwork.h +++ b/include/TNetwork.h @@ -20,6 +20,7 @@ #include "BoostAliases.h" #include "Compat.h" +#include "TConnectionLimiter.h" #include "TResourceManager.h" #include "TServer.h" #include @@ -40,7 +41,7 @@ public: void DisconnectClient(const std::weak_ptr& c, const std::string& R); void DisconnectClient(TClient& c, const std::string& R); [[nodiscard]] bool SyncClient(const std::weak_ptr& c); - void Identify(TConnection&& client); + void Identify(TConnection&& client, TConnectionLimiter::TGuard&&); std::shared_ptr Authentication(TConnection&& ClientConnection); void SyncResources(TClient& c); [[nodiscard]] bool UDPSend(TClient& Client, std::vector Data); @@ -61,8 +62,7 @@ private: std::thread mUDPThread; std::thread mTCPThread; std::mutex mOpenIDMutex; - std::map mClientMap; - std::mutex mClientMapMutex; + TConnectionLimiter mConnectionLimiter; std::vector UDPRcvFromClient(boost::asio::ip::udp::endpoint& ClientEndpoint); void OnConnect(const std::weak_ptr& c); diff --git a/src/TConnectionLimiter.cpp b/src/TConnectionLimiter.cpp new file mode 100644 index 0000000..c8b8c4b --- /dev/null +++ b/src/TConnectionLimiter.cpp @@ -0,0 +1,61 @@ +#include "TConnectionLimiter.h" +#include +#include +#include "Common.h" + +TConnectionLimiter::TConnectionLimiter(size_t maxPerIp, size_t maxGlobal) + : mMaxPerIp(maxPerIp) + , mMaxGlobal(maxGlobal) { + if (maxPerIp == 0) { + beammp_errorf("Max connections count per IP is set to zero; the server would reject ALL connections"); + throw std::runtime_error("Invalid maximum connections per IP setting"); + } + if (maxGlobal == 0) { + beammp_errorf("Max connection count is set to zero; the server would reject ALL connections"); + throw std::runtime_error("Invalid maximum connections setting"); + } +} + +std::optional TConnectionLimiter::TryAcquire(const std::string& ip) { + std::unique_lock Lock { mMutex }; + if (mGlobal >= mMaxGlobal) return std::nullopt; + // `It` is the inserted element (so if insertion worked, its 0), or its the element that + // was already there, in which case we must check the ip connection limit + auto [It, _] = mPerIp.try_emplace(ip, 0); + if (It->second >= mMaxPerIp) { + return std::nullopt; + } + // now increment the counter finally + ++It->second; + ++mGlobal; + // RAII guard will drop the count once destructed + return TGuard(this, ip); +} + +TConnectionLimiter::TGuard::TGuard(TConnectionLimiter* owner, std::string ip) + : mOwner(owner) + , mIp(std::move(ip)) { + beammp_debugf("Acquired connection guard for {}", ip); +} + +TConnectionLimiter::TGuard& TConnectionLimiter::TGuard::operator=(TGuard&& other) noexcept { + // identity check + if (this != &other) { + Release(); + mOwner = other.mOwner; + mIp = std::move(other.mIp); + other.mOwner = nullptr; + } + return *this; +} + +void TConnectionLimiter::TGuard::Release() { + beammp_debugf("Trying to release connection guard for {} ...", mIp); + if (mOwner) { + mOwner->Release(mIp); + mOwner = nullptr; + beammp_debugf("... Released connection guard for {}", mIp); + } else { + beammp_debugf("... Connection guard for {} was already released (nothing happened)", mIp); + } +} diff --git a/src/TNetwork.cpp b/src/TNetwork.cpp index b76829a..71e8f44 100644 --- a/src/TNetwork.cpp +++ b/src/TNetwork.cpp @@ -20,6 +20,7 @@ #include "Client.h" #include "Common.h" #include "LuaAPI.h" +#include "TConnectionLimiter.h" #include "THeartbeatThread.h" #include "TLuaEngine.h" #include "TScopedTimer.h" @@ -59,7 +60,8 @@ TNetwork::TNetwork(TServer& Server, TPPSMonitor& PPSMonitor, TResourceManager& R : mServer(Server) , mPPSMonitor(PPSMonitor) , mUDPSock(Server.IoCtx()) - , mResourceManager(ResourceManager) { + , mResourceManager(ResourceManager) + , mConnectionLimiter(MAX_CONCURRENT_CONNECTIONS, MAX_GLOBAL_CONNECTIONS){ Application::SetSubsystemStatus("TCPNetwork", Application::Status::Starting); Application::SetSubsystemStatus("UDPNetwork", Application::Status::Starting); Application::RegisterShutdownHandler([&] { @@ -247,22 +249,17 @@ void TNetwork::TCPServerMain() { boost::asio::ip::tcp::socket ClientSocket = Acceptor.accept(ClientEp, ec); std::string ClientIP = ClientEp.address().to_string(); if (!ec) { - mClientMapMutex.lock(); - if (mClientMap[ClientIP] >= MAX_CONCURRENT_CONNECTIONS) { - beammp_debugf("The connection was rejected for {}, as it had {} concurrent connections.", ClientIP, mClientMap[ClientIP]); - } - else if (mClientMap.size() >= MAX_GLOBAL_CONNECTIONS) { - beammp_debugf("The connection was rejected for {}, as there are {} global connections.", ClientIP, mClientMap.size()); - } - else { + auto MaybeGuard = mConnectionLimiter.TryAcquire(ClientEp.address().to_v6().to_string()); + if (MaybeGuard.has_value()) { + // move-swap to avoid copy ctor (deleted) + auto Guard = std::move(MaybeGuard.value()); TConnection Conn { std::move(ClientSocket), ClientEp }; - std::thread ID(&TNetwork::Identify, this, std::move(Conn)); + std::thread ID(&TNetwork::Identify, this, std::move(Conn), std::move(Guard)); ID.detach(); // TODO: Add to a queue and attempt to join periodically - mClientMap[ClientIP]++; + } else { + beammp_errorf("Connection rejected for {} due to the global or concurrent connection limit", ClientIP); } - mClientMapMutex.unlock(); - } - else { + } else { beammp_errorf("Failed to accept() new client: {}", ec.message()); } } catch (const std::exception& e) { @@ -276,7 +273,7 @@ void TNetwork::TCPServerMain() { #include "Json.h" namespace json = rapidjson; -void TNetwork::Identify(TConnection&& RawConnection) { +void TNetwork::Identify(TConnection&& RawConnection, TConnectionLimiter::TGuard&& Guard) { RegisterThreadAuto(); char Code; @@ -285,17 +282,8 @@ void TNetwork::Identify(TConnection&& RawConnection) { // TODO: is this right?! beammp_debug("Error occured reading code"); RawConnection.Socket.shutdown(socket_base::shutdown_both, ec); - mClientMapMutex.lock(); - { - std::string ClientIP = RawConnection.SockAddr.address().to_string(); - if (mClientMap[ClientIP] > 0) { - mClientMap[ClientIP]--; - } - if (mClientMap[ClientIP] == 0) { - mClientMap.erase(ClientIP); - } - } - mClientMapMutex.unlock(); + // TODO: is this right too? + RawConnection.Socket.close(ec); return; } std::shared_ptr Client { nullptr }; @@ -326,17 +314,6 @@ void TNetwork::Identify(TConnection&& RawConnection) { beammp_errorf("Error during handling of code {} - client left in invalid state, closing socket: {}", Code, e.what()); boost::system::error_code ec; RawConnection.Socket.shutdown(boost::asio::socket_base::shutdown_both, ec); - mClientMapMutex.lock(); - { - std::string ClientIP = RawConnection.SockAddr.address().to_string(); - if (mClientMap[ClientIP] > 0) { - mClientMap[ClientIP]--; - } - if (mClientMap[ClientIP] == 0) { - mClientMap.erase(ClientIP); - } - } - mClientMapMutex.unlock(); if (ec) { beammp_debugf("Failed to shutdown client socket: {}", ec.message()); } @@ -661,19 +638,6 @@ void TNetwork::DisconnectClient(const std::weak_ptr &c, const std::stri void TNetwork::DisconnectClient(TClient &c, const std::string &R) { if (c.IsDisconnected()) return; - try { - std::string ClientIP = c.GetTCPSock().remote_endpoint().address().to_string(); - mClientMapMutex.lock(); - if (mClientMap[ClientIP] > 0) { - mClientMap[ClientIP]--; - } - if (mClientMap[ClientIP] == 0) { - mClientMap.erase(ClientIP); - } - mClientMapMutex.unlock(); - } catch (const std::exception& e) { - beammp_debugf("Failed to disconnect client (already disconnected?). There might be a lingering client in the IP-to-client map. This is not an error."); - } c.Disconnect(R); }