From 0cd9a2cee3757e97be55f1e716627fe9ce96902f Mon Sep 17 00:00:00 2001
From: antao <antao2002@gmail.com>
Date: Sat, 29 Sep 2018 15:30:27 +0800
Subject: [PATCH] Add config option max_connections_per_ip

---
 config.json.example               |  2 ++
 lib/inc/drogon/HttpAppFramework.h |  1 +
 lib/src/ConfigLoader.cc           |  7 +++++++
 lib/src/HttpAppFrameworkImpl.cc   | 24 ++++++++++++++++++++++++
 lib/src/HttpAppFrameworkImpl.h    |  7 +++++++
 5 files changed, 41 insertions(+)

diff --git a/config.json.example b/config.json.example
index 12e39dde..a3174729 100644
--- a/config.json.example
+++ b/config.json.example
@@ -41,6 +41,8 @@
     "file_types":["gif","png","jpg","js","css","html","ico","swf","xap","apk","cur","xml"],
     //max_connections:max connections number,100000 by default
     "max_connections":100000,
+    //max_connections_per_ip:max connections number per clinet,0 by default which means no limit
+    "max_connections_per_ip":0,
     //Load_dynamic_views: false by default, when set to true, drogon will
     //compile and load dynamically "CSP View Files" in directories defined
     //by "dynamic_views_path"
diff --git a/lib/inc/drogon/HttpAppFramework.h b/lib/inc/drogon/HttpAppFramework.h
index d9c794c6..5fdf7b44 100755
--- a/lib/inc/drogon/HttpAppFramework.h
+++ b/lib/inc/drogon/HttpAppFramework.h
@@ -97,6 +97,7 @@ namespace drogon
         virtual void enableDynamicViewsLoading(const std::vector<std::string> &libPaths)=0;
 
         virtual void setMaxConnectionNum(size_t maxConnections)=0;
+        virtual void setMaxConnectionNumPerIP(size_t maxConnectionsPerIP)=0;
 
         virtual void loadConfigFile(const std::string &fileName)=0;
 
diff --git a/lib/src/ConfigLoader.cc b/lib/src/ConfigLoader.cc
index 651e1590..7f7106f2 100644
--- a/lib/src/ConfigLoader.cc
+++ b/lib/src/ConfigLoader.cc
@@ -103,6 +103,13 @@ static void loadApp(const Json::Value &app)
     {
         HttpAppFramework::instance().setMaxConnectionNum(maxConns);
     }
+    //max connections per IP
+    auto maxConnsPerIP=app.get("max_connections_per_ip",0).asUInt64();
+    if(maxConnsPerIP>0)
+    {
+        HttpAppFramework::instance().setMaxConnectionNumPerIP(maxConnsPerIP);
+    }
+
     //dynamic views
     auto enableDynamicViews=app.get("load_dynamic_views",false).asBool();
     if(enableDynamicViews)
diff --git a/lib/src/HttpAppFrameworkImpl.cc b/lib/src/HttpAppFrameworkImpl.cc
index daa49075..55445e28 100755
--- a/lib/src/HttpAppFrameworkImpl.cc
+++ b/lib/src/HttpAppFrameworkImpl.cc
@@ -229,6 +229,10 @@ void HttpAppFrameworkImpl::setMaxConnectionNum(size_t maxConnections)
 {
     _maxConnectionNum=maxConnections;
 }
+void HttpAppFrameworkImpl::setMaxConnectionNumPerIP(size_t maxConnectionsPerIP)
+{
+    _maxConnectionNumPerIP=maxConnectionsPerIP;
+}
 void HttpAppFrameworkImpl::loadConfigFile(const std::string &fileName)
 {
     ConfigLoader loader(fileName);
@@ -506,10 +510,30 @@ void HttpAppFrameworkImpl::onConnection(const TcpConnectionPtr &conn)
             LOG_ERROR<<"too much connections!force close!";
             conn->forceClose();
         }
+        else if(_maxConnectionNumPerIP>0)
+        {
+            {
+                auto iter=_connectionsNumMap.find(conn->peerAddr().toIp());
+                if(iter==_connectionsNumMap.end())
+                {
+                    _connectionsNumMap[conn->peerAddr().toIp()]=0;
+                }
+                if(_connectionsNumMap[conn->peerAddr().toIp()]++>=_maxConnectionNumPerIP)
+                {
+                    conn->forceClose();
+                }
+            }
+        }
     }
     else
     {
         _connectionNum--;
+
+        if(_maxConnectionNumPerIP>0&&_connectionsNumMap.find(conn->peerAddr().toIp())!=_connectionsNumMap.end())
+        {
+            std::lock_guard<std::mutex> guard(_connectionsNumMapMutex);
+            _connectionsNumMap[conn->peerAddr().toIp()]--;
+        }
     }
 }
 std::string parseWebsockFrame(trantor::MsgBuffer *buffer)
diff --git a/lib/src/HttpAppFrameworkImpl.h b/lib/src/HttpAppFrameworkImpl.h
index eb7f66f1..c3850d70 100644
--- a/lib/src/HttpAppFrameworkImpl.h
+++ b/lib/src/HttpAppFrameworkImpl.h
@@ -60,6 +60,7 @@ namespace drogon
         virtual void setFileTypes(const std::vector<std::string> &types) override;
         virtual void enableDynamicViewsLoading(const std::vector<std::string> &libPaths) override;
         virtual void setMaxConnectionNum(size_t maxConnections) override;
+        virtual void setMaxConnectionNumPerIP(size_t maxConnectionsPerIP) override;
         virtual void loadConfigFile(const std::string &fileName) override;
         virtual void enableRunAsDaemon() override {_runAsDaemon=true;}
         virtual void enableRelaunchOnError() override {_relaunchOnError=true;}
@@ -164,7 +165,13 @@ namespace drogon
         std::string _sslKeyPath;
 
         size_t _maxConnectionNum=100000;
+        size_t _maxConnectionNumPerIP=0;
+
         std::atomic<uint64_t> _connectionNum;
+        std::unordered_map<std::string,size_t> _connectionsNumMap;
+
+        std::mutex _connectionsNumMapMutex;
+
 
         bool _runAsDaemon=false;
         bool _relaunchOnError=false;