diff --git a/lib/src/HttpAppFrameworkImpl.cc b/lib/src/HttpAppFrameworkImpl.cc index e451e02d..d2f239a0 100755 --- a/lib/src/HttpAppFrameworkImpl.cc +++ b/lib/src/HttpAppFrameworkImpl.cc @@ -441,6 +441,7 @@ void HttpAppFrameworkImpl::onWebsockDisconnect(const WebSocketConnectionPtr &wsC } void HttpAppFrameworkImpl::onConnection(const TcpConnectionPtr &conn) { + static std::mutex mtx; if (conn->connected()) { if (_connectionNum.fetch_add(1) >= _maxConnectionNum) @@ -451,14 +452,17 @@ void HttpAppFrameworkImpl::onConnection(const TcpConnectionPtr &conn) else if (_maxConnectionNumPerIP > 0) { { + std::lock_guard lock(mtx); auto iter = _connectionsNumMap.find(conn->peerAddr().toIp()); if (iter == _connectionsNumMap.end()) { - _connectionsNumMap[conn->peerAddr().toIp()] = 0; + _connectionsNumMap[conn->peerAddr().toIp()] = 1; } - if (_connectionsNumMap[conn->peerAddr().toIp()]++ >= _maxConnectionNumPerIP) + else if (iter->second++ > _maxConnectionNumPerIP) { - conn->forceClose(); + conn->getLoop()->queueInLoop([conn]() { + conn->forceClose(); + }); } } } @@ -466,9 +470,18 @@ void HttpAppFrameworkImpl::onConnection(const TcpConnectionPtr &conn) else { _connectionNum--; - if (_maxConnectionNumPerIP > 0 && _connectionsNumMap.find(conn->peerAddr().toIp()) != _connectionsNumMap.end()) + if (_maxConnectionNumPerIP > 0) { - _connectionsNumMap[conn->peerAddr().toIp()]--; + std::lock_guard lock(mtx); + auto iter = _connectionsNumMap.find(conn->peerAddr().toIp()); + if (iter != _connectionsNumMap.end()) + { + iter->second--; + if (iter->second <= 0) + { + _connectionsNumMap.erase(iter); + } + } } } } diff --git a/lib/src/HttpAppFrameworkImpl.h b/lib/src/HttpAppFrameworkImpl.h index 6f049919..7e05e1f9 100644 --- a/lib/src/HttpAppFrameworkImpl.h +++ b/lib/src/HttpAppFrameworkImpl.h @@ -184,7 +184,7 @@ class HttpAppFrameworkImpl : public HttpAppFramework size_t _maxConnectionNumPerIP = 0; std::atomic _connectionNum; - std::unordered_map> _connectionsNumMap; + std::unordered_map _connectionsNumMap; bool _runAsDaemon = false; bool _relaunchOnError = false; diff --git a/trantor b/trantor index 50f5401a..e4da5dee 160000 --- a/trantor +++ b/trantor @@ -1 +1 @@ -Subproject commit 50f5401a0b97fcf94ced323d0913c2cb193f26d7 +Subproject commit e4da5dee4293475ea176a6615db0feb8bf98f833