crates/utils/src/rate_limit/mod.rs

   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};
   6 use std::{
   7   future::Future,
   8   pin::Pin,
   9   rc::Rc,
  10   sync::{Arc, Mutex},
  11   task::{Context, Poll},
  12 };
  13 use typed_builder::TypedBuilder;
  14
  15 pub mod rate_limiter;
  16
  17 #[derive(Debug, Deserialize, Serialize, Clone, TypedBuilder)]
  18 pub struct RateLimitConfig {
  19   #[builder(default = 180)]
  20   /// Maximum number of messages created in interval
  21   pub message: i32,
  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
  27   pub post: i32,
  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
  33   pub register: i32,
  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
  39   pub image: i32,
  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
  45   pub comment: i32,
  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
  51   pub search: i32,
  52   #[builder(default = 600)]
  53   /// Interval length for search limit, in seconds
  54   pub search_per_second: i32,
  55 }
  56
  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,
  63 }
  64
  65 #[derive(Debug, Clone)]
  66 pub struct RateLimited {
  67   rate_limiter: Arc<Mutex<RateLimiter>>,
  68   rate_limit_config: RateLimitConfig,
  69   type_: RateLimitType,
  70 }
  71
  72 pub struct RateLimitedMiddleware<S> {
  73   rate_limited: RateLimited,
  74   service: Rc<S>,
  75 }
  76
  77 impl RateLimit {
  78   pub fn message(&self) -> RateLimited {
  79     self.kind(RateLimitType::Message)
  80   }
  81
  82   pub fn post(&self) -> RateLimited {
  83     self.kind(RateLimitType::Post)
  84   }
  85
  86   pub fn register(&self) -> RateLimited {
  87     self.kind(RateLimitType::Register)
  88   }
  89
  90   pub fn image(&self) -> RateLimited {
  91     self.kind(RateLimitType::Image)
  92   }
  93
  94   pub fn comment(&self) -> RateLimited {
  95     self.kind(RateLimitType::Comment)
  96   }
  97
  98   pub fn search(&self) -> RateLimited {
  99     self.kind(RateLimitType::Search)
 100   }
 101
 102   fn kind(&self, type_: RateLimitType) -> RateLimited {
 103     RateLimited {
 104       rate_limiter: self.rate_limiter.clone(),
 105       rate_limit_config: self.rate_limit_config.clone(),
 106       type_,
 107     }
 108   }
 109 }
 110
 111 impl RateLimited {
 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;
 117
 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),
 125     };
 126     let mut limiter = self.rate_limiter.lock().expect("mutex poison error");
 127
 128     limiter.check_rate_limit_full(self.type_, &ip_addr, kind, interval)
 129   }
 130 }
 131
 132 impl<S> Transform<S, ServiceRequest> for RateLimited
 133 where
 134   S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
 135   S::Future: 'static,
 136 {
 137   type Response = S::Response;
 138   type Error = actix_web::Error;
 139   type InitError = ();
 140   type Transform = RateLimitedMiddleware<S>;
 141   type Future = Ready<Result<Self::Transform, Self::InitError>>;
 142
 143   fn new_transform(&self, service: S) -> Self::Future {
 144     ok(RateLimitedMiddleware {
 145       rate_limited: self.clone(),
 146       service: Rc::new(service),
 147     })
 148   }
 149 }
 150
 151 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
 152
 153 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
 154 where
 155   S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error> + 'static,
 156   S::Future: 'static,
 157 {
 158   type Response = S::Response;
 159   type Error = actix_web::Error;
 160   type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
 161
 162   fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
 163     self.service.poll_ready(cx)
 164   }
 165
 166   fn call(&self, req: ServiceRequest) -> Self::Future {
 167     let ip_addr = get_ip(&req.connection_info());
 168
 169     let rate_limited = self.rate_limited.clone();
 170     let service = self.service.clone();
 171
 172     Box::pin(async move {
 173       if rate_limited.check(ip_addr) {
 174         service.call(req).await
 175       } else {
 176         let (http_req, _) = req.into_parts();
 177         Ok(ServiceResponse::from_err(
 178           LemmyError::from_message("rate_limit_error"),
 179           http_req,
 180         ))
 181       }
 182     })
 183   }
 184 }