Added configurable rate limiter

src/utils/parceyaml.rs

@@ -16,6 +16,7 @@ pub fn load_configuration(d: &str, kind: &str) -> Option<Configuration> {
             sticky_sessions: false,
             to_https: None,
             authentication: DashMap::new(),
+            rate_limit: None,
         },
     };
     toreturn.upstreams = UpstreamsDashMap::new();
@@ -55,9 +56,9 @@ pub fn load_configuration(d: &str, kind: &str) -> Option<Configuration> {
                 }
                 global_headers.insert("/".to_string(), hl);
                 toreturn.headers.insert("GLOBAL_HEADERS".to_string(), global_headers);
-
                 toreturn.extraparams.sticky_sessions = parsed.sticky_sessions;
                 toreturn.extraparams.to_https = parsed.to_https;
+                toreturn.extraparams.rate_limit = parsed.rate_limit;
             }
             if let Some(auth) = &parsed.authorization {
                 let name = auth.get("type").unwrap().to_string();
@@ -67,7 +68,6 @@ pub fn load_configuration(d: &str, kind: &str) -> Option<Configuration> {
             } else {
                 toreturn.extraparams.authentication = DashMap::new();
             }
-
             match parsed.provider.as_str() {
                 "file" => {
                     toreturn.typecfg = "file".to_string();

src/utils/structs.rs

@@ -19,6 +19,7 @@ pub struct Extraparams {
     pub sticky_sessions: bool,
     pub to_https: Option<bool>,
     pub authentication: DashMap<String, Vec<String>>,
+    pub rate_limit: Option<isize>,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -37,11 +38,13 @@ pub struct Config {
     pub headers: Option<Vec<String>>,
     pub authorization: Option<HashMap<String, String>>,
     pub consul: Option<Consul>,
+    pub rate_limit: Option<isize>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct HostConfig {
     pub paths: HashMap<String, PathConfig>,
+    pub rate_limit: Option<isize>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -49,6 +52,7 @@ pub struct PathConfig {
     pub servers: Vec<String>,
     pub to_https: Option<bool>,
     pub headers: Option<Vec<String>>,
+    pub rate_limit: Option<isize>,
 }
 #[derive(Debug)]
 pub struct Configuration {

src/web/bgservice.rs

@@ -67,6 +67,7 @@ impl BackgroundService for LB {
                             new.sticky_sessions = ss.extraparams.sticky_sessions;
                             new.to_https = ss.extraparams.to_https;
                             new.authentication = ss.extraparams.authentication.clone();
+                            new.rate_limit = ss.extraparams.rate_limit;
                             self.extraparams.store(Arc::new(new));
                             self.headers.clear();
 

src/web/proxyhttp.rs

@@ -6,13 +6,16 @@ use arc_swap::ArcSwap;
 use async_trait::async_trait;
 use axum::body::Bytes;
 use log::{debug, warn};
+use once_cell::sync::Lazy;
 use pingora::http::{RequestHeader, ResponseHeader, StatusCode};
 use pingora::prelude::*;
 use pingora::ErrorSource::Upstream;
 use pingora_core::listeners::ALPN;
 use pingora_core::prelude::HttpPeer;
+use pingora_limits::rate::Rate;
 use pingora_proxy::{ProxyHttp, Session};
 use std::sync::Arc;
+use std::time::Duration;
 use tokio::time::Instant;
 
 #[derive(Clone)]
@@ -30,7 +33,12 @@ pub struct Context {
     to_https: bool,
     redirect_to: String,
     start_time: Instant,
+    hostname: Option<String>,
 }
+// Rate limiter
+static RATE_LIMITER: Lazy<Rate> = Lazy::new(|| Rate::new(Duration::from_secs(1)));
+// max request per second per client
+// static MAX_REQ_PER_SEC: isize = 1;
 
 #[async_trait]
 impl ProxyHttp for LB {
@@ -41,6 +49,7 @@ impl ProxyHttp for LB {
             to_https: false,
             redirect_to: String::new(),
             start_time: Instant::now(),
+            hostname: None,
         }
     }
     async fn request_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<bool> {
@@ -52,11 +61,32 @@ impl ProxyHttp for LB {
                 return Ok(true);
             }
         };
+
+        let hostname = return_header_host(&session);
+        _ctx.hostname = hostname.clone();
+        if let Some(rate) = self.extraparams.load().rate_limit {
+            match hostname {
+                None => return Ok(false),
+                Some(host) => {
+                    let curr_window_requests = RATE_LIMITER.observe(&host, 1);
+                    if curr_window_requests > rate {
+                        let mut header = ResponseHeader::build(429, None).unwrap();
+                        header.insert_header("X-Rate-Limit-Limit", rate.to_string()).unwrap();
+                        header.insert_header("X-Rate-Limit-Remaining", "0").unwrap();
+                        header.insert_header("X-Rate-Limit-Reset", "1").unwrap();
+                        session.set_keepalive(None);
+                        session.write_response_header(Box::new(header), true).await?;
+                        debug!("Rate limited: {:?}, {}", session.client_addr(), rate);
+                        return Ok(true);
+                    }
+                }
+            };
+        }
         Ok(false)
     }
     async fn upstream_peer(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<Box<HttpPeer>> {
-        let host_name = return_header_host(&session);
-        match host_name {
+        // let host_name = return_header_host(&session);
+        match ctx.hostname.as_ref() {
             Some(hostname) => {
                 let mut backend_id = None;
 
@@ -84,7 +114,7 @@ impl ProxyHttp for LB {
                             peer.options.alpn = ALPN::H2;
                         }
                         if ssl {
-                            peer.sni = hostname.to_string();
+                            peer.sni = hostname.clone();
                             peer.options.verify_cert = false;
                             peer.options.verify_hostname = false;
                         }
@@ -172,7 +202,8 @@ impl ProxyHttp for LB {
             redirect_response.insert_header("Content-Length", "0")?;
             session.write_response_header(Box::new(redirect_response), false).await?;
         }
-        match return_header_host(&session) {
+        // match return_header_host(&session) {
+        match ctx.hostname.as_ref() {
             Some(host) => {
                 let path = session.req_header().uri.path();
                 let host_header = host;
@@ -215,17 +246,17 @@ impl ProxyHttp for LB {
     }
 }
 
-fn return_header_host(session: &Session) -> Option<&str> {
+fn return_header_host(session: &Session) -> Option<String> {
     if session.is_http2() {
         match session.req_header().uri.host() {
-            Some(host) => Option::from(host),
+            Some(host) => Option::from(host.to_string()),
             None => None,
         }
     } else {
         match session.req_header().headers.get("host") {
             Some(host) => {
                 let header_host = host.to_str().unwrap().splitn(2, ':').collect::<Vec<&str>>();
-                Option::from(header_host[0])
+                Option::from(header_host[0].to_string())
             }
             None => None,
         }

src/web/start.rs

@@ -33,6 +33,7 @@ pub fn run() {
         sticky_sessions: false,
         to_https: None,
         authentication: DashMap::new(),
+        rate_limit: None,
     }));
 
     let cfg = Arc::new(maincfg);

src/web/webserver.rs

@@ -91,7 +91,7 @@ async fn conf(State(mut st): State<AppState>, Query(params): Query<HashMap<Strin
     }
 
     if let Some(s) = params.get("key") {
-        if s.to_owned() == st.master_key.to_owned() {
+        if s.to_owned() == st.master_key {
             if let Some(serverlist) = crate::utils::parceyaml::load_configuration(content.as_str(), "content") {
                 st.config_sender.send(serverlist).await.unwrap();
                 return Response::builder().status(StatusCode::OK).body(Body::from("Config, conf file, updated !\n")).unwrap();