add dos guard

server/DOSGuard.h

@@ -0,0 +1,86 @@
+#include <boost/asio.hpp>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+class DOSGuard
+{
+    std::unordered_map<std::string, uint32_t> ipFetchCount_;
+    uint32_t maxFetches_ = 100;
+    uint32_t sweepInterval_ = 1;
+    std::unordered_set<std::string> whitelist_;
+    boost::asio::io_context& ctx_;
+    std::mutex mtx_;
+
+public:
+    DOSGuard(boost::json::object const& config, boost::asio::io_context& ctx)
+        : ctx_(ctx)
+    {
+        if (config.contains("dos_guard"))
+        {
+            auto dosGuardConfig = config.at("dos_guard").as_object();
+            if (dosGuardConfig.contains("max_fetches") &&
+                dosGuardConfig.contains("sweep_interval"))
+            {
+                maxFetches_ = dosGuardConfig.at("max_fetches").as_int64();
+                sweepInterval_ = dosGuardConfig.at("sweep_interval").as_int64();
+            }
+            if (dosGuardConfig.contains("whitelist"))
+            {
+                auto whitelist = dosGuardConfig.at("whitelist").as_array();
+                for (auto& ip : whitelist)
+                    whitelist_.insert(ip.as_string().c_str());
+            }
+        }
+        createTimer();
+    }
+
+    void
+    createTimer()
+    {
+        auto wait = std::chrono::seconds(sweepInterval_);
+        std::shared_ptr<boost::asio::steady_timer> timer =
+            std::make_shared<boost::asio::steady_timer>(
+                ctx_, std::chrono::steady_clock::now() + wait);
+        timer->async_wait(
+            [timer, this](const boost::system::error_code& error) {
+                clear();
+                createTimer();
+            });
+    }
+
+    bool
+    isOk(std::string const& ip)
+    {
+        if (whitelist_.count(ip) > 0)
+            return true;
+        std::unique_lock lck(mtx_);
+        auto it = ipFetchCount_.find(ip);
+        if (it == ipFetchCount_.end())
+            return true;
+        return it->second < maxFetches_;
+    }
+
+    bool
+    add(std::string const& ip, uint32_t numObjects)
+    {
+        if (whitelist_.count(ip) > 0)
+            return true;
+        {
+            std::unique_lock lck(mtx_);
+            auto it = ipFetchCount_.find(ip);
+            if (it == ipFetchCount_.end())
+                ipFetchCount_[ip] = numObjects;
+            else
+                it->second += numObjects;
+        }
+        return isOk(ip);
+    }
+
+    void
+    clear()
+    {
+        std::unique_lock lck(mtx_);
+        ipFetchCount_.clear();
+    }
+};