iuna

iuna - experimental devnet protocol
git clone https://iuna.jhx.app/git/iuna.git
Log | Files | Refs | README | LICENSE

inbound_limiter.rs (4753B)

      1 use std::{
      2     collections::{BTreeMap, VecDeque},
      3     net::IpAddr,
      4     sync::{Arc, Mutex as StdMutex},
      5 };
      6 
      7 use super::{
      8     INBOUND_ACCEPT_RATE_WINDOW_MS, MAX_INBOUND_ACCEPTS_PER_IP_PER_WINDOW, MAX_INBOUND_SESSIONS,
      9     MAX_INBOUND_SESSIONS_PER_IP,
     10 };
     11 
     12 #[derive(Default)]
     13 pub(super) struct InboundConnectionLimiter {
     14     active: usize,
     15     peers: BTreeMap<IpAddr, InboundPeerLimit>,
     16 }
     17 
     18 #[derive(Default)]
     19 struct InboundPeerLimit {
     20     active: usize,
     21     accepted_at_ms: VecDeque<u64>,
     22 }
     23 
     24 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
     25 pub(super) enum InboundSessionRejection {
     26     GlobalActive,
     27     PeerActive,
     28     PeerRate,
     29 }
     30 
     31 impl InboundSessionRejection {
     32     pub(super) fn label(self) -> &'static str {
     33         match self {
     34             Self::GlobalActive => "global active inbound session limit",
     35             Self::PeerActive => "per-IP active inbound session limit",
     36             Self::PeerRate => "per-IP inbound accept rate limit",
     37         }
     38     }
     39 }
     40 
     41 pub(super) struct InboundSessionPermit {
     42     pub(super) limiter: Arc<StdMutex<InboundConnectionLimiter>>,
     43     pub(super) ip: IpAddr,
     44 }
     45 
     46 impl Drop for InboundSessionPermit {
     47     fn drop(&mut self) {
     48         if let Ok(mut limiter) = self.limiter.lock() {
     49             limiter.release(self.ip);
     50         }
     51     }
     52 }
     53 
     54 impl InboundConnectionLimiter {
     55     pub(super) fn try_acquire(
     56         &mut self,
     57         ip: IpAddr,
     58         now_ms: u64,
     59     ) -> std::result::Result<(), InboundSessionRejection> {
     60         self.prune_stale_accepts(now_ms);
     61         if self.active >= MAX_INBOUND_SESSIONS {
     62             return Err(InboundSessionRejection::GlobalActive);
     63         }
     64 
     65         let peer = self.peers.entry(ip).or_default();
     66         prune_peer_accepts(peer, now_ms);
     67         if peer.active >= MAX_INBOUND_SESSIONS_PER_IP {
     68             return Err(InboundSessionRejection::PeerActive);
     69         }
     70         if peer.accepted_at_ms.len() >= MAX_INBOUND_ACCEPTS_PER_IP_PER_WINDOW {
     71             return Err(InboundSessionRejection::PeerRate);
     72         }
     73 
     74         peer.active += 1;
     75         peer.accepted_at_ms.push_back(now_ms);
     76         self.active += 1;
     77         Ok(())
     78     }
     79 
     80     pub(super) fn release(&mut self, ip: IpAddr) {
     81         if self.active > 0 {
     82             self.active -= 1;
     83         }
     84         if let Some(peer) = self.peers.get_mut(&ip) {
     85             if peer.active > 0 {
     86                 peer.active -= 1;
     87             }
     88         }
     89     }
     90 
     91     fn prune_stale_accepts(&mut self, now_ms: u64) {
     92         self.peers.retain(|_, peer| {
     93             prune_peer_accepts(peer, now_ms);
     94             peer.active > 0 || !peer.accepted_at_ms.is_empty()
     95         });
     96     }
     97 }
     98 
     99 fn prune_peer_accepts(peer: &mut InboundPeerLimit, now_ms: u64) {
    100     while peer.accepted_at_ms.front().is_some_and(|accepted_ms| {
    101         now_ms.saturating_sub(*accepted_ms) >= INBOUND_ACCEPT_RATE_WINDOW_MS
    102     }) {
    103         peer.accepted_at_ms.pop_front();
    104     }
    105 }
    106 
    107 #[cfg(test)]
    108 mod tests {
    109     use super::super::{
    110         INBOUND_ACCEPT_RATE_WINDOW_MS, MAX_INBOUND_ACCEPTS_PER_IP_PER_WINDOW, MAX_INBOUND_SESSIONS,
    111         MAX_INBOUND_SESSIONS_PER_IP,
    112     };
    113     use super::{InboundConnectionLimiter, InboundSessionRejection};
    114 
    115     #[test]
    116     fn inbound_limiter_enforces_per_ip_active_limit() {
    117         let ip = "203.0.113.10".parse().unwrap();
    118         let mut limiter = InboundConnectionLimiter::default();
    119         for _ in 0..MAX_INBOUND_SESSIONS_PER_IP {
    120             limiter.try_acquire(ip, 1_000).unwrap();
    121         }
    122 
    123         assert_eq!(
    124             limiter.try_acquire(ip, 1_000).unwrap_err(),
    125             InboundSessionRejection::PeerActive
    126         );
    127 
    128         limiter.release(ip);
    129         limiter.try_acquire(ip, 1_000).unwrap();
    130     }
    131 
    132     #[test]
    133     fn inbound_limiter_enforces_global_active_limit() {
    134         let mut limiter = InboundConnectionLimiter::default();
    135         for index in 0..MAX_INBOUND_SESSIONS {
    136             let ip = format!("198.51.100.{index}").parse().unwrap();
    137             limiter.try_acquire(ip, 1_000).unwrap();
    138         }
    139 
    140         assert_eq!(
    141             limiter
    142                 .try_acquire("203.0.113.200".parse().unwrap(), 1_000)
    143                 .unwrap_err(),
    144             InboundSessionRejection::GlobalActive
    145         );
    146     }
    147 
    148     #[test]
    149     fn inbound_limiter_enforces_per_ip_accept_rate() {
    150         let ip = "203.0.113.20".parse().unwrap();
    151         let mut limiter = InboundConnectionLimiter::default();
    152         for _ in 0..MAX_INBOUND_ACCEPTS_PER_IP_PER_WINDOW {
    153             limiter.try_acquire(ip, 1_000).unwrap();
    154             limiter.release(ip);
    155         }
    156 
    157         assert_eq!(
    158             limiter.try_acquire(ip, 1_000).unwrap_err(),
    159             InboundSessionRejection::PeerRate
    160         );
    161         limiter
    162             .try_acquire(ip, 1_000 + INBOUND_ACCEPT_RATE_WINDOW_MS)
    163             .unwrap();
    164     }
    165 }