1 use crate::{ApiError, IpAddr, LemmyError};
3 use std::{collections::HashMap, time::SystemTime};
4 use strum::IntoEnumIterator;
6 #[derive(Debug, Clone)]
7 struct RateLimitBucket {
8 last_checked: SystemTime,
12 #[derive(Eq, PartialEq, Hash, Debug, EnumIter, Copy, Clone, AsRefStr)]
13 pub(crate) enum RateLimitType {
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>>,
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());
33 if let Some(bucket) = self.buckets.get_mut(&rate_limit_type) {
34 if bucket.get(ip).is_none() {
38 last_checked: SystemTime::now(),
47 #[allow(clippy::float_cmp)]
48 pub(super) fn check_rate_limit_full(
55 ) -> Result<(), LemmyError> {
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;
63 if rate_limit.allowance == -2f64 {
64 rate_limit.allowance = rate as f64;
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;
73 if rate_limit.allowance < 1.0 {
75 "Rate limited type: {}, IP: {}, time_passed: {}, allowance: {}",
84 "Too many requests. type: {}, IP: {}, {} per {} seconds",
95 rate_limit.allowance -= 1.0;