From 5afc07b0b3ea2f9c2879d1253f4d8fd2a816c700 Mon Sep 17 00:00:00 2001
From: sla-ppy <sla-ppy-pm@protonmail.com>
Date: Thu, 20 Jun 2024 16:39:03 +0200
Subject: [PATCH] add RateLimiter class

---
 include/RateLimiter.h | 20 +++++++++++++++++
 src/RateLimiter.cpp   | 50 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)
 create mode 100644 include/RateLimiter.h
 create mode 100644 src/RateLimiter.cpp

diff --git a/include/RateLimiter.h b/include/RateLimiter.h
new file mode 100644
index 0000000..8e0ab2c
--- /dev/null
+++ b/include/RateLimiter.h
@@ -0,0 +1,20 @@
+#include "Common.h"
+#include <chrono>
+#include <fstream>
+#include <iostream>
+#include <mutex>
+#include <unordered_map>
+#include <unordered_set>
+
+class RateLimiter {
+public:
+    RateLimiter();
+    bool isConnectionAllowed(const std::string& client_address);
+
+private:
+    std::unordered_map<std::string, std::vector<std::chrono::time_point<std::chrono::high_resolution_clock>>> m_connection;
+    std::mutex m_connection_mutex;
+
+    void blockIP(const std::string& client_address);
+    bool isIPBlocked(const std::string& client_address);
+};
\ No newline at end of file
diff --git a/src/RateLimiter.cpp b/src/RateLimiter.cpp
new file mode 100644
index 0000000..18028d7
--- /dev/null
+++ b/src/RateLimiter.cpp
@@ -0,0 +1,50 @@
+#include "RateLimiter.h"
+
+RateLimiter::RateLimiter() {};
+
+bool RateLimiter::isConnectionAllowed(const std::string& client_address) {
+    if (RateLimiter::isIPBlocked(client_address)) {
+        return false;
+    }
+    std::lock_guard<std::mutex> lock(m_connection_mutex);
+    auto current_time = std::chrono::high_resolution_clock::now();
+    auto& violations = m_connection[client_address];
+
+    // Deleting old violations (older than 5 seconds)
+    violations.erase(std::remove_if(violations.begin(), violations.end(),
+                         [&](const auto& timestamp) {
+                             return std::chrono::duration_cast<std::chrono::seconds>(current_time - timestamp).count() > 5;
+                         }),
+        violations.end());
+
+    violations.push_back(current_time);
+
+    if (violations.size() >= 4) {
+        RateLimiter::blockIP(client_address);
+        beammp_errorf("[DoS Protection] Client with the IP: {} surpassed the violation treshhold and is now on the blocked list", client_address);
+        return false;
+    }
+
+    return true; // We allow the connection
+}
+
+void RateLimiter::blockIP(const std::string& client_address) {
+    std::ofstream block_file("blocked_ips.txt", std::ios::app);
+    if (block_file.is_open()) {
+        block_file << client_address << std::endl;
+    }
+}
+
+bool RateLimiter::isIPBlocked(const std::string& client_address) {
+    std::ifstream block_file("blocked_ips.txt");
+    std::unordered_set<std::string> blockedIPs;
+
+    if (block_file.is_open()) {
+        std::string line;
+        while (std::getline(block_file, line)) {
+            blockedIPs.insert(line);
+        }
+    }
+
+    return blockedIPs.contains(client_address);
+};
\ No newline at end of file