1 use enum_map::{enum_map, EnumMap};
2 use once_cell::sync::Lazy;
5 net::{IpAddr, Ipv4Addr, Ipv6Addr},
6 time::{Duration, Instant},
10 const UNINITIALIZED_TOKEN_AMOUNT: f32 = -2.0;
12 static START_TIME: Lazy<Instant> = Lazy::new(Instant::now);
14 /// Smaller than `std::time::Instant` because it uses a smaller integer for seconds and doesn't
16 #[derive(PartialEq, Debug, Clone, Copy)]
17 pub struct InstantSecs {
22 pub fn now() -> Self {
24 secs: u32::try_from(START_TIME.elapsed().as_secs())
25 .expect("server has been running for over 136 years"),
29 fn secs_since(self, earlier: Self) -> u32 {
30 self.secs.saturating_sub(earlier.secs)
33 fn to_instant(self) -> Instant {
34 *START_TIME + Duration::from_secs(self.secs.into())
38 #[derive(PartialEq, Debug, Clone)]
39 struct RateLimitBucket {
40 last_checked: InstantSecs,
41 /// This field stores the amount of tokens that were present at `last_checked`.
42 /// The amount of tokens steadily increases until it reaches the bucket's capacity.
43 /// Performing the rate-limited action consumes 1 token.
47 #[derive(Debug, enum_map::Enum, Copy, Clone, AsRefStr)]
48 pub(crate) enum RateLimitType {
57 type Map<K, C> = HashMap<K, RateLimitedGroup<C>>;
59 #[derive(PartialEq, Debug, Clone)]
60 struct RateLimitedGroup<C> {
61 total: EnumMap<RateLimitType, RateLimitBucket>,
65 impl<C: Default> RateLimitedGroup<C> {
66 fn new(now: InstantSecs) -> Self {
69 _ => RateLimitBucket {
71 tokens: UNINITIALIZED_TOKEN_AMOUNT,
74 children: Default::default(),
85 let capacity = capacity as f32;
86 let secs_to_refill = secs_to_refill as f32;
88 #[allow(clippy::indexing_slicing)] // `EnumMap` has no `get` funciton
89 let bucket = &mut self.total[type_];
91 if bucket.tokens == UNINITIALIZED_TOKEN_AMOUNT {
92 bucket.tokens = capacity;
95 let secs_since_last_checked = now.secs_since(bucket.last_checked) as f32;
96 bucket.last_checked = now;
98 // For `secs_since_last_checked` seconds, increase `bucket.tokens`
99 // by `capacity` every `secs_to_refill` seconds
101 let tokens_per_sec = capacity / secs_to_refill;
102 secs_since_last_checked * tokens_per_sec
105 // Prevent `bucket.tokens` from exceeding `capacity`
106 if bucket.tokens > capacity {
107 bucket.tokens = capacity;
110 if bucket.tokens < 1.0 {
111 // Not enough tokens yet
113 "Rate limited type: {}, time_passed: {}, allowance: {}",
115 secs_since_last_checked,
121 bucket.tokens -= 1.0;
127 /// Rate limiting based on rate type and IP addr
128 #[derive(PartialEq, Debug, Clone, Default)]
129 pub struct RateLimitStorage {
130 /// One bucket per individual IPv4 address
131 ipv4_buckets: Map<Ipv4Addr, ()>,
132 /// Seperate buckets for 48, 56, and 64 bit prefixes of IPv6 addresses
133 ipv6_buckets: Map<[u8; 6], Map<u8, Map<u8, ()>>>,
136 impl RateLimitStorage {
137 /// Rate limiting Algorithm described here: https://stackoverflow.com/a/668327/1655478
139 /// Returns true if the request passed the rate limit, false if it failed and should be rejected.
140 pub(super) fn check_rate_limit_full(
142 type_: RateLimitType,
148 let mut result = true;
151 IpAddr::V4(ipv4) => {
152 // Only used by one address.
156 .or_insert(RateLimitedGroup::new(now));
158 result &= group.check_total(type_, now, capacity, secs_to_refill);
161 IpAddr::V6(ipv6) => {
162 let (key_48, key_56, key_64) = split_ipv6(ipv6);
164 // Contains all addresses with the same first 48 bits. These addresses might be part of the same network.
168 .or_insert(RateLimitedGroup::new(now));
169 result &= group_48.check_total(type_, now, capacity.saturating_mul(16), secs_to_refill);
171 // Contains all addresses with the same first 56 bits. These addresses might be part of the same network.
172 let group_56 = group_48
175 .or_insert(RateLimitedGroup::new(now));
176 result &= group_56.check_total(type_, now, capacity.saturating_mul(4), secs_to_refill);
178 // A group with no children. It is shared by all addresses with the same first 64 bits. These addresses are always part of the same network.
179 let group_64 = group_56
182 .or_insert(RateLimitedGroup::new(now));
184 result &= group_64.check_total(type_, now, capacity, secs_to_refill);
189 debug!("Rate limited IP: {ip}");
195 /// Remove buckets older than the given duration
196 pub(super) fn remove_older_than(&mut self, duration: Duration, now: InstantSecs) {
197 // Only retain buckets that were last used after `instant`
198 let Some(instant) = now.to_instant().checked_sub(duration) else { return };
200 let is_recently_used = |group: &RateLimitedGroup<_>| {
204 .all(|bucket| bucket.last_checked.to_instant() > instant)
207 self.ipv4_buckets.retain(|_, group| is_recently_used(group));
209 self.ipv6_buckets.retain(|_, group_48| {
210 group_48.children.retain(|_, group_56| {
213 .retain(|_, group_64| is_recently_used(group_64));
214 !group_56.children.is_empty()
216 !group_48.children.is_empty()
221 fn split_ipv6(ip: Ipv6Addr) -> ([u8; 6], u8, u8) {
222 let [a0, a1, a2, a3, a4, a5, b, c, ..] = ip.octets();
223 ([a0, a1, a2, a3, a4, a5], b, c)
229 fn test_split_ipv6() {
230 let ip = std::net::Ipv6Addr::new(
231 0x0011, 0x2233, 0x4455, 0x6677, 0x8899, 0xAABB, 0xCCDD, 0xEEFF,
234 super::split_ipv6(ip),
235 ([0x00, 0x11, 0x22, 0x33, 0x44, 0x55], 0x66, 0x77)
240 fn test_rate_limiter() {
241 let mut rate_limiter = super::RateLimitStorage::default();
242 let mut now = super::InstantSecs::now();
252 let ip = ip.parse().unwrap();
254 rate_limiter.check_rate_limit_full(super::RateLimitType::Message, ip, 2, 1, now);
256 rate_limiter.check_rate_limit_full(super::RateLimitType::Post, ip, 3, 1, now);
257 assert!(message_passed);
258 assert!(post_passed);
261 #[allow(clippy::indexing_slicing)]
262 let expected_buckets = |factor: f32, tokens_consumed: f32| {
263 let mut buckets = super::RateLimitedGroup::<()>::new(now).total;
264 buckets[super::RateLimitType::Message] = super::RateLimitBucket {
266 tokens: (2.0 * factor) - tokens_consumed,
268 buckets[super::RateLimitType::Post] = super::RateLimitBucket {
270 tokens: (3.0 * factor) - tokens_consumed,
275 let bottom_group = |tokens_consumed| super::RateLimitedGroup {
276 total: expected_buckets(1.0, tokens_consumed),
282 super::RateLimitStorage {
283 ipv4_buckets: [([123, 123, 123, 123].into(), bottom_group(1.0)),].into(),
286 super::RateLimitedGroup {
287 total: expected_buckets(16.0, 4.0),
291 super::RateLimitedGroup {
292 total: expected_buckets(4.0, 1.0),
293 children: [(0, bottom_group(1.0)),].into(),
298 super::RateLimitedGroup {
299 total: expected_buckets(4.0, 3.0),
300 children: [(0, bottom_group(1.0)), (5, bottom_group(2.0)),].into(),
312 rate_limiter.remove_older_than(std::time::Duration::from_secs(1), now);
313 assert!(rate_limiter.ipv4_buckets.is_empty());
314 assert!(rate_limiter.ipv6_buckets.is_empty());