Add inbound connection admission control to the QUIC endpoint

Author: aschran

crates/anemo/src/config.rs

@@ -8,7 +8,7 @@ use rcgen::{CertificateParams, KeyPair};
 use rustls::pki_types::CertificateDer;
 use rustls::pki_types::PrivateKeyDer;
 use serde::{Deserialize, Serialize};
-use std::{sync::Arc, time::Duration};
+use std::{num::NonZeroU32, sync::Arc, time::Duration};
 
 /// Configuration for a [`Network`](crate::Network).
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
@@ -81,6 +81,36 @@ pub struct Config {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub max_concurrent_connections: Option<usize>,
 
+    /// Whether to require QUIC source-address validation (via a stateless Retry
+    /// packet) before beginning the TLS handshake for an inbound connection whose
+    /// source address has not yet been validated.
+    ///
+    /// When enabled, the first attempt from an unvalidated address is answered with
+    /// a Retry instead of starting the handshake, forcing the peer to prove it can
+    /// receive traffic at its claimed address.
+    ///
+    /// If unspecified, this defaults to `true`.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub require_inbound_address_validation: Option<bool>,
+
+    /// Maximum sustained rate, per source IP, of new inbound connections admitted to
+    /// the TLS handshake, as a token-bucket refill rate in connections per second.
+    /// Inbound attempts from a source IP exceeding this rate are dropped before the
+    /// handshake begins.
+    ///
+    /// If unspecified, this defaults to `10`. Setting it to `0` disables per-IP
+    /// inbound connection rate limiting.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub inbound_connection_rate_limit_per_ip: Option<u32>,
+
+    /// Token-bucket burst size for [`Config::inbound_connection_rate_limit_per_ip`],
+    /// i.e. the maximum number of inbound connections from a single source IP that
+    /// may be admitted in an instantaneous burst.
+    ///
+    /// If unspecified, this defaults to `100`.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub inbound_connection_rate_limit_burst_per_ip: Option<u32>,
+
     /// Size of the broadcast channel use for subscribing to
     /// [`PeerEvent`](crate::types::PeerEvent)s via
     /// [`Network::subscribe`](crate::Network::subscribe).
@@ -278,6 +308,34 @@ impl Config {
         self.max_concurrent_connections
     }
 
+    pub(crate) fn require_inbound_address_validation(&self) -> bool {
+        const DEFAULT: bool = true;
+
+        self.require_inbound_address_validation.unwrap_or(DEFAULT)
+    }
+
+    /// Sustained per-source-IP inbound connection admission rate in connections per
+    /// second. Returns `None` only when explicitly disabled by configuring `0`.
+    pub(crate) fn inbound_connection_rate_limit_per_ip(&self) -> Option<NonZeroU32> {
+        const DEFAULT_RATE_PER_SEC: u32 = 10;
+
+        // `NonZeroU32::new` yields `None` for a configured `0`, i.e. rate limiting off.
+        NonZeroU32::new(
+            self.inbound_connection_rate_limit_per_ip
+                .unwrap_or(DEFAULT_RATE_PER_SEC),
+        )
+    }
+
+    pub(crate) fn inbound_connection_rate_limit_burst_per_ip(&self) -> NonZeroU32 {
+        const DEFAULT_BURST: u32 = 100;
+
+        // A configured burst of `0` is nonsensical (it would admit nothing); fall back
+        // to the default in that case as well as when unset.
+        self.inbound_connection_rate_limit_burst_per_ip
+            .and_then(NonZeroU32::new)
+            .unwrap_or_else(|| NonZeroU32::new(DEFAULT_BURST).expect("DEFAULT_BURST is nonzero"))
+    }
+
     pub(crate) fn peer_event_broadcast_channel_capacity(&self) -> usize {
         const PEER_EVENT_BROADCAST_CHANNEL_CAPACITY: usize = 128;
 

crates/anemo/src/endpoint.rs

@@ -154,14 +154,68 @@ pin_project_lite::pin_project! {
 }
 
 impl Future for Accept<'_> {
-    type Output = Option<Connecting>;
+    type Output = Option<Incoming>;
 
     fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
-        self.project().inner.poll(ctx).map(|maybe_connecting| {
-            maybe_connecting
-                .and_then(|incoming| incoming.accept().ok())
-                .map(Connecting::new_inbound)
-        })
+        self.project()
+            .inner
+            .poll(ctx)
+            .map(|maybe_incoming| maybe_incoming.map(Incoming::new))
+    }
+}
+
+/// An inbound connection attempt for which the server has not yet begun the TLS
+/// handshake.
+///
+/// Yielded by [`Endpoint::accept`]. Admission decisions (source-address validation,
+/// rate limiting) are made against this type before calling [`Incoming::accept`],
+/// which is what begins the handshake.
+#[derive(Debug)]
+#[must_use = "an Incoming must be accepted, retried, or ignored"]
+pub(crate) struct Incoming(quinn::Incoming);
+
+impl Incoming {
+    pub(crate) fn new(inner: quinn::Incoming) -> Self {
+        Self(inner)
+    }
+
+    /// The remote peer's UDP socket address.
+    pub(crate) fn remote_address(&self) -> SocketAddr {
+        self.0.remote_address()
+    }
+
+    /// Whether the peer has proved it can receive traffic at `remote_address()`.
+    pub(crate) fn remote_address_validated(&self) -> bool {
+        self.0.remote_address_validated()
+    }
+
+    /// Ensure the peer's source address is validated before the handshake.
+    ///
+    /// If the address is already validated, returns the attempt so the caller can
+    /// proceed. Otherwise sends a stateless Retry packet and consumes the attempt
+    /// (returning `None`); the peer must re-initiate from a validated address.
+    #[must_use = "a returned Incoming must be accepted or ignored"]
+    pub(crate) fn validate_source(self) -> Option<Incoming> {
+        if self.remote_address_validated() {
+            return Some(self);
+        }
+        // The address is unvalidated, so a Retry is always legal here: quinn
+        // guarantees may_retry() whenever the address is unvalidated.
+        let _ = self.0.retry();
+        None
+    }
+
+    /// Drop the attempt without sending any response packet.
+    pub(crate) fn ignore(self) {
+        self.0.ignore()
+    }
+
+    /// Begin the handshake for this connection.
+    pub(crate) fn accept(self) -> Result<Connecting> {
+        self.0
+            .accept()
+            .map_err(anyhow::Error::from)
+            .map(Connecting::new_inbound)
     }
 }
 
@@ -237,7 +291,14 @@ mod test {
         };
 
         let peer_2 = async move {
-            let connection = endpoint_2.accept().await.unwrap().await.unwrap();
+            let connection = endpoint_2
+                .accept()
+                .await
+                .unwrap()
+                .accept()
+                .unwrap()
+                .await
+                .unwrap();
             assert_eq!(connection.peer_id(), peer_id_1);
 
             let mut recv = connection.accept_uni().await.unwrap();
@@ -299,10 +360,24 @@ mod test {
         };
 
         let peer_2 = async move {
-            let connection_1 = endpoint_2.accept().await.unwrap().await.unwrap();
+            let connection_1 = endpoint_2
+                .accept()
+                .await
+                .unwrap()
+                .accept()
+                .unwrap()
+                .await
+                .unwrap();
             assert_eq!(connection_1.peer_id(), peer_id_1);
 
-            let connection_2 = endpoint_2.accept().await.unwrap().await.unwrap();
+            let connection_2 = endpoint_2
+                .accept()
+                .await
+                .unwrap()
+                .accept()
+                .unwrap()
+                .await
+                .unwrap();
             assert_eq!(connection_2.peer_id(), peer_id_1);
             assert_ne!(connection_1.stable_id(), connection_2.stable_id());
 
@@ -352,7 +427,16 @@ mod test {
 
         let (connection_1_to_2, connection_2_to_1) = timeout(join(
             async { endpoint_1.connect(addr_2).unwrap().await.unwrap() },
-            async { endpoint_2.accept().await.unwrap().await.unwrap() },
+            async {
+                endpoint_2
+                    .accept()
+                    .await
+                    .unwrap()
+                    .accept()
+                    .unwrap()
+                    .await
+                    .unwrap()
+            },
         ))
         .await
         .unwrap();

crates/anemo/src/network/connection_manager.rs

@@ -1,8 +1,9 @@
+use super::inbound_rate_limit::InboundIpRateLimiter;
 use super::request_handler::InboundRequestHandler;
 use crate::{
     config::Config,
     connection::Connection,
-    endpoint::{Connecting, Endpoint},
+    endpoint::{Connecting, Endpoint, Incoming},
     types::{Address, DisconnectReason, PeerAffinity, PeerEvent, PeerInfo},
     ConnectionOrigin, PeerId, Request, Response, Result,
 };
@@ -11,6 +12,7 @@ use std::{
     collections::{hash_map::Entry, HashMap},
     convert::Infallible,
     sync::{Arc, RwLock},
+    time::Instant,
 };
 use tokio::{
     sync::{broadcast, mpsc, oneshot},
@@ -54,6 +56,10 @@ pub(crate) struct ConnectionManager {
     pending_dials: HashMap<PeerId, oneshot::Receiver<Result<PeerId>>>,
     dial_backoff_states: HashMap<PeerId, DialBackoffState>,
 
+    /// Per-source-IP rate limiter for admitting new inbound connections.
+    /// `None` when inbound rate limiting is disabled.
+    inbound_rate_limiter: Option<InboundIpRateLimiter>,
+
     active_peers: ActivePeers,
     known_peers: KnownPeers,
 
@@ -75,6 +81,9 @@ impl ConnectionManager {
         service: BoxCloneService<Request<Bytes>, Response<Bytes>, Infallible>,
     ) -> (Self, mpsc::Sender<ConnectionManagerRequest>) {
         let (sender, receiver) = mpsc::channel(config.connection_manager_channel_capacity());
+        let inbound_rate_limiter = config.inbound_connection_rate_limit_per_ip().map(|rate| {
+            InboundIpRateLimiter::new(rate, config.inbound_connection_rate_limit_burst_per_ip())
+        });
         (
             Self {
                 config,
@@ -84,6 +93,7 @@ impl ConnectionManager {
                 connection_handlers: JoinSet::new(),
                 pending_dials: HashMap::default(),
                 dial_backoff_states: HashMap::default(),
+                inbound_rate_limiter,
                 active_peers,
                 known_peers,
                 service,
@@ -142,9 +152,9 @@ impl ConnectionManager {
                         }
                     }
                 }
-                connecting = self.endpoint.accept() => {
-                    if let Some(connecting) = connecting {
-                        self.handle_incoming(connecting);
+                incoming = self.endpoint.accept() => {
+                    if let Some(incoming) = incoming {
+                        self.handle_incoming(incoming);
                     }
                 },
                 Some(connecting_output) = self.pending_connections.join_next() => {
@@ -229,15 +239,48 @@ impl ConnectionManager {
         self.dial_peer(address, peer_id, oneshot);
     }
 
-    fn handle_incoming(&mut self, connecting: Connecting) {
+    fn handle_incoming(&mut self, incoming: Incoming) {
         trace!("received new incoming connection");
 
-        self.pending_connections.spawn(Self::handle_incoming_task(
-            connecting,
-            self.config.clone(),
-            self.active_peers.clone(),
-            self.known_peers.clone(),
-        ));
+        let remote_address = incoming.remote_address();
+
+        // Validate source address, if enabled.
+        let incoming = if self.config.require_inbound_address_validation() {
+            match incoming.validate_source() {
+                Some(incoming) => incoming,
+                None => return,
+            }
+        } else {
+            incoming
+        };
+
+        // Apply per-source-IP inbound rate limit. Loopback sources are
+        // exempt: they can only originate on this host (a remote attacker cannot
+        // present a loopback source address), and exempting them avoids throttling
+        // local/test topologies where many peers share a loopback address.
+        if let Some(limiter) = self.inbound_rate_limiter.as_mut() {
+            let ip = remote_address.ip();
+            if !ip.is_loopback() && !limiter.check(ip, Instant::now()) {
+                debug!(%remote_address, "dropping inbound connection: per-source-IP rate limit exceeded");
+                incoming.ignore();
+                return;
+            }
+        }
+
+        // Admit the connection and begin the handshake.
+        match incoming.accept() {
+            Ok(connecting) => {
+                self.pending_connections.spawn(Self::handle_incoming_task(
+                    connecting,
+                    self.config.clone(),
+                    self.active_peers.clone(),
+                    self.known_peers.clone(),
+                ));
+            }
+            Err(e) => {
+                debug!(%remote_address, "failed to accept inbound connection: {e}");
+            }
+        }
     }
 
     async fn handle_incoming_task(
@@ -420,6 +463,11 @@ impl ConnectionManager {
             self.dial_peer(address, Some(peer.peer_id), sender);
             self.pending_dials.insert(peer.peer_id, receiver);
         }
+
+        // Clean up rate limiter buckets.
+        if let Some(limiter) = self.inbound_rate_limiter.as_mut() {
+            limiter.evict_idle(now);
+        }
     }
 
     #[instrument(level = "trace", skip_all, fields(peer_id = ?peer_id, address = ?address))]