crates/utils/src/rate_limit/rate_limiter.rs

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