From 00eb1ddb3fa4f73a64b5c60b33cab994faff096b Mon Sep 17 00:00:00 2001
From: Jerboa-app <dev@jerboa.app>
Date: Sat, 10 Feb 2024 12:13:37 +0000
Subject: [PATCH] clear throttle ips periodically

configure throttle in config.json
---
 Cargo.toml          |  2 +-
 src/server/http.rs  |  5 +++--
 src/server/https.rs |  5 +++--
 src/server/model.rs | 34 +++++++++++++++++++++++++++++++++-
 src/web/throttle.rs | 39 ++++++++++++++++++++++++++++-----------
 5 files changed, 68 insertions(+), 17 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index dcb6040..1f1289c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pulse"
-version = "0.0.5"
+version = "0.0.6"
 authors = ["Jerboa"]
 
 edition="2021"
diff --git a/src/server/http.rs b/src/server/http.rs
index 2c22306..1b5cfb2 100644
--- a/src/server/http.rs
+++ b/src/server/http.rs
@@ -71,8 +71,9 @@ impl ServerHttp
 
         let requests: IpThrottler = IpThrottler::new
         (
-            10.0, 
-            5000
+            config.get_throttle_config().get_max_requests_per_second(), 
+            config.get_throttle_config().get_timeout_millis(),
+            config.get_throttle_config().get_clear_period_seconds()
         );
 
         let throttle_state = Arc::new(Mutex::new(requests));
diff --git a/src/server/https.rs b/src/server/https.rs
index d071122..d343a5e 100644
--- a/src/server/https.rs
+++ b/src/server/https.rs
@@ -74,8 +74,9 @@ impl Server
 
         let requests: IpThrottler = IpThrottler::new
         (
-            10.0, 
-            5000
+            config.get_throttle_config().get_max_requests_per_second(), 
+            config.get_throttle_config().get_timeout_millis(),
+            config.get_throttle_config().get_clear_period_seconds()
         );
 
         let throttle_state = Arc::new(Mutex::new(requests));
diff --git a/src/server/model.rs b/src/server/model.rs
index c18174b..7cde055 100644
--- a/src/server/model.rs
+++ b/src/server/model.rs
@@ -4,13 +4,40 @@ use crate::web::{github::model::GithubStats, discord::request::model::Webhook};
 
 pub const CONFIG_PATH: &str = "config.json";
 
+#[derive(Clone, Serialize, Deserialize)]
+pub struct ThrottleConfig
+{
+    max_requests_per_second: f64,
+    timeout_millis: u128,
+    clear_period_seconds: u64
+}
+
+impl ThrottleConfig 
+{
+    pub fn get_max_requests_per_second(&self) -> f64
+    {
+        self.max_requests_per_second
+    }
+
+    pub fn get_timeout_millis(&self) -> u128
+    {
+        self.timeout_millis
+    }
+
+    pub fn get_clear_period_seconds(&self) -> u64
+    {
+        self.clear_period_seconds
+    }
+}
+
 #[derive(Clone, Serialize, Deserialize)]
 pub struct Config
 {
     port: u16,
     stats_endpoint: Webhook,
     cert_path: String,
-    key_path: String
+    key_path: String,
+    throttle: ThrottleConfig
 }
 
 impl Config 
@@ -34,6 +61,11 @@ impl Config
     {
         self.key_path.clone()
     }
+
+    pub fn get_throttle_config(&self) -> ThrottleConfig
+    {
+        self.throttle.clone()
+    }
     
 }
 
diff --git a/src/web/throttle.rs b/src/web/throttle.rs
index 12fb449..a8fb20a 100644
--- a/src/web/throttle.rs
+++ b/src/web/throttle.rs
@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 use std::net::{SocketAddr, Ipv4Addr, IpAddr};
-use std::time::Instant;
+use std::time::{Instant, Duration};
 use std::sync::Arc;
 use tokio::sync::Mutex;
 
@@ -32,17 +32,30 @@ pub struct IpThrottler
     requests_from: HashMap<Ipv4Addr, Requests>,
     max_requests_per_second: f64,
     timeout_millis: u128,
+    clear_period: Duration,
+    last_clear: Instant 
 }
 
 impl IpThrottler
 {
-    pub fn new(max_requests_per_second: f64, timeout_millis: u128) -> IpThrottler
+    pub fn new(max_requests_per_second: f64, timeout_millis: u128, clear_period_seconds: u64) -> IpThrottler
     {
         IpThrottler 
         {
             requests_from: HashMap::new(), 
             max_requests_per_second: max_requests_per_second,
-            timeout_millis: timeout_millis
+            timeout_millis: timeout_millis,
+            clear_period: Duration::from_secs(clear_period_seconds),
+            last_clear: Instant::now()
+        }
+    }
+
+    pub fn check_clear(&mut self)
+    {
+        if self.last_clear.elapsed() > self.clear_period
+        {
+            self.requests_from.clear();
+            self.last_clear = Instant::now();
         }
     }
 
@@ -105,15 +118,19 @@ pub async fn handle_throttle<B>
 ) -> Result<Response, StatusCode>
 {
 
-    if state.lock().await.is_limited(addr)
-    {
-        Err(StatusCode::TOO_MANY_REQUESTS)
-    }
-    else 
     {
-        println!("passing on");
-        let response = next.run(request).await;
-        Ok(response)
+        let mut throttler = state.lock().await;
+        throttler.check_clear();
+        if throttler.is_limited(addr)
+        {
+            Err(StatusCode::TOO_MANY_REQUESTS)
+        }
+        else 
+        {
+            crate::debug(format!("Allowing: {}", addr), None);
+            let response = next.run(request).await;
+            Ok(response)
+        }
     }
     
 }