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, 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 ) -> Result { - 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) + } } }