Added configurable rate limiter

Cargo.lock

@@ -128,9 +128,11 @@ dependencies = [
  "log",
  "mimalloc",
  "notify",
+ "once_cell",
  "pingora",
  "pingora-core",
  "pingora-http",
+ "pingora-limits",
  "pingora-proxy",
  "prometheus 0.14.0",
  "rand 0.9.1",
@@ -2085,6 +2087,15 @@ dependencies = [
  "crc32fast",
 ]
 
+[[package]]
+name = "pingora-limits"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a719a8cb5558ca06bd6076c97b8905d500ea556da89e132ba53d4272844f95b9"
+dependencies = [
+ "ahash",
+]
+
 [[package]]
 name = "pingora-load-balancing"
 version = "0.5.0"

Cargo.toml

@@ -19,6 +19,7 @@ dashmap = "7.0.0-rc2"
 pingora-core = "0.5.0"
 pingora-proxy = "0.5.0"
 pingora-http = "0.5.0"
+pingora-limits = "0.5.0"
 #pingora-pool = "0.5.0"
 async-trait = "0.1.88"
 env_logger = "0.11.8"
@@ -48,5 +49,6 @@ lazy_static = "1.5.0"
 x509-parser = "0.17.0"
 rustls-pemfile = "2.2.0"
 tower-http = { version = "0.6.6", features = ["fs"] }
+once_cell = "1.20.2"
 
 

README.md

@@ -15,6 +15,7 @@ Built on Rust, on top of **Cloudflare’s Pingora engine**, **Aralez** delivers
 - **TLS Termination** — Built-in OpenSSL support.
     - **Automatic load of certificates** — Automatically reads and loads certificates from a folder, without a restart.
 - **Upstreams TLS detection** — Aralez will automatically detect if upstreams uses secure connection.
+- **Built in rate limiter** — Limit requests to server, by setting up upper limit for requests per seconds, per virtualhost.
 - **Authentication** — Supports Basic Auth, API tokens, and JWT verification.
     - **Basic Auth**
     - **API Key** via `x-api-key` header
@@ -147,6 +148,7 @@ A sample `upstreams.yaml` entry:
 provider: "file"
 sticky_sessions: false
 to_https: false
+rate_limit: 10
 headers:
   - "Access-Control-Allow-Origin:*"
   - "Access-Control-Allow-Methods:POST, GET, OPTIONS"
@@ -177,6 +179,9 @@ myhost.mydomain.com:
 
 - Sticky sessions are disabled globally. This setting applies to all upstreams. If enabled all requests will be 301 redirected to HTTPS.
 - HTTP to HTTPS redirect disabled globally, but can be overridden by `to_https` setting per upstream.
+- Requests to each hosted domains will be limited to 10 requests per second per virtualhost.
+    - The limiter is per virtualhost so requests and limits will be calculated per virtualhost individually.
+    - Optional. Rate limiter will be disabled if the parameter is entirely removed from config.
 - Requests to `myhost.mydomain.com/` will be proxied to `127.0.0.1` and `127.0.0.2`.
 - Plain HTTP to `myhost.mydomain.com/foo` will get 301 redirect to configured TLS port of Aralez.
 - Requests to `myhost.mydomain.com/foo` will be proxied to `127.0.0.4` and `127.0.0.5`.

etc/upstreams.yaml

@@ -2,6 +2,7 @@
 provider: "file" # consul
 sticky_sessions: false
 to_ssl: false
+#rate_limit: 100
 headers:
   - "Access-Control-Allow-Origin:*"
   - "Access-Control-Allow-Methods:POST, GET, OPTIONS"

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);
+        // let host_name = return_header_host(&session);
-        match host_name {
+        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();