1 use crate::{error::LemmyError, utils::get_ip, IpAddr};
2 use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
3 use futures::future::{ok, Ready};
4 use rate_limiter::{RateLimitType, RateLimiter};
5 use serde::{Deserialize, Serialize};
11 task::{Context, Poll},
13 use typed_builder::TypedBuilder;
17 #[derive(Debug, Deserialize, Serialize, Clone, TypedBuilder)]
18 pub struct RateLimitConfig {
19 #[builder(default = 180)]
20 /// Maximum number of messages created in interval
22 #[builder(default = 60)]
23 /// Interval length for message limit, in seconds
24 pub message_per_second: i32,
25 #[builder(default = 6)]
26 /// Maximum number of posts created in interval
28 #[builder(default = 300)]
29 /// Interval length for post limit, in seconds
30 pub post_per_second: i32,
31 #[builder(default = 3)]
32 /// Maximum number of registrations in interval
34 #[builder(default = 3600)]
35 /// Interval length for registration limit, in seconds
36 pub register_per_second: i32,
37 #[builder(default = 6)]
38 /// Maximum number of image uploads in interval
40 #[builder(default = 3600)]
41 /// Interval length for image uploads, in seconds
42 pub image_per_second: i32,
43 #[builder(default = 6)]
44 /// Maximum number of comments created in interval
46 #[builder(default = 600)]
47 /// Interval length for comment limit, in seconds
48 pub comment_per_second: i32,
49 #[builder(default = 60)]
50 /// Maximum number of searches created in interval
52 #[builder(default = 600)]
53 /// Interval length for search limit, in seconds
54 pub search_per_second: i32,
57 #[derive(Debug, Clone)]
58 pub struct RateLimit {
59 // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
60 // across await points
61 pub rate_limiter: Arc<Mutex<RateLimiter>>,
62 pub rate_limit_config: RateLimitConfig,
65 #[derive(Debug, Clone)]
66 pub struct RateLimited {
67 rate_limiter: Arc<Mutex<RateLimiter>>,
68 rate_limit_config: RateLimitConfig,
72 pub struct RateLimitedMiddleware<S> {
73 rate_limited: RateLimited,
78 pub fn message(&self) -> RateLimited {
79 self.kind(RateLimitType::Message)
82 pub fn post(&self) -> RateLimited {
83 self.kind(RateLimitType::Post)
86 pub fn register(&self) -> RateLimited {
87 self.kind(RateLimitType::Register)
90 pub fn image(&self) -> RateLimited {
91 self.kind(RateLimitType::Image)
94 pub fn comment(&self) -> RateLimited {
95 self.kind(RateLimitType::Comment)
98 pub fn search(&self) -> RateLimited {
99 self.kind(RateLimitType::Search)
102 fn kind(&self, type_: RateLimitType) -> RateLimited {
104 rate_limiter: self.rate_limiter.clone(),
105 rate_limit_config: self.rate_limit_config.clone(),
112 /// Returns true if the request passed the rate limit, false if it failed and should be rejected.
113 pub fn check(self, ip_addr: IpAddr) -> bool {
114 // Does not need to be blocking because the RwLock in settings never held across await points,
115 // and the operation here locks only long enough to clone
116 let rate_limit = self.rate_limit_config;
118 let (kind, interval) = match self.type_ {
119 RateLimitType::Message => (rate_limit.message, rate_limit.message_per_second),
120 RateLimitType::Post => (rate_limit.post, rate_limit.post_per_second),
121 RateLimitType::Register => (rate_limit.register, rate_limit.register_per_second),
122 RateLimitType::Image => (rate_limit.image, rate_limit.image_per_second),
123 RateLimitType::Comment => (rate_limit.comment, rate_limit.comment_per_second),
124 RateLimitType::Search => (rate_limit.search, rate_limit.search_per_second),
126 let mut limiter = self.rate_limiter.lock().expect("mutex poison error");
128 limiter.check_rate_limit_full(self.type_, &ip_addr, kind, interval)
132 impl<S> Transform<S, ServiceRequest> for RateLimited
134 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
137 type Response = S::Response;
138 type Error = actix_web::Error;
140 type Transform = RateLimitedMiddleware<S>;
141 type Future = Ready<Result<Self::Transform, Self::InitError>>;
143 fn new_transform(&self, service: S) -> Self::Future {
144 ok(RateLimitedMiddleware {
145 rate_limited: self.clone(),
146 service: Rc::new(service),
151 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
153 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
155 S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
158 type Response = S::Response;
159 type Error = actix_web::Error;
160 type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
162 fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
163 self.service.poll_ready(cx)
166 fn call(&self, req: ServiceRequest) -> Self::Future {
167 let ip_addr = get_ip(&req.connection_info());
169 let rate_limited = self.rate_limited.clone();
170 let service = self.service.clone();
172 Box::pin(async move {
173 if rate_limited.check(ip_addr) {
174 service.call(req).await
176 let (http_req, _) = req.into_parts();
177 Ok(ServiceResponse::from_err(
178 LemmyError::from_message("rate_limit_error"),