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   pub fn comment(&self) -> RateLimited {
  53     self.kind(RateLimitType::Comment)
  54   }
  55
  56   fn kind(&self, type_: RateLimitType) -> RateLimited {
  57     RateLimited {
  58       rate_limiter: self.rate_limiter.clone(),
  59       rate_limit_config: self.rate_limit_config.clone(),
  60       type_,
  61     }
  62   }
  63 }
  64
  65 impl RateLimited {
  66   pub async fn wrap<T, E>(
  67     self,
  68     ip_addr: IpAddr,
  69     fut: impl Future<Output = Result<T, E>>,
  70   ) -> Result<T, E>
  71   where
  72     E: From<LemmyError>,
  73   {
  74     // Does not need to be blocking because the RwLock in settings never held across await points,
  75     // and the operation here locks only long enough to clone
  76     let rate_limit = self.rate_limit_config;
  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         RateLimitType::Comment => {
 123           limiter.check_rate_limit_full(
 124             self.type_,
 125             &ip_addr,
 126             rate_limit.comment,
 127             rate_limit.comment_per_second,
 128             false,
 129           )?;
 130         }
 131       };
 132     }
 133
 134     let res = fut.await;
 135
 136     // after
 137     {
 138       let mut limiter = self.rate_limiter.lock().await;
 139       if res.is_ok() {
 140         match self.type_ {
 141           RateLimitType::Post => {
 142             limiter.check_rate_limit_full(
 143               self.type_,
 144               &ip_addr,
 145               rate_limit.post,
 146               rate_limit.post_per_second,
 147               false,
 148             )?;
 149           }
 150           RateLimitType::Register => {
 151             limiter.check_rate_limit_full(
 152               self.type_,
 153               &ip_addr,
 154               rate_limit.register,
 155               rate_limit.register_per_second,
 156               false,
 157             )?;
 158           }
 159           _ => (),
 160         };
 161       }
 162     }
 163
 164     res
 165   }
 166 }
 167
 168 impl<S> Transform<S, ServiceRequest> for RateLimited
 169 where
 170   S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
 171   S::Future: 'static,
 172 {
 173   type Response = S::Response;
 174   type Error = actix_web::Error;
 175   type InitError = ();
 176   type Transform = RateLimitedMiddleware<S>;
 177   type Future = Ready<Result<Self::Transform, Self::InitError>>;
 178
 179   fn new_transform(&self, service: S) -> Self::Future {
 180     ok(RateLimitedMiddleware {
 181       rate_limited: self.clone(),
 182       service,
 183     })
 184   }
 185 }
 186
 187 type FutResult<T, E> = dyn Future<Output = Result<T, E>>;
 188
 189 impl<S> Service<ServiceRequest> for RateLimitedMiddleware<S>
 190 where
 191   S: Service<ServiceRequest, Response = ServiceResponse, Error = actix_web::Error>,
 192   S::Future: 'static,
 193 {
 194   type Response = S::Response;
 195   type Error = actix_web::Error;
 196   type Future = Pin<Box<FutResult<Self::Response, Self::Error>>>;
 197
 198   fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
 199     self.service.poll_ready(cx)
 200   }
 201
 202   fn call(&self, req: ServiceRequest) -> Self::Future {
 203     let ip_addr = get_ip(&req.connection_info());
 204
 205     let fut = self
 206       .rate_limited
 207       .clone()
 208       .wrap(ip_addr, self.service.call(req));
 209
 210     Box::pin(async move { fut.await.map_err(actix_web::Error::from) })
 211   }
 212 }