1
0
Fork 0
mirror of https://gitlab.com/famedly/conduit.git synced 2025-07-22 17:18:35 +00:00

more rate limit targets

This commit is contained in:
Matthias Ahouansou 2024-06-26 09:16:54 +01:00
parent 02cea0bb93
commit d6abf5472b
No known key found for this signature in database
4 changed files with 123 additions and 67 deletions

View file

@ -1,7 +1,8 @@
use std::{
collections::BTreeMap,
iter::FromIterator,
str::{self},
net::IpAddr,
str::{self, FromStr},
};
use axum::{
@ -102,44 +103,6 @@ where
Token::None
};
// doesn't work when Conduit is behind proxy
// let remote_addr: ConnectInfo<SocketAddr> = parts.extract().await?;
let target = match &token {
Token::User((user_id, _)) => Some(Target::User(user_id.clone())),
Token::None => {
let header = parts
.headers
.get("x-forwarded-for")
.ok_or_else(|| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting."))?;
let s = header
.to_str()
.map_err(|_| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting."))?;
Some(
s.parse()
.map(Target::Ip)
.map_err(|_| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting.")),
)
.transpose()?
}
_ => None,
};
if let Err(retry_after_ms) = target.map_or(Ok(()), |t| {
let key = (t, (&metadata).into());
services()
.rate_limiting
.update_or_reject(&key)
.map_err(Some)
}) {
return Err(Error::BadRequest(
ErrorKind::LimitExceeded { retry_after_ms },
"Rate limit exceeded.",
));
}
let mut json_body = serde_json::from_slice::<CanonicalJsonValue>(&body).ok();
let (sender_user, sender_device, sender_servername, appservice_info) =
@ -350,8 +313,64 @@ where
}
};
// doesn't work when Conduit is behind proxy
// let remote_addr: ConnectInfo<SocketAddr> = parts.extract().await?;
let headers = parts.headers;
let target = if let Some(server) = sender_servername.clone() {
Target::Server(server)
// Token::User((user_id, _)) => Some(Target::User(user_id.clone())),
// Token::None => {
// let header = parts
// .headers
// .get("x-forwarded-for")
// .ok_or_else(|| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting."))?;
// let s = header
// .to_str()
// .map_err(|_| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting."))?;
// Some(
// s.parse()
// .map(Target::Ip)
// .map_err(|_| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting.")),
// )
// .transpose()?
// }
// _ => None,
} else if let Some(appservice) = appservice_info.clone() {
Target::Appservice(appservice.registration.id)
} else if let Some(user) = sender_user.clone() {
Target::User(user)
} else {
let ip = headers
.get("X-Forwarded-For")
.and_then(|header| header.to_str().ok())
.map(|header| header.split_once(",").map(|(ip, _)| ip).unwrap_or(header))
.and_then(|ip| IpAddr::from_str(ip).ok());
if let Some(ip) = ip {
Target::Ip(ip)
} else {
Target::None
}
};
if let Err(retry_after_ms) = {
services()
.rate_limiting
.update_or_reject(target, metadata)
.map_err(Some)
} {
return Err(Error::BadRequest(
ErrorKind::LimitExceeded { retry_after_ms },
"Rate limit exceeded.",
));
}
let mut http_request = Request::builder().uri(parts.uri).method(parts.method);
*http_request.headers_mut().unwrap() = parts.headers;
*http_request.headers_mut().unwrap() = headers;
if let Some(CanonicalJsonValue::Object(json_body)) = &mut json_body {
let user_id = sender_user.clone().unwrap_or_else(|| {