crates/utils/src/rate_limit/rate_limiter.rs

   1 use crate::{ApiError, IpAddr, LemmyError};
   2 use log::debug;
   3 use std::{collections::HashMap, time::SystemTime};
   4 use strum::IntoEnumIterator;
   5
   6 #[derive(Debug, Clone)]
   7 struct RateLimitBucket {
   8   last_checked: SystemTime,
   9   allowance: f64,
  10 }
  11
  12 #[derive(Eq, PartialEq, Hash, Debug, EnumIter, Copy, Clone, AsRefStr)]
  13 pub(crate) enum RateLimitType {
  14   Message,
  15   Register,
  16   Post,
  17   Image,
  18 }
  19
  20 /// Rate limiting based on rate type and IP addr
  21 #[derive(Debug, Clone, Default)]
  22 pub struct RateLimiter {
  23   buckets: HashMap<RateLimitType, HashMap<IpAddr, RateLimitBucket>>,
  24 }
  25
  26 impl RateLimiter {
  27   fn insert_ip(&mut self, ip: &IpAddr) {
  28     for rate_limit_type in RateLimitType::iter() {
  29       if self.buckets.get(&rate_limit_type).is_none() {
  30         self.buckets.insert(rate_limit_type, HashMap::new());
  31       }
  32
  33       if let Some(bucket) = self.buckets.get_mut(&rate_limit_type) {
  34         if bucket.get(ip).is_none() {
  35           bucket.insert(
  36             ip.clone(),
  37             RateLimitBucket {
  38               last_checked: SystemTime::now(),
  39               allowance: -2f64,
  40             },
  41           );
  42         }
  43       }
  44     }
  45   }
  46
  47   #[allow(clippy::float_cmp)]
  48   pub(super) fn check_rate_limit_full(
  49     &mut self,
  50     type_: RateLimitType,
  51     ip: &IpAddr,
  52     rate: i32,
  53     per: i32,
  54     check_only: bool,
  55   ) -> Result<(), LemmyError> {
  56     self.insert_ip(ip);
  57     if let Some(bucket) = self.buckets.get_mut(&type_) {
  58       if let Some(rate_limit) = bucket.get_mut(ip) {
  59         let current = SystemTime::now();
  60         let time_passed = current.duration_since(rate_limit.last_checked)?.as_secs() as f64;
  61
  62         // The initial value
  63         if rate_limit.allowance == -2f64 {
  64           rate_limit.allowance = rate as f64;
  65         };
  66
  67         rate_limit.last_checked = current;
  68         rate_limit.allowance += time_passed * (rate as f64 / per as f64);
  69         if !check_only && rate_limit.allowance > rate as f64 {
  70           rate_limit.allowance = rate as f64;
  71         }
  72
  73         if rate_limit.allowance < 1.0 {
  74           debug!(
  75             "Rate limited type: {}, IP: {}, time_passed: {}, allowance: {}",
  76             type_.as_ref(),
  77             ip,
  78             time_passed,
  79             rate_limit.allowance
  80           );
  81           Err(
  82             ApiError {
  83               message: format!(
  84                 "Too many requests. type: {}, IP: {}, {} per {} seconds",
  85                 type_.as_ref(),
  86                 ip,
  87                 rate,
  88                 per
  89               ),
  90             }
  91             .into(),
  92           )
  93         } else {
  94           if !check_only {
  95             rate_limit.allowance -= 1.0;
  96           }
  97           Ok(())
  98         }
  99       } else {
 100         Ok(())
 101       }
 102     } else {
 103       Ok(())
 104     }
 105   }
 106 }