bridge/
rate_limiter.rs

1use std::collections::HashMap;
2use std::net::IpAddr;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5use tokio::sync::Mutex;
6
7/// Simple rate limiter to prevent abuse
8pub struct RateLimiter {
9    /// Maximum concurrent connections per IP
10    max_connections_per_ip: usize,
11    /// Maximum connection attempts per minute per IP
12    max_attempts_per_minute: usize,
13    /// Current connection counts per IP
14    connections: Arc<Mutex<HashMap<IpAddr, usize>>>,
15    /// Recent connection attempts per IP (timestamp of each attempt)
16    attempts: Arc<Mutex<HashMap<IpAddr, Vec<Instant>>>>,
17}
18
19impl RateLimiter {
20    pub fn new(max_connections_per_ip: usize, max_attempts_per_minute: usize) -> Self {
21        Self {
22            max_connections_per_ip,
23            max_attempts_per_minute,
24            connections: Arc::new(Mutex::new(HashMap::new())),
25            attempts: Arc::new(Mutex::new(HashMap::new())),
26        }
27    }
28
29    /// Check if a new connection is allowed from this IP
30    /// Returns Ok(()) if allowed, Err with reason if denied
31    pub async fn check_connection(&self, ip: IpAddr) -> Result<(), RateLimitError> {
32        // Check rate limit (attempts per minute)
33        {
34            let mut attempts = self.attempts.lock().await;
35            let now = Instant::now();
36            let minute_ago = now - Duration::from_secs(60);
37            
38            // Get or create attempt list for this IP
39            let ip_attempts = attempts.entry(ip).or_insert_with(Vec::new);
40            
41            // Remove old attempts (older than 1 minute)
42            ip_attempts.retain(|t| *t > minute_ago);
43            
44            // Check if we've exceeded the rate limit
45            if ip_attempts.len() >= self.max_attempts_per_minute {
46                return Err(RateLimitError::TooManyAttempts {
47                    attempts: ip_attempts.len(),
48                    max: self.max_attempts_per_minute,
49                });
50            }
51            
52            // Record this attempt
53            ip_attempts.push(now);
54        }
55
56        // Check concurrent connection limit
57        {
58            let connections = self.connections.lock().await;
59            if let Some(&count) = connections.get(&ip) {
60                if count >= self.max_connections_per_ip {
61                    return Err(RateLimitError::TooManyConnections {
62                        current: count,
63                        max: self.max_connections_per_ip,
64                    });
65                }
66            }
67        }
68
69        Ok(())
70    }
71
72    /// Register a new active connection from this IP
73    pub async fn add_connection(&self, ip: IpAddr) {
74        let mut connections = self.connections.lock().await;
75        *connections.entry(ip).or_insert(0) += 1;
76    }
77
78    /// Remove an active connection from this IP
79    pub async fn remove_connection(&self, ip: IpAddr) {
80        let mut connections = self.connections.lock().await;
81        if let Some(count) = connections.get_mut(&ip) {
82            if *count > 0 {
83                *count -= 1;
84            }
85            if *count == 0 {
86                connections.remove(&ip);
87            }
88        }
89    }
90}
91
92#[derive(Debug)]
93pub enum RateLimitError {
94    TooManyConnections { current: usize, max: usize },
95    TooManyAttempts { attempts: usize, max: usize },
96}
97
98impl std::fmt::Display for RateLimitError {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        match self {
101            RateLimitError::TooManyConnections { current, max } => {
102                write!(f, "Too many concurrent connections ({}/{})", current, max)
103            }
104            RateLimitError::TooManyAttempts { attempts, max } => {
105                write!(f, "Too many connection attempts ({}/{} per minute)", attempts, max)
106            }
107        }
108    }
109}
110
111impl std::error::Error for RateLimitError {}
bridge/rate_limiter.rs

bridge/
rate_limiter.rs
