implement connection limiter to replace manual limiting code

This commit is contained in:
Lion Kortlepel
2026-04-08 18:44:37 +02:00
parent 7925caf1a2
commit 6fca901aa2
5 changed files with 142 additions and 53 deletions
+2
View File
@@ -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)
+62
View File
@@ -0,0 +1,62 @@
#pragma once
#include <mutex>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
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<TGuard> 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<std::string, size_t> mPerIp { };
size_t mGlobal = 0;
};
+3 -3
View File
@@ -20,6 +20,7 @@
#include "BoostAliases.h"
#include "Compat.h"
#include "TConnectionLimiter.h"
#include "TResourceManager.h"
#include "TServer.h"
#include <boost/asio/io_context.hpp>
@@ -40,7 +41,7 @@ public:
void DisconnectClient(const std::weak_ptr<TClient>& c, const std::string& R);
void DisconnectClient(TClient& c, const std::string& R);
[[nodiscard]] bool SyncClient(const std::weak_ptr<TClient>& c);
void Identify(TConnection&& client);
void Identify(TConnection&& client, TConnectionLimiter::TGuard&&);
std::shared_ptr<TClient> Authentication(TConnection&& ClientConnection);
void SyncResources(TClient& c);
[[nodiscard]] bool UDPSend(TClient& Client, std::vector<uint8_t> Data);
@@ -61,8 +62,7 @@ private:
std::thread mUDPThread;
std::thread mTCPThread;
std::mutex mOpenIDMutex;
std::map<std::string, uint16_t> mClientMap;
std::mutex mClientMapMutex;
TConnectionLimiter mConnectionLimiter;
std::vector<uint8_t> UDPRcvFromClient(boost::asio::ip::udp::endpoint& ClientEndpoint);
void OnConnect(const std::weak_ptr<TClient>& c);
+61
View File
@@ -0,0 +1,61 @@
#include "TConnectionLimiter.h"
#include <mutex>
#include <optional>
#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::TGuard> 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);
}
}
+14 -50
View File
@@ -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<TClient> 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<TClient> &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);
}