Implement rate limiting

Author: jedisct1

example-encrypted-dns.toml

@@ -259,3 +259,22 @@ enabled = false
 
 tokens = ["Y2oHkDJNHz", "G5zY3J5cHQtY", "C5zZWN1cmUuZG5z"]
 
+
+################################
+#         Rate limiting        #
+################################
+
+[rate_limit]
+
+# Enable per-client rate limiting
+
+enabled = false
+
+# Maximum queries per second per client IP
+
+max_queries_per_second = 100
+
+# Maximum number of client IPs to track (uses SIEVE cache for automatic eviction)
+
+capacity = 10000
+

src/config.rs

@@ -15,6 +15,13 @@ pub struct AccessControlConfig {
     pub tokens: Vec<String>,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct RateLimitConfig {
+    pub enabled: bool,
+    pub max_queries_per_second: u32,
+    pub capacity: usize,
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct AnonymizedDNSConfig {
     pub enabled: bool,
@@ -111,6 +118,7 @@ pub struct Config {
     pub metrics: Option<MetricsConfig>,
     pub anonymized_dns: Option<AnonymizedDNSConfig>,
     pub access_control: Option<AccessControlConfig>,
+    pub rate_limit: Option<RateLimitConfig>,
 }
 
 impl Config {

src/globals.rs

@@ -14,6 +14,7 @@ use crate::blacklist::*;
 use crate::cache::*;
 use crate::crypto::*;
 use crate::dnscrypt_certs::*;
+use crate::rate_limiter::*;
 #[cfg(feature = "metrics")]
 use crate::varz::*;
 
@@ -52,6 +53,8 @@ pub struct Globals {
     pub access_control_tokens: Option<Vec<String>>,
     pub client_ttl_holdon: u32,
     pub my_ip: Option<Vec<u8>>,
+    #[educe(Debug(ignore))]
+    pub rate_limiter: SharedRateLimiter,
     #[cfg(feature = "metrics")]
     #[educe(Debug(ignore))]
     pub varz: Varz,

src/main.rs

@@ -31,6 +31,7 @@ mod errors;
 mod globals;
 #[cfg(feature = "metrics")]
 mod metrics;
+mod rate_limiter;
 mod resolver;
 #[cfg(feature = "metrics")]
 mod varz;
@@ -61,6 +62,7 @@ use futures::join;
 use futures::prelude::*;
 use globals::*;
 use parking_lot::Mutex;
+use rate_limiter::*;
 use parking_lot::RwLock;
 #[cfg(target_family = "unix")]
 use privdrop::PrivDrop;
@@ -298,7 +300,7 @@ async fn tcp_acceptor(globals: Arc<Globals>, tcp_listener: TcpListener) -> Resul
     let concurrent_connections = globals.tcp_concurrent_connections.clone();
     let active_connections = globals.tcp_active_connections.clone();
     loop {
-        let (mut client_connection, _client_addr) = match tcp_listener.accept().await {
+        let (mut client_connection, client_addr) = match tcp_listener.accept().await {
             Ok(x) => x,
             Err(e) => {
                 if e.kind() == std::io::ErrorKind::WouldBlock {
@@ -333,6 +335,16 @@ async fn tcp_acceptor(globals: Arc<Globals>, tcp_listener: TcpListener) -> Resul
                 continue;
             }
         };
+
+        if let Some(ref rate_limiter) = globals.rate_limiter {
+            if !rate_limiter.is_allowed(client_addr.ip()) {
+                debug!("Rate limit exceeded for {}", client_addr.ip());
+                #[cfg(feature = "metrics")]
+                globals.varz.client_queries_rate_limited.inc();
+                continue;
+            }
+        }
+
         let (tx, rx) = oneshot::channel::<()>();
         let tx_channel_index = {
             let mut active_connections = active_connections.lock();
@@ -409,11 +421,18 @@ async fn udp_acceptor(
         if packet_len < DNS_HEADER_SIZE {
             continue;
         }
-        // Create a socket clone only when we've checked the packet is valid
-        // This helps avoid resource exhaustion
+
+        if let Some(ref rate_limiter) = globals.rate_limiter {
+            if !rate_limiter.is_allowed(client_addr.ip()) {
+                debug!("Rate limit exceeded for {}", client_addr.ip());
+                #[cfg(feature = "metrics")]
+                globals.varz.client_queries_rate_limited.inc();
+                continue;
+            }
+        }
+
         packet.truncate(packet_len);
 
-        // Only create a new socket if there's capacity for a new connection
         let active_count = concurrent_connections.load(Ordering::Relaxed);
         if active_count >= globals.udp_max_active_connections {
             debug!("UDP connection limit reached, dropping packet");
@@ -835,6 +854,19 @@ fn main() -> Result<(), Error> {
         }
         _ => None,
     };
+    let rate_limiter: SharedRateLimiter = match config.rate_limit {
+        Some(rate_limit) if rate_limit.enabled => {
+            info!(
+                "Rate limiting enabled: {} queries/second per client",
+                rate_limit.max_queries_per_second
+            );
+            Some(Arc::new(RateLimiter::new(
+                rate_limit.capacity,
+                rate_limit.max_queries_per_second,
+            )))
+        }
+        _ => None,
+    };
     let runtime_handle = runtime.handle();
     let globals = Arc::new(Globals {
         runtime_handle: runtime_handle.clone(),
@@ -875,6 +907,7 @@ fn main() -> Result<(), Error> {
         access_control_tokens,
         my_ip: config.my_ip.map(|ip| ip.as_bytes().to_ascii_lowercase()),
         client_ttl_holdon: config.client_ttl_holdon.unwrap_or(60),
+        rate_limiter,
         #[cfg(feature = "metrics")]
         varz: Varz::default(),
     });

src/rate_limiter.rs

@@ -0,0 +1,111 @@
+use std::net::IpAddr;
+use std::sync::Arc;
+
+use coarsetime::Instant;
+use parking_lot::Mutex;
+use sieve_cache::SieveCache;
+
+const DEFAULT_CAPACITY: usize = 10000;
+const DEFAULT_MAX_QPS: u32 = 100;
+const MICROTOKENS_PER_TOKEN: u64 = 1_000_000;
+
+struct ClientState {
+    microtokens: u64,
+    last_update: Instant,
+}
+
+pub struct RateLimiter {
+    clients: Mutex<SieveCache<IpAddr, ClientState>>,
+    max_microtokens: u64,
+    refill_rate: u64, // microtokens per microsecond (equals max_qps)
+}
+
+impl RateLimiter {
+    pub fn new(capacity: usize, max_queries_per_second: u32) -> Self {
+        let capacity = if capacity == 0 {
+            DEFAULT_CAPACITY
+        } else {
+            capacity
+        };
+        let max_qps = if max_queries_per_second == 0 {
+            DEFAULT_MAX_QPS
+        } else {
+            max_queries_per_second
+        };
+        RateLimiter {
+            clients: Mutex::new(
+                SieveCache::new(capacity).expect("Failed to create rate limiter cache"),
+            ),
+            max_microtokens: (max_qps as u64).saturating_mul(MICROTOKENS_PER_TOKEN),
+            refill_rate: max_qps as u64,
+        }
+    }
+
+    pub fn is_allowed(&self, client_ip: IpAddr) -> bool {
+        let now = Instant::now();
+        let mut clients = self.clients.lock();
+
+        if let Some(state) = clients.get_mut(&client_ip) {
+            let elapsed_us = now.as_ticks().saturating_sub(state.last_update.as_ticks());
+            let refill = elapsed_us.saturating_mul(self.refill_rate);
+            state.microtokens = state.microtokens.saturating_add(refill).min(self.max_microtokens);
+            state.last_update = now;
+
+            if state.microtokens >= MICROTOKENS_PER_TOKEN {
+                state.microtokens -= MICROTOKENS_PER_TOKEN;
+                true
+            } else {
+                false
+            }
+        } else {
+            let state = ClientState {
+                microtokens: self.max_microtokens.saturating_sub(MICROTOKENS_PER_TOKEN),
+                last_update: now,
+            };
+            clients.insert(client_ip, state);
+            true
+        }
+    }
+}
+
+pub type SharedRateLimiter = Option<Arc<RateLimiter>>;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::net::{IpAddr, Ipv4Addr};
+
+    #[test]
+    fn test_rate_limiter_allows_initial_requests() {
+        let limiter = RateLimiter::new(100, 10);
+        let ip = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));
+
+        assert!(limiter.is_allowed(ip));
+    }
+
+    #[test]
+    fn test_rate_limiter_exhausts_tokens() {
+        let limiter = RateLimiter::new(100, 3);
+        let ip = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2));
+
+        let mut allowed = 0;
+        for _ in 0..10 {
+            if limiter.is_allowed(ip) {
+                allowed += 1;
+            }
+        }
+        assert!(allowed >= 3 && allowed <= 5);
+    }
+
+    #[test]
+    fn test_rate_limiter_separate_clients() {
+        let limiter = RateLimiter::new(100, 100);
+        let ip1 = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));
+        let ip2 = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2));
+
+        assert!(limiter.is_allowed(ip1));
+        assert!(limiter.is_allowed(ip2));
+        assert!(limiter.is_allowed(ip1));
+        assert!(limiter.is_allowed(ip2));
+    }
+}

src/varz.rs

@@ -20,6 +20,7 @@ pub struct Inner {
     pub client_queries_blocked: IntCounter,
     pub client_queries_resolved: IntCounter,
     pub client_queries_rcode_nxdomain: IntCounter,
+    pub client_queries_rate_limited: IntCounter,
     pub inflight_udp_queries: IntGauge,
     pub inflight_tcp_queries: IntGauge,
     pub upstream_errors: IntCounter,
@@ -117,6 +118,12 @@ impl Inner {
                 labels! {"handler" => "all",}
             ))
             .unwrap(),
+            client_queries_rate_limited: register_int_counter!(opts!(
+                "encrypted_dns_client_queries_rate_limited",
+                "Number of client queries dropped due to rate limiting",
+                labels! {"handler" => "all",}
+            ))
+            .unwrap(),
             inflight_udp_queries: register_int_gauge!(opts!(
                 "encrypted_dns_inflight_udp_queries",
                 "Number of UDP queries currently waiting for a response",