crates/utils/src/rate_limit/mod.rs

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