server/src/rate_limit/mod.rs

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