crates/utils/src/rate_limit/mod.rs

   1 use crate::{
   2   settings::structs::{RateLimitConfig, Settings},
   3   utils::get_ip,
   4   IpAddr,
   5   LemmyError,
   6 };
   7 use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
   8 use futures::future::{ok, Ready};
   9 use rate_limiter::{RateLimitType, RateLimiter};
  10 use std::{
  11   future::Future,
  12   pin::Pin,
  13   sync::Arc,
  14   task::{Context, Poll},
  15 };
  16 use tokio::sync::Mutex;
  17
  18 pub mod rate_limiter;
  19
  20 #[derive(Debug, Clone)]
  21 pub struct RateLimit {
  22   // it might be reasonable to use a std::sync::Mutex here, since we don't need to lock this
  23   // across await points
  24   pub rate_limiter: Arc<Mutex<RateLimiter>>,
  25 }
  26
  27 #[derive(Debug, Clone)]
  28 pub struct RateLimited {
  29   rate_limiter: Arc<Mutex<RateLimiter>>,
  30   type_: RateLimitType,
  31 }
  32
  33 pub struct RateLimitedMiddleware<S> {
  34   rate_limited: RateLimited,
  35   service: S,
  36 }
  37
  38 impl RateLimit {
  39   pub fn message(&self) -> RateLimited {
  40     self.kind(RateLimitType::Message)
  41   }
  42
  43   pub fn post(&self) -> RateLimited {
  44     self.kind(RateLimitType::Post)
  45   }
  46
  47   pub fn register(&self) -> RateLimited {
  48     self.kind(RateLimitType::Register)
  49   }
  50
  51   pub fn image(&self) -> RateLimited {
  52     self.kind(RateLimitType::Image)
  53   }
  54
  55   fn kind(&self, type_: RateLimitType) -> RateLimited {
  56     RateLimited {
  57       rate_limiter: self.rate_limiter.clone(),
  58       type_,
  59     }
  60   }
  61 }
  62
  63 impl RateLimited {
  64   pub async fn wrap<T, E>(
  65     self,
  66     ip_addr: IpAddr,
  67     fut: impl Future<Output = Result<T, E>>,
  68   ) -> Result<T, E>
  69   where
  70     E: From<LemmyError>,
  71   {
  72     // Does not need to be blocking because the RwLock in settings never held across await points,
  73     // and the operation here locks only long enough to clone
  74     let rate_limit: RateLimitConfig = Settings::get()
  75       .rate_limit
  76       .unwrap_or_else(RateLimitConfig::default);
  77
  78     // before
  79     {
  80       let mut limiter = self.rate_limiter.lock().await;
  81
  82       match self.type_ {
  83         RateLimitType::Message => {
  84           limiter.check_rate_limit_full(
  85             self.type_,
  86             &ip_addr,
  87             rate_limit.message,
  88             rate_limit.message_per_second,
  89             false,
  90           )?;
  91
  92           drop(limiter);
  93           return fut.await;
  94         }
  95         RateLimitType::Post => {
  96           limiter.check_rate_limit_full(
  97             self.type_,
  98             &ip_addr,
  99             rate_limit.post,
 100             rate_limit.post_per_second,
 101             true,
 102           )?;
 103         }
 104         RateLimitType::Register => {
 105           limiter.check_rate_limit_full(
 106             self.type_,
 107             &ip_addr,
 108             rate_limit.register,
 109             rate_limit.register_per_second,
 110             true,
 111           )?;
 112         }
 113         RateLimitType::Image => {
 114           limiter.check_rate_limit_full(
 115             self.type_,
 116             &ip_addr,
 117             rate_limit.image,
 118             rate_limit.image_per_second,
 119             false,
 120           )?;
 121         }
 122       };
 123     }
 124
 125     let res = fut.await;
 126
 127     // after
 128     {
 129       let mut limiter = self.rate_limiter.lock().await;
 130       if res.is_ok() {
 131         match self.type_ {
 132           RateLimitType::Post => {
 133             limiter.check_rate_limit_full(
 134               self.type_,
 135               &ip_addr,
 136               rate_limit.post,
 137               rate_limit.post_per_second,
 138               false,
 139             )?;
 140           }
 141           RateLimitType::Register => {
 142             limiter.check_rate_limit_full(
 143               self.type_,
 144               &ip_addr,
 145               rate_limit.register,
 146               rate_limit.register_per_second,
 147               false,
 148             )?;
 149           }
 150           _ => (),
 151         };
 152       }
 153     }
 154
 155     res
 156   }
 157 }
 158
 159 impl<S> Transform<S, ServiceRequest> for RateLimited
 160 where
 161   S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
 162   S::Future: 'static,
 163 {
 164   type Response = S::Response;
 165   type Error = actix_web::Error;
 166   type InitError = ();
 167   type Transform = RateLimitedMiddleware<S>;
 168   type Future = Ready<Result<Self::Transform, Self::InitError>>;
 169
 170   fn new_transform(&self, service: S) -> Self::Future {
 171     ok(RateLimitedMiddleware {
 172       rate_limited: self.clone(),
 173       service,
 174     })
 175   }
 176 }
 177
 178 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
 179
 180 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
 181 where
 182   S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
 183   S::Future: 'static,
 184 {
 185   type Response = S::Response;
 186   type Error = actix_web::Error;
 187   type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
 188
 189   fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
 190     self.service.poll_ready(cx)
 191   }
 192
 193   fn call(&self, req: ServiceRequest) -> Self::Future {
 194     let ip_addr = get_ip(&req.connection_info());
 195
 196     let fut = self
 197       .rate_limited
 198       .clone()
 199       .wrap(ip_addr, self.service.call(req));
 200
 201     Box::pin(async move { fut.await.map_err(actix_web::Error::from) })
 202   }
 203 }