2023-07-10 14:50:07 +00:00
|
|
|
use crate::error::{LemmyError, LemmyErrorType};
|
2023-02-16 04:05:14 +00:00
|
|
|
use actix_web::dev::{ConnectionInfo, Service, ServiceRequest, ServiceResponse, Transform};
|
2023-10-19 13:31:51 +00:00
|
|
|
use enum_map::{enum_map, EnumMap};
|
2020-12-21 23:27:42 +00:00
|
|
|
use futures::future::{ok, Ready};
|
2023-10-19 13:31:51 +00:00
|
|
|
pub use rate_limiter::{ActionType, BucketConfig};
|
|
|
|
use rate_limiter::{InstantSecs, RateLimitState};
|
2020-05-16 14:04:08 +00:00
|
|
|
use std::{
|
|
|
|
future::Future,
|
2023-06-21 08:28:20 +00:00
|
|
|
net::{IpAddr, Ipv4Addr, SocketAddr},
|
2020-05-16 14:04:08 +00:00
|
|
|
pin::Pin,
|
2022-03-25 15:41:38 +00:00
|
|
|
rc::Rc,
|
2023-06-21 08:28:20 +00:00
|
|
|
str::FromStr,
|
2022-07-11 17:12:12 +00:00
|
|
|
sync::{Arc, Mutex},
|
2020-05-16 14:04:08 +00:00
|
|
|
task::{Context, Poll},
|
2023-06-21 08:28:20 +00:00
|
|
|
time::Duration,
|
2020-05-16 14:04:08 +00:00
|
|
|
};
|
2020-04-19 22:08:25 +00:00
|
|
|
|
2020-05-16 14:04:08 +00:00
|
|
|
pub mod rate_limiter;
|
|
|
|
|
2020-04-20 03:59:07 +00:00
|
|
|
#[derive(Debug, Clone)]
|
2023-10-19 13:31:51 +00:00
|
|
|
pub struct RateLimitChecker {
|
|
|
|
state: Arc<Mutex<RateLimitState>>,
|
|
|
|
action_type: ActionType,
|
2020-04-20 17:51:42 +00:00
|
|
|
}
|
2020-04-20 03:59:07 +00:00
|
|
|
|
2022-11-16 19:06:22 +00:00
|
|
|
/// Single instance of rate limit config and buckets, which is shared across all threads.
|
|
|
|
#[derive(Clone)]
|
|
|
|
pub struct RateLimitCell {
|
2023-10-19 13:31:51 +00:00
|
|
|
state: Arc<Mutex<RateLimitState>>,
|
2020-04-20 17:51:42 +00:00
|
|
|
}
|
2020-04-20 03:59:07 +00:00
|
|
|
|
2022-11-16 19:06:22 +00:00
|
|
|
impl RateLimitCell {
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn new(rate_limit_config: EnumMap<ActionType, BucketConfig>) -> Self {
|
|
|
|
let state = Arc::new(Mutex::new(RateLimitState::new(rate_limit_config)));
|
2022-11-16 19:06:22 +00:00
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
let state_weak_ref = Arc::downgrade(&state);
|
2022-11-16 19:06:22 +00:00
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
tokio::spawn(async move {
|
2023-12-12 18:06:17 +00:00
|
|
|
let interval = Duration::from_secs(120);
|
2023-06-21 08:28:20 +00:00
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
// This loop stops when all other references to `state` are dropped
|
|
|
|
while let Some(state) = state_weak_ref.upgrade() {
|
2023-12-12 18:06:17 +00:00
|
|
|
tokio::time::sleep(interval).await;
|
2023-10-19 13:31:51 +00:00
|
|
|
state
|
|
|
|
.lock()
|
|
|
|
.expect("Failed to lock rate limit mutex for reading")
|
|
|
|
.remove_full_buckets(InstantSecs::now());
|
|
|
|
}
|
|
|
|
});
|
2023-06-21 08:28:20 +00:00
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
RateLimitCell { state }
|
|
|
|
}
|
2023-06-21 08:28:20 +00:00
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn set_config(&self, config: EnumMap<ActionType, BucketConfig>) {
|
|
|
|
self
|
|
|
|
.state
|
|
|
|
.lock()
|
|
|
|
.expect("Failed to lock rate limit mutex for updating")
|
|
|
|
.set_config(config);
|
2023-06-21 08:28:20 +00:00
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn message(&self) -> RateLimitChecker {
|
|
|
|
self.new_checker(ActionType::Message)
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn post(&self) -> RateLimitChecker {
|
|
|
|
self.new_checker(ActionType::Post)
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn register(&self) -> RateLimitChecker {
|
|
|
|
self.new_checker(ActionType::Register)
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn image(&self) -> RateLimitChecker {
|
|
|
|
self.new_checker(ActionType::Image)
|
2020-08-05 16:00:00 +00:00
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn comment(&self) -> RateLimitChecker {
|
|
|
|
self.new_checker(ActionType::Comment)
|
2021-11-11 20:40:25 +00:00
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn search(&self) -> RateLimitChecker {
|
|
|
|
self.new_checker(ActionType::Search)
|
2022-03-29 15:46:03 +00:00
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
pub fn import_user_settings(&self) -> RateLimitChecker {
|
|
|
|
self.new_checker(ActionType::ImportUserSettings)
|
2023-10-11 14:47:22 +00:00
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
fn new_checker(&self, action_type: ActionType) -> RateLimitChecker {
|
|
|
|
RateLimitChecker {
|
|
|
|
state: self.state.clone(),
|
|
|
|
action_type,
|
2020-04-20 17:51:42 +00:00
|
|
|
}
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
2023-10-19 13:31:51 +00:00
|
|
|
|
|
|
|
pub fn with_test_config() -> Self {
|
|
|
|
Self::new(enum_map! {
|
|
|
|
ActionType::Message => BucketConfig {
|
|
|
|
capacity: 180,
|
|
|
|
secs_to_refill: 60,
|
|
|
|
},
|
|
|
|
ActionType::Post => BucketConfig {
|
|
|
|
capacity: 6,
|
|
|
|
secs_to_refill: 300,
|
|
|
|
},
|
|
|
|
ActionType::Register => BucketConfig {
|
|
|
|
capacity: 3,
|
|
|
|
secs_to_refill: 3600,
|
|
|
|
},
|
|
|
|
ActionType::Image => BucketConfig {
|
|
|
|
capacity: 6,
|
|
|
|
secs_to_refill: 3600,
|
|
|
|
},
|
|
|
|
ActionType::Comment => BucketConfig {
|
|
|
|
capacity: 6,
|
|
|
|
secs_to_refill: 600,
|
|
|
|
},
|
|
|
|
ActionType::Search => BucketConfig {
|
|
|
|
capacity: 60,
|
|
|
|
secs_to_refill: 600,
|
|
|
|
},
|
|
|
|
ActionType::ImportUserSettings => BucketConfig {
|
|
|
|
capacity: 1,
|
|
|
|
secs_to_refill: 24 * 60 * 60,
|
|
|
|
},
|
|
|
|
})
|
|
|
|
}
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
|
2022-11-16 19:06:22 +00:00
|
|
|
pub struct RateLimitedMiddleware<S> {
|
2023-10-19 13:31:51 +00:00
|
|
|
checker: RateLimitChecker,
|
2022-11-16 19:06:22 +00:00
|
|
|
service: Rc<S>,
|
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
impl RateLimitChecker {
|
2022-03-25 15:41:38 +00:00
|
|
|
/// Returns true if the request passed the rate limit, false if it failed and should be rejected.
|
2022-03-27 00:29:05 +00:00
|
|
|
pub fn check(self, ip_addr: IpAddr) -> bool {
|
2020-07-01 12:54:29 +00:00
|
|
|
// Does not need to be blocking because the RwLock in settings never held across await points,
|
|
|
|
// and the operation here locks only long enough to clone
|
2023-10-19 13:31:51 +00:00
|
|
|
let mut state = self
|
|
|
|
.state
|
2022-11-16 19:06:22 +00:00
|
|
|
.lock()
|
|
|
|
.expect("Failed to lock rate limit mutex for reading");
|
2022-03-27 00:29:05 +00:00
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
state.check(self.action_type, ip_addr, InstantSecs::now())
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
impl<S> Transform<S, ServiceRequest> for RateLimitChecker
|
2020-04-20 03:59:07 +00:00
|
|
|
where
|
2022-03-25 15:41:38 +00:00
|
|
|
S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
|
2020-04-20 03:59:07 +00:00
|
|
|
S::Future: 'static,
|
|
|
|
{
|
|
|
|
type Response = S::Response;
|
|
|
|
type Error = actix_web::Error;
|
|
|
|
type InitError = ();
|
|
|
|
type Transform = RateLimitedMiddleware<S>;
|
|
|
|
type Future = Ready<Result<Self::Transform, Self::InitError>>;
|
|
|
|
|
|
|
|
fn new_transform(&self, service: S) -> Self::Future {
|
2020-04-20 17:51:42 +00:00
|
|
|
ok(RateLimitedMiddleware {
|
2023-10-19 13:31:51 +00:00
|
|
|
checker: self.clone(),
|
2022-03-25 15:41:38 +00:00
|
|
|
service: Rc::new(service),
|
2020-04-20 17:51:42 +00:00
|
|
|
})
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-04-20 04:43:30 +00:00
|
|
|
type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
|
|
|
|
|
2021-07-06 13:26:46 +00:00
|
|
|
impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
|
2020-04-20 03:59:07 +00:00
|
|
|
where
|
2022-03-25 15:41:38 +00:00
|
|
|
S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
|
2020-04-20 03:59:07 +00:00
|
|
|
S::Future: 'static,
|
|
|
|
{
|
|
|
|
type Response = S::Response;
|
|
|
|
type Error = actix_web::Error;
|
2020-04-20 04:43:30 +00:00
|
|
|
type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
|
2020-04-20 03:59:07 +00:00
|
|
|
|
2021-07-06 13:26:46 +00:00
|
|
|
fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
2020-04-20 17:51:42 +00:00
|
|
|
self.service.poll_ready(cx)
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
|
|
|
|
2021-07-06 13:26:46 +00:00
|
|
|
fn call(&self, req: ServiceRequest) -> Self::Future {
|
2020-04-20 18:02:25 +00:00
|
|
|
let ip_addr = get_ip(&req.connection_info());
|
2020-04-20 03:59:07 +00:00
|
|
|
|
2023-10-19 13:31:51 +00:00
|
|
|
let checker = self.checker.clone();
|
2022-03-25 15:41:38 +00:00
|
|
|
let service = self.service.clone();
|
|
|
|
|
|
|
|
Box::pin(async move {
|
2023-10-19 13:31:51 +00:00
|
|
|
if checker.check(ip_addr) {
|
2022-03-25 15:41:38 +00:00
|
|
|
service.call(req).await
|
|
|
|
} else {
|
|
|
|
let (http_req, _) = req.into_parts();
|
2022-11-03 18:13:40 +00:00
|
|
|
Ok(ServiceResponse::from_err(
|
2023-07-10 14:50:07 +00:00
|
|
|
LemmyError::from(LemmyErrorType::RateLimitError),
|
2022-03-25 15:41:38 +00:00
|
|
|
http_req,
|
|
|
|
))
|
|
|
|
}
|
|
|
|
})
|
2020-04-20 03:59:07 +00:00
|
|
|
}
|
2020-04-19 22:08:25 +00:00
|
|
|
}
|
2023-02-16 04:05:14 +00:00
|
|
|
|
|
|
|
fn get_ip(conn_info: &ConnectionInfo) -> IpAddr {
|
2023-06-21 08:28:20 +00:00
|
|
|
conn_info
|
|
|
|
.realip_remote_addr()
|
|
|
|
.and_then(parse_ip)
|
|
|
|
.unwrap_or(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))
|
|
|
|
}
|
|
|
|
|
|
|
|
fn parse_ip(addr: &str) -> Option<IpAddr> {
|
|
|
|
if let Some(s) = addr.strip_suffix(']') {
|
|
|
|
IpAddr::from_str(s.get(1..)?).ok()
|
|
|
|
} else if let Ok(ip) = IpAddr::from_str(addr) {
|
|
|
|
Some(ip)
|
|
|
|
} else if let Ok(socket) = SocketAddr::from_str(addr) {
|
|
|
|
Some(socket.ip())
|
|
|
|
} else {
|
|
|
|
None
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[cfg(test)]
|
2024-03-25 18:42:10 +00:00
|
|
|
#[allow(clippy::unwrap_used)]
|
|
|
|
#[allow(clippy::indexing_slicing)]
|
2023-06-21 08:28:20 +00:00
|
|
|
mod tests {
|
2023-07-17 15:04:14 +00:00
|
|
|
|
2023-06-21 08:28:20 +00:00
|
|
|
#[test]
|
|
|
|
fn test_parse_ip() {
|
|
|
|
let ip_addrs = [
|
|
|
|
"1.2.3.4",
|
|
|
|
"1.2.3.4:8000",
|
|
|
|
"2001:db8::",
|
|
|
|
"[2001:db8::]",
|
|
|
|
"[2001:db8::]:8000",
|
|
|
|
];
|
|
|
|
for addr in ip_addrs {
|
|
|
|
assert!(super::parse_ip(addr).is_some(), "failed to parse {addr}");
|
|
|
|
}
|
|
|
|
}
|
2023-02-16 04:05:14 +00:00
|
|
|
}
|