lemmy/crates/utils/src/rate_limit/mod.rs

185 lines
5.4 KiB
Rust
Raw Normal View History

use crate::{error::LemmyError, utils::get_ip, IpAddr};
use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
use futures::future::{ok, Ready};
use rate_limiter::{RateLimitType, RateLimiter};
use serde::{Deserialize, Serialize};
2020-05-16 14:04:08 +00:00
use std::{
future::Future,
pin::Pin,
rc::Rc,
2022-07-11 17:12:12 +00:00
sync::{Arc, Mutex},
2020-05-16 14:04:08 +00:00
task::{Context, Poll},
};
use typed_builder::TypedBuilder;
2020-05-16 14:04:08 +00:00
pub mod rate_limiter;
#[derive(Debug, Deserialize, Serialize, Clone, TypedBuilder)]
pub struct RateLimitConfig {
#[builder(default = 180)]
/// Maximum number of messages created in interval
pub message: i32,
#[builder(default = 60)]
/// Interval length for message limit, in seconds
pub message_per_second: i32,
#[builder(default = 6)]
/// Maximum number of posts created in interval
pub post: i32,
#[builder(default = 300)]
/// Interval length for post limit, in seconds
pub post_per_second: i32,
#[builder(default = 3)]
/// Maximum number of registrations in interval
pub register: i32,
#[builder(default = 3600)]
/// Interval length for registration limit, in seconds
pub register_per_second: i32,
#[builder(default = 6)]
/// Maximum number of image uploads in interval
pub image: i32,
#[builder(default = 3600)]
/// Interval length for image uploads, in seconds
pub image_per_second: i32,
#[builder(default = 6)]
/// Maximum number of comments created in interval
pub comment: i32,
#[builder(default = 600)]
/// Interval length for comment limit, in seconds
pub comment_per_second: i32,
#[builder(default = 60)]
/// Maximum number of searches created in interval
pub search: i32,
#[builder(default = 600)]
/// Interval length for search limit, in seconds
pub search_per_second: i32,
}
#[derive(Debug, Clone)]
2020-04-20 17:51:42 +00:00
pub struct RateLimit {
// it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
// across await points
pub rate_limiter: Arc<Mutex<RateLimiter>>,
pub rate_limit_config: RateLimitConfig,
2020-04-20 17:51:42 +00:00
}
#[derive(Debug, Clone)]
2020-04-20 17:51:42 +00:00
pub struct RateLimited {
rate_limiter: Arc<Mutex<RateLimiter>>,
rate_limit_config: RateLimitConfig,
2020-04-20 17:51:42 +00:00
type_: RateLimitType,
}
2020-04-20 17:51:42 +00:00
pub struct RateLimitedMiddleware<S> {
rate_limited: RateLimited,
service: Rc<S>,
2020-04-20 17:51:42 +00:00
}
impl RateLimit {
pub fn message(&self) -> RateLimited {
self.kind(RateLimitType::Message)
}
pub fn post(&self) -> RateLimited {
self.kind(RateLimitType::Post)
}
pub fn register(&self) -> RateLimited {
self.kind(RateLimitType::Register)
}
pub fn image(&self) -> RateLimited {
self.kind(RateLimitType::Image)
}
2021-11-11 20:40:25 +00:00
pub fn comment(&self) -> RateLimited {
self.kind(RateLimitType::Comment)
}
pub fn search(&self) -> RateLimited {
self.kind(RateLimitType::Search)
}
fn kind(&self, type_: RateLimitType) -> RateLimited {
2020-04-20 17:51:42 +00:00
RateLimited {
rate_limiter: self.rate_limiter.clone(),
rate_limit_config: self.rate_limit_config.clone(),
2020-04-20 17:51:42 +00:00
type_,
}
}
}
impl RateLimited {
/// Returns true if the request passed the rate limit, false if it failed and should be rejected.
pub fn check(self, ip_addr: IpAddr) -> bool {
// 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
let rate_limit = self.rate_limit_config;
let (kind, interval) = match self.type_ {
RateLimitType::Message => (rate_limit.message, rate_limit.message_per_second),
RateLimitType::Post => (rate_limit.post, rate_limit.post_per_second),
RateLimitType::Register => (rate_limit.register, rate_limit.register_per_second),
RateLimitType::Image => (rate_limit.image, rate_limit.image_per_second),
RateLimitType::Comment => (rate_limit.comment, rate_limit.comment_per_second),
RateLimitType::Search => (rate_limit.search, rate_limit.search_per_second),
};
2022-07-11 17:12:12 +00:00
let mut limiter = self.rate_limiter.lock().expect("mutex poison error");
limiter.check_rate_limit_full(self.type_, &ip_addr, kind, interval)
}
}
impl<S> Transform<S, ServiceRequest> for RateLimited
where
S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
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 {
rate_limited: self.clone(),
service: Rc::new(service),
2020-04-20 17:51:42 +00:00
})
}
}
2020-04-20 04:43:30 +00:00
type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
where
S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
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>>>;
fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
2020-04-20 17:51:42 +00:00
self.service.poll_ready(cx)
}
fn call(&self, req: ServiceRequest) -> Self::Future {
2020-04-20 18:02:25 +00:00
let ip_addr = get_ip(&req.connection_info());
let rate_limited = self.rate_limited.clone();
let service = self.service.clone();
Box::pin(async move {
if rate_limited.check(ip_addr) {
service.call(req).await
} else {
let (http_req, _) = req.into_parts();
Ok(ServiceResponse::from_err(
LemmyError::from_message("rate_limit_error"),
http_req,
))
}
})
}
}