Github/aralez
1
0
Fork 0
mirror of https://github.com/sadoyan/aralez.git synced 2026-04-29 22:38:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Per path rate limiter

Browse Source
This commit is contained in:
Ara Sadoyan
2025-07-24 13:34:15 +02:00
parent 6f012cee69
commit 2ad3a059ab
9 changed files with 114 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/utils/consul.rs
View File

@@ -133,6 +133,7 @@ async fn get_by_http(url: String, token: Option<String>) -> Option<DashMap<Strin
is_ssl: false,
is_http2: false,
to_https: false,
rate_limit: None,
};
values.push(to_add);
}

2
src/utils/healthcheck.rs
View File

@@ -44,6 +44,7 @@ pub async fn hc2(upslist: Arc<UpstreamsDashMap>, fullist: Arc<UpstreamsDashMap>,
is_ssl: tls.0,
is_http2: is_h2,
to_https: k.1.to_https,
rate_limit: k.1.rate_limit,
};
let resp = http_request(_link.as_str(), params.0, "", &client).await;
match resp.0 {
@@ -55,6 +56,7 @@ pub async fn hc2(upslist: Arc<UpstreamsDashMap>, fullist: Arc<UpstreamsDashMap>,
is_ssl: tls.0,
is_http2: is_h2,
to_https: k.1.to_https,
rate_limit: k.1.rate_limit,
};
}
innervec.push(scheme);

19
src/utils/parceyaml.rs
View File

@@ -80,6 +80,10 @@ fn populate_headers_and_auth(config: &mut Configuration, parsed: &Config) {
config.extraparams.to_https = parsed.to_https;
config.extraparams.rate_limit = parsed.rate_limit;
if let Some(rate) = &parsed.rate_limit {
info!("Applied Global Rate Limit : {} request per second", rate);
}
if let Some(auth) = &parsed.authorization {
let name = auth.get("type").unwrap_or(&"".to_string()).to_string();
let creds = auth.get("creds").unwrap_or(&"".to_string()).to_string();
@@ -94,8 +98,11 @@ fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
for (hostname, host_config) in upstreams {
let path_map = DashMap::new();
let header_list = DashMap::new();
for (path, path_config) in &host_config.paths {
if let Some(rate) = &path_config.rate_limit {
info!("Applied Rate Limit for {} : {} request per second", hostname, rate);
}
let mut server_list = Vec::new();
let mut hl = Vec::new();
@@ -109,6 +116,14 @@ fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
header_list.insert(path.clone(), hl);
for server in &path_config.servers {
// let mut rate: Option<isize> = None;
// let size: isize = path_config.servers.len() as isize;
// if let Some(limit) = &path_config.rate_limit {
// if size > 0 {
// rate = Some(limit / size);
// }
// }
if let Some((ip, port_str)) = server.split_once(':') {
if let Ok(port) = port_str.parse::<u16>() {
server_list.push(InnerMap {
@@ -117,6 +132,8 @@ fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
is_ssl: true,
is_http2: false,
to_https: path_config.to_https.unwrap_or(false),
// rate_limit: rate,
rate_limit: path_config.rate_limit,
});
}
}

19
src/utils/structs.rs
View File

@@ -5,6 +5,23 @@ use std::sync::atomic::AtomicUsize;
// pub type InnerMap = BackendConfig;
pub type UpstreamsDashMap = DashMap<String, DashMap<String, (Vec<InnerMap>, AtomicUsize)>>;
// #[derive(Debug, Default)]
// pub struct UpstreamsMap {
// pub upstreams: DashMap<String, DashMap<String, (Vec<InnerMap>, AtomicUsize)>>,
// pub ratelimit: DashMap<String, Option<isize>>,
// }
// impl UpstreamsMap {
// pub fn new() -> Self {
// Self {
// upstreams: Default::default(),
// ratelimit: Default::default(),
// }
// }
// }
//
// pub type XUpstreamsDashMap = DashMap<String, UpstreamsMap>;
pub type UpstreamsIdMap = DashMap<String, InnerMap>;
pub type Headers = DashMap<String, DashMap<String, Vec<(String, String)>>>;
@@ -97,6 +114,7 @@ pub struct InnerMap {
pub is_ssl: bool,
pub is_http2: bool,
pub to_https: bool,
pub rate_limit: Option<isize>,
}
impl InnerMap {
@@ -107,6 +125,7 @@ impl InnerMap {
is_ssl: Default::default(),
is_http2: Default::default(),
to_https: Default::default(),
rate_limit: Default::default(),
}
}
}

1
src/utils/tools.rs
View File

@@ -154,6 +154,7 @@ pub fn clone_idmap_into(original: &UpstreamsDashMap, cloned: &UpstreamsIdMap) {
is_ssl: false,
is_http2: false,
to_https: false,
rate_limit: None,
};
cloned.insert(id, to_add);
cloned.insert(hh, x.to_owned());

92
src/web/proxyhttp.rs
View File

@@ -1,6 +1,6 @@
use crate::utils::auth::authenticate;
use crate::utils::metrics::*;
use crate::utils::structs::{AppConfig, Extraparams, Headers, UpstreamsDashMap, UpstreamsIdMap};
use crate::utils::structs::{AppConfig, Extraparams, Headers, InnerMap, UpstreamsDashMap, UpstreamsIdMap};
use crate::web::gethosts::GetHost;
use arc_swap::ArcSwap;
use async_trait::async_trait;
@@ -36,6 +36,8 @@ pub struct Context {
redirect_to: String,
start_time: Instant,
hostname: Option<String>,
upstream_peer: Option<InnerMap>,
extraparams: arc_swap::Guard<Arc<Extraparams>>,
}
#[async_trait]
@@ -48,14 +50,18 @@ impl ProxyHttp for LB {
redirect_to: String::new(),
start_time: Instant::now(),
hostname: None,
upstream_peer: None,
extraparams: self.extraparams.load(),
}
}
async fn request_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<bool> {
if let Some(auth) = self.extraparams.load().authentication.get("authorization") {
let ep = _ctx.extraparams.clone();
if let Some(auth) = ep.authentication.get("authorization") {
let authenticated = authenticate(&auth.value(), &session);
if !authenticated {
let _ = session.respond_error(401).await;
warn!("Forbidden: {:?}, {}", session.client_addr(), session.req_header().uri.path().to_string());
warn!("Forbidden: {:?}, {}", session.client_addr(), session.req_header().uri.path());
return Ok(true);
}
};
@@ -63,23 +69,48 @@ impl ProxyHttp for LB {
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);
let mut backend_id = None;
if ep.sticky_sessions {
if let Some(cookies) = session.req_header().headers.get("cookie") {
if let Ok(cookie_str) = cookies.to_str() {
for cookie in cookie_str.split(';') {
let trimmed = cookie.trim();
if let Some(value) = trimmed.strip_prefix("backend_id=") {
backend_id = Some(value);
break;
}
}
}
};
}
}
match hostname {
None => return Ok(false),
Some(host) => {
let optioninnermap = self.get_host(host.as_str(), host.as_str(), backend_id);
match optioninnermap {
None => return Ok(false),
Some(ref innermap) => {
if let Some(rate) = innermap.rate_limit.or(ep.rate_limit) {
// let rate_key = session.client_addr().and_then(|addr| addr.as_inet()).map(|inet| inet.ip().to_string()).unwrap_or_else(|| host.to_string());
let rate_key = session.client_addr().and_then(|addr| addr.as_inet()).map(|inet| inet.ip());
let curr_window_requests = RATE_LIMITER.observe(&rate_key, 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: {:?}, {}", rate_key, rate);
return Ok(true);
}
}
}
}
_ctx.upstream_peer = optioninnermap;
}
}
Ok(false)
}
@@ -87,25 +118,7 @@ impl ProxyHttp for LB {
// let host_name = return_header_host(&session);
match ctx.hostname.as_ref() {
Some(hostname) => {
let mut backend_id = None;
if self.extraparams.load().sticky_sessions {
if let Some(cookies) = session.req_header().headers.get("cookie") {
if let Ok(cookie_str) = cookies.to_str() {
for cookie in cookie_str.split(';') {
let trimmed = cookie.trim();
if let Some(value) = trimmed.strip_prefix("backend_id=") {
backend_id = Some(value);
break;
}
}
}
}
}
let optioninnermap = self.get_host(hostname, hostname, backend_id);
match optioninnermap {
match ctx.upstream_peer.as_ref() {
// Some((address, port, ssl, is_h2, to_https)) => {
Some(innermap) => {
let mut peer = Box::new(HttpPeer::new((innermap.address.clone(), innermap.port.clone()), innermap.is_ssl, String::new()));
@@ -118,8 +131,7 @@ impl ProxyHttp for LB {
peer.options.verify_cert = false;
peer.options.verify_hostname = false;
}
if self.extraparams.load().to_https.unwrap_or(false) || innermap.to_https {
if ctx.to_https || innermap.to_https {
if let Some(stream) = session.stream() {
if stream.get_ssl().is_none() {
if let Some(addr) = session.server_addr() {
@@ -190,7 +202,7 @@ impl ProxyHttp for LB {
// }
async fn response_filter(&self, session: &mut Session, _upstream_response: &mut ResponseHeader, ctx: &mut Self::CTX) -> Result<()> {
// _upstream_response.insert_header("X-Proxied-From", "Fooooooooooooooo").unwrap();
if self.extraparams.load().sticky_sessions {
if ctx.extraparams.sticky_sessions {
let backend_id = ctx.backend_id.clone();
if let Some(bid) = self.ump_byid.get(&backend_id) {
let _ = _upstream_response.insert_header("set-cookie", format!("backend_id={}; Path=/; Max-Age=600; HttpOnly; SameSite=Lax", bid.address));