diff --git a/crates/anemo/src/config.rs b/crates/anemo/src/config.rs index 4bfceee5..d95c3819 100644 --- a/crates/anemo/src/config.rs +++ b/crates/anemo/src/config.rs @@ -1,5 +1,5 @@ use crate::{ - crypto::{CertVerifier, ExpectedCertVerifier}, + crypto::{CertVerifier, ExpectedCertVerifier, PROBE_SERVER_NAME}, PeerId, Result, }; use pkcs8::EncodePrivateKey; @@ -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)] @@ -63,13 +63,13 @@ pub struct Config { /// /// This limit is applied in the following ways: /// - Inbound connections from [`KnownPeers`] with [`PeerAffinity::High`] or - /// [`PeerAffinity::Allowed`] bypass this limit. All other inbound - /// connections are only accepted if the total number of inbound and outbound - /// connections, irrespective of affinity, is less than this limit. + /// [`PeerAffinity::Allowed`] bypass this limit. All other inbound + /// connections are only accepted if the total number of inbound and outbound + /// connections, irrespective of affinity, is less than this limit. /// - Outbound connections explicitly made by the application via [`Network::connect`] or - /// [`Network::connect_with_peer_id`] bypass this limit. + /// [`Network::connect_with_peer_id`] bypass this limit. /// - Outbound connections made in the background, due to configured [`KnownPeers`], to peers with - /// [`PeerAffinity::High`] bypass this limit and are always attempted. + /// [`PeerAffinity::High`] bypass this limit and are always attempted. /// /// If unspecified, there will be no limit on the number of concurrent connections. /// @@ -81,6 +81,36 @@ pub struct Config { #[serde(skip_serializing_if = "Option::is_none")] pub max_concurrent_connections: Option, + /// 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, + + /// 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, + + /// Token-bucket burst size for `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, + /// Size of the broadcast channel use for subscribing to /// [`PeerEvent`](crate::types::PeerEvent)s via /// [`Network::subscribe`](crate::Network::subscribe). @@ -91,12 +121,26 @@ pub struct Config { /// Set the maximum frame size in bytes. /// - /// This controls the maximum size of a request or response. + /// This controls the maximum size of a request or response. Acts as the default + /// for both directions when [`Config::max_request_frame_size`] or + /// [`Config::max_response_frame_size`] are not set. /// /// If unspecified, there will be no limit. #[serde(skip_serializing_if = "Option::is_none")] pub max_frame_size: Option, + /// Set the maximum frame size in bytes for request messages. + /// + /// If unspecified, falls back to [`Config::max_frame_size`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub max_request_frame_size: Option, + + /// Set the maximum frame size in bytes for response messages. + /// + /// If unspecified, falls back to [`Config::max_frame_size`]. + #[serde(skip_serializing_if = "Option::is_none")] + pub max_response_frame_size: Option, + /// Set a timeout, in milliseconds, for all inbound requests. /// /// When an inbound timeout is hit when processing a request a Response is sent to the @@ -264,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 { + 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; @@ -271,8 +343,12 @@ impl Config { .unwrap_or(PEER_EVENT_BROADCAST_CHANNEL_CAPACITY) } - pub(crate) fn max_frame_size(&self) -> Option { - self.max_frame_size + pub(crate) fn request_frame_size(&self) -> Option { + self.max_request_frame_size.or(self.max_frame_size) + } + + pub(crate) fn response_frame_size(&self) -> Option { + self.max_response_frame_size.or(self.max_frame_size) } pub(crate) fn inbound_request_timeout(&self) -> Option { @@ -429,6 +505,12 @@ impl EndpointConfigBuilder { transport_config.clone(), )?; + // Always register a dedicated probe certificate/server-name in the cert resolver so any node + // can be probed (the client selects it via SNI). It is intentionally not added to the + // client-cert verifier: probing clients still authenticate with their primary certificate. + let (probe_certificate, _) = Self::generate_cert(&keypair, PROBE_SERVER_NAME); + let probe_cert_entry = (PROBE_SERVER_NAME.to_owned(), probe_certificate); + let alternate_server_name = self.alternate_server_name; let server_config = match alternate_server_name { Some(alternate_server_name) => { @@ -441,6 +523,7 @@ impl EndpointConfigBuilder { vec![ (primary_server_name.clone(), primary_certificate.clone()), (alternate_server_name, alternate_certificate), + probe_cert_entry, ], pkcs8_der.clone_key(), cert_verifier, @@ -448,7 +531,10 @@ impl EndpointConfigBuilder { ) } _ => Self::server_config( - vec![(primary_server_name.clone(), primary_certificate.clone())], + vec![ + (primary_server_name.clone(), primary_certificate.clone()), + probe_cert_entry, + ], pkcs8_der.clone_key(), cert_verifier, transport_config.clone(), @@ -585,23 +671,14 @@ impl EndpointConfig { // deprecated. Before that happens, we use the attribute to // keep clippy happy. #[allow(deprecated)] - pub fn client_config_with_expected_server_identity( - &self, - peer_id: PeerId, - ) -> quinn::ClientConfig { - let server_cert_verifier = ExpectedCertVerifier( - CertVerifier { - server_names: vec![self.server_name().into()], - }, - peer_id, - ); + fn client_config_with_verifier(&self, verifier: ExpectedCertVerifier) -> quinn::ClientConfig { let client_crypto = rustls::ClientConfig::builder_with_provider(Arc::new( rustls::crypto::ring::default_provider(), )) .with_safe_default_protocol_versions() .unwrap() .dangerous() - .with_custom_certificate_verifier(Arc::new(server_cert_verifier)) + .with_custom_certificate_verifier(Arc::new(verifier)) .with_client_auth_cert( vec![self.client_certificate.clone()], self.pkcs8_der.clone_key(), @@ -615,6 +692,30 @@ impl EndpointConfig { client } + pub fn client_config_with_expected_server_identity( + &self, + peer_id: PeerId, + ) -> quinn::ClientConfig { + self.client_config_with_verifier(ExpectedCertVerifier( + CertVerifier { + server_names: vec![self.server_name().into()], + }, + peer_id, + )) + } + + /// Client config for a probe connection: verifies the server's identity matches `peer_id` and + /// uses the dedicated probe server-name so the server can recognize and decline the probe (see + /// [`PROBE_SERVER_NAME`]). The probe must also pass `PROBE_SERVER_NAME` as the SNI when dialing. + pub fn client_config_for_probe(&self, peer_id: PeerId) -> quinn::ClientConfig { + self.client_config_with_verifier(ExpectedCertVerifier( + CertVerifier { + server_names: vec![PROBE_SERVER_NAME.into()], + }, + peer_id, + )) + } + #[cfg(test)] pub(crate) fn random(server_name: &str) -> Self { Self::builder() diff --git a/crates/anemo/src/connection.rs b/crates/anemo/src/connection.rs index fe5e9703..1c9781d8 100644 --- a/crates/anemo/src/connection.rs +++ b/crates/anemo/src/connection.rs @@ -58,6 +58,15 @@ impl Connection { self.origin } + /// The TLS server-name (SNI) negotiated for this connection, if any. + pub fn server_name(&self) -> Option { + let handshake_data = self.inner.handshake_data()?; + let handshake_data = handshake_data + .downcast::() + .ok()?; + handshake_data.server_name + } + /// Time the Connection was established #[allow(unused)] pub fn time_established(&self) -> std::time::Instant { diff --git a/crates/anemo/src/crypto.rs b/crates/anemo/src/crypto.rs index 50d21588..137b6019 100644 --- a/crates/anemo/src/crypto.rs +++ b/crates/anemo/src/crypto.rs @@ -20,6 +20,13 @@ static SUPPORTED_ALGORITHMS: WebPkiSupportedAlgorithms = WebPkiSupportedAlgorith mapping: &[(rustls::SignatureScheme::ED25519, SUPPORTED_SIG_ALGS)], }; +/// Dedicated TLS server-name (SNI) used to mark probe connections. +/// +/// A probe verifies that an address is reachable and that the server's cryptographic identity +/// matches an expected [`PeerId`], at the TLS layer only. The dialing client selects this +/// server-name before the handshake. +pub(crate) const PROBE_SERVER_NAME: &str = "anemo-probe"; + #[derive(Clone, Debug)] pub(crate) struct CertVerifier { pub(crate) server_names: Vec, @@ -326,9 +333,12 @@ fn pki_error(error: webpki::Error) -> rustls::Error { BadDer | BadDerTime => { rustls::Error::InvalidCertificate(rustls::CertificateError::BadEncoding) } + #[allow(deprecated)] InvalidSignatureForPublicKey | UnsupportedSignatureAlgorithm - | UnsupportedSignatureAlgorithmForPublicKey => { + | UnsupportedSignatureAlgorithmForPublicKey + | UnsupportedSignatureAlgorithmContext(..) + | UnsupportedSignatureAlgorithmForPublicKeyContext(..) => { rustls::Error::InvalidCertificate(rustls::CertificateError::BadSignature) } e => { diff --git a/crates/anemo/src/endpoint.rs b/crates/anemo/src/endpoint.rs index 191540cc..45d72a38 100644 --- a/crates/anemo/src/endpoint.rs +++ b/crates/anemo/src/endpoint.rs @@ -78,6 +78,22 @@ impl Endpoint { .map(Connecting::new_outbound) } + /// Open a probe connection that verifies `address` is reachable and the server's identity + /// matches `peer_id`, without admitting the result as a peer. + /// + /// Identity is enforced by the TLS verifier, so a successful connect means + /// the identity matched. + pub fn connect_for_probe( + &self, + address: SocketAddr, + peer_id: PeerId, + ) -> Result { + let config = self.config.client_config_for_probe(peer_id); + self.inner + .connect_with(config, address, crate::crypto::PROBE_SERVER_NAME) + .map_err(Into::into) + } + /// Returns the socket address that this Endpoint is bound to. pub fn local_addr(&self) -> SocketAddr { *self.local_addr.read().unwrap() @@ -154,14 +170,68 @@ pin_project_lite::pin_project! { } impl Future for Accept<'_> { - type Output = Option; + type Output = Option; fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { - 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 { + 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 { + self.0 + .accept() + .map_err(anyhow::Error::from) + .map(Connecting::new_inbound) } } @@ -237,7 +307,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 +376,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 +443,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(); diff --git a/crates/anemo/src/lib.rs b/crates/anemo/src/lib.rs index a36093c4..90535fff 100644 --- a/crates/anemo/src/lib.rs +++ b/crates/anemo/src/lib.rs @@ -11,8 +11,8 @@ pub mod types; pub use config::{Config, QuicConfig}; pub use error::{Error, Result}; -pub use network::{Builder, KnownPeers, Network, NetworkRef, Peer}; -pub use routing::Router; +pub use network::{Builder, KnownPeers, Network, NetworkRef, Peer, ProbeOutcome}; +pub use routing::{Router, ServicesOpen, ServicesSealed}; #[doc(inline)] pub use types::{request::Request, response::Response, ConnectionOrigin, Direction, PeerId}; diff --git a/crates/anemo/src/network/connection_manager.rs b/crates/anemo/src/network/connection_manager.rs index eaec5aca..698c374c 100644 --- a/crates/anemo/src/network/connection_manager.rs +++ b/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>>, dial_backoff_states: HashMap, + /// Per-source-IP rate limiter for admitting new inbound connections. + /// `None` when inbound rate limiting is disabled. + inbound_rate_limiter: Option, + active_peers: ActivePeers, known_peers: KnownPeers, @@ -75,6 +81,9 @@ impl ConnectionManager { service: BoxCloneService, Response, Infallible>, ) -> (Self, mpsc::Sender) { 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,13 +152,24 @@ 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() => { - self.handle_connecting_result(connecting_output.unwrap()); + match connecting_output { + Ok(out) => { + self.handle_connecting_result(out); + }, + Err(e) => { + if e.is_cancelled() { + debug!("Pending connection cancelled"); + } else if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } + }, + } }, Some(connection_handler_output) = self.connection_handlers.join_next() => { // If a task panics, just propagate it @@ -229,15 +250,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( @@ -249,6 +303,17 @@ impl ConnectionManager { let fut = async { let connection = connecting.await?; + // A probe connection only needs to verify reachability and our identity at the TLS + // layer, which has now completed. Decline to admit it as a peer. + if connection.server_name().as_deref() == Some(crate::crypto::PROBE_SERVER_NAME) { + trace!( + peer_id = %connection.peer_id(), + "closing completed probe connection" + ); + connection.close(); + return Err(anyhow::anyhow!("completed probe connection")); + } + // TODO close the connection explicitly with a reason once we have machine // readable errors. See https://github.com/MystenLabs/anemo/issues/13 for more info. match known_peers.get(&connection.peer_id()) { @@ -420,6 +485,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))] @@ -743,6 +813,19 @@ impl KnownPeers { self.inner_mut().insert(peer_info.peer_id, peer_info) } + pub fn batch_update<'a>( + &self, + to_remove: impl Iterator, + to_insert: impl Iterator, + ) -> (Vec>, Vec>) { + let mut inner = self.inner_mut(); + let removed = to_remove.map(|peer_id| inner.remove(peer_id)).collect(); + let inserted = to_insert + .map(|peer_info| inner.insert(peer_info.peer_id, peer_info)) + .collect(); + (removed, inserted) + } + fn inner(&self) -> std::sync::RwLockReadGuard<'_, HashMap> { self.0.read().unwrap() } diff --git a/crates/anemo/src/network/inbound_rate_limit.rs b/crates/anemo/src/network/inbound_rate_limit.rs new file mode 100644 index 00000000..cc54ec74 --- /dev/null +++ b/crates/anemo/src/network/inbound_rate_limit.rs @@ -0,0 +1,182 @@ +use std::collections::HashMap; +use std::net::IpAddr; +use std::num::NonZeroU32; +use std::time::Instant; + +/// Hard cap on the number of distinct source IPs tracked at once. +/// At a few dozen bytes per entry this bounds the map near ~2 MiB. +const MAX_TRACKED_IPS: usize = 16_384; + +/// A token-bucket rate limiter keyed by source IP, used to bound the rate at which +/// new inbound connections from any single source are admitted to the TLS handshake. +pub(crate) struct InboundIpRateLimiter { + /// Sustained refill rate, in tokens (connections) per second. + rate_per_sec: f64, + /// Bucket capacity, i.e. the largest instantaneous burst. + burst: f64, + buckets: HashMap, +} + +struct TokenBucket { + tokens: f64, + last_refill: Instant, +} + +impl InboundIpRateLimiter { + pub(crate) fn new(rate_per_sec: NonZeroU32, burst: NonZeroU32) -> Self { + Self { + rate_per_sec: f64::from(rate_per_sec.get()), + burst: f64::from(burst.get()), + buckets: HashMap::new(), + } + } + + /// Returns `true` if a new inbound connection from `ip` is permitted, consuming + /// one token. `now` is taken as an argument to keep the type unit-testable. + pub(crate) fn check(&mut self, ip: IpAddr, now: Instant) -> bool { + // Enforce size cap. Idle buckets are reclaimed separately by `evict_idle`. + if self.buckets.len() >= MAX_TRACKED_IPS && !self.buckets.contains_key(&ip) { + return false; + } + + let rate = self.rate_per_sec; + let burst = self.burst; + let bucket = self.buckets.entry(ip).or_insert(TokenBucket { + tokens: burst, + last_refill: now, + }); + + bucket.tokens = Self::refilled_tokens(bucket, now, rate, burst); + bucket.last_refill = now; + + if bucket.tokens >= 1.0 { + bucket.tokens -= 1.0; + true + } else { + false + } + } + + /// Drop buckets that have refilled to capacity (i.e. have seen no recent + /// activity); a full bucket is indistinguishable from a never-seen IP, so + /// forgetting it is safe. + pub(crate) fn evict_idle(&mut self, now: Instant) { + let rate = self.rate_per_sec; + let burst = self.burst; + self.buckets + .retain(|_ip, bucket| Self::refilled_tokens(bucket, now, rate, burst) < burst); + } + + /// Tokens `bucket` holds after refilling for the time elapsed since its last + /// update, clamped to the burst capacity. + fn refilled_tokens(bucket: &TokenBucket, now: Instant, rate: f64, burst: f64) -> f64 { + let elapsed = now + .saturating_duration_since(bucket.last_refill) + .as_secs_f64(); + (bucket.tokens + elapsed * rate).min(burst) + } + + #[cfg(test)] + fn tracked_ips(&self) -> usize { + self.buckets.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::Ipv4Addr; + use std::time::Duration; + + fn ip(n: u8) -> IpAddr { + IpAddr::V4(Ipv4Addr::new(127, 0, 0, n)) + } + + fn nz(n: u32) -> NonZeroU32 { + NonZeroU32::new(n).unwrap() + } + + #[test] + fn admits_full_burst_then_throttles() { + let now = Instant::now(); + let mut limiter = InboundIpRateLimiter::new(nz(10), nz(10)); + for _ in 0..10 { + assert!(limiter.check(ip(1), now)); + } + // The next attempt at the same instant is over budget. + assert!(!limiter.check(ip(1), now)); + } + + #[test] + fn refills_at_the_configured_rate() { + let now = Instant::now(); + let mut limiter = InboundIpRateLimiter::new(nz(10), nz(10)); + for _ in 0..10 { + assert!(limiter.check(ip(1), now)); + } + assert!(!limiter.check(ip(1), now)); + + // After half a second, ~5 tokens (0.5s * 10/s) have refilled. + let later = now + Duration::from_millis(500); + let admitted = (0..10).filter(|_| limiter.check(ip(1), later)).count(); + assert_eq!(admitted, 5); + } + + #[test] + fn buckets_are_per_ip() { + let now = Instant::now(); + let mut limiter = InboundIpRateLimiter::new(nz(1), nz(1)); + assert!(limiter.check(ip(1), now)); + assert!(!limiter.check(ip(1), now)); + // A different source IP has an independent bucket. + assert!(limiter.check(ip(2), now)); + assert!(!limiter.check(ip(2), now)); + } + + #[test] + fn evict_idle_reclaims_fully_refilled_buckets() { + let now = Instant::now(); + let mut limiter = InboundIpRateLimiter::new(nz(10), nz(10)); + for n in 0u8..50 { + assert!(limiter.check(ip(n), now)); + } + assert_eq!(limiter.tracked_ips(), 50); + + // After enough time for the buckets to refill to capacity, a sweep reclaims + // them all -- they're indistinguishable from never-seen IPs. + limiter.evict_idle(now + Duration::from_secs(60)); + assert_eq!(limiter.tracked_ips(), 0); + } + + #[test] + fn tracked_ips_never_exceeds_the_cap() { + let now = Instant::now(); + // rate = burst = 1 leaves every bucket drained (non-full) at `now`, so eviction + // can reclaim nothing and the cap must be held by rejecting new sources. + let mut limiter = InboundIpRateLimiter::new(nz(1), nz(1)); + for n in 0..MAX_TRACKED_IPS as u32 { + assert!(limiter.check(IpAddr::V4(Ipv4Addr::from(n)), now)); + } + assert_eq!(limiter.tracked_ips(), MAX_TRACKED_IPS); + + // A further distinct source is rejected rather than growing the map past the cap. + let overflow = IpAddr::V4(Ipv4Addr::from(MAX_TRACKED_IPS as u32)); + assert!(!limiter.check(overflow, now)); + assert_eq!(limiter.tracked_ips(), MAX_TRACKED_IPS); + } + + #[test] + fn idle_refill_is_capped_at_burst() { + let now = Instant::now(); + let mut limiter = InboundIpRateLimiter::new(nz(2), nz(3)); + for _ in 0..3 { + assert!(limiter.check(ip(1), now)); + } + assert!(!limiter.check(ip(1), now)); + + // A long idle period refills only up to the burst capacity, not beyond. + let later = now + Duration::from_secs(100); + let admitted = (0..10).filter(|_| limiter.check(ip(1), later)).count(); + assert_eq!(admitted, 3); + } +} diff --git a/crates/anemo/src/network/mod.rs b/crates/anemo/src/network/mod.rs index abe540f9..3bd35126 100644 --- a/crates/anemo/src/network/mod.rs +++ b/crates/anemo/src/network/mod.rs @@ -15,7 +15,7 @@ use tower::{ util::{BoxLayer, BoxService}, Layer, Service, ServiceBuilder, ServiceExt, }; -use tracing::warn; +use tracing::{trace, warn}; mod connection_manager; pub use connection_manager::KnownPeers; @@ -26,6 +26,7 @@ use connection_manager::{ mod peer; pub use peer::Peer; +mod inbound_rate_limit; mod request_handler; mod wire; @@ -268,6 +269,35 @@ impl Builder { } } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ProbeOutcome { + /// The address was reachable and the server's identity matched the expected [`PeerId`]. + Reachable, + /// A connection could not be established (e.g. no route, nothing listening). + Unreachable, + /// A connection was reached at the transport level but TLS verification failed — for a probe + /// with an explicit expected peer id, this means the server's identity did not match. + WrongIdentity, + /// The address could not be resolved or used to initiate a connection. + BadAddress, + /// The probe did not complete within the configured connect timeout. + Timeout, +} + +impl ProbeOutcome { + /// Classify a failed probe connection. A transport/crypto-level failure means we reached an + /// endpoint speaking QUIC+TLS but verification failed; since a probe presents a valid client + /// cert and an explicit expected peer id, that indicates the server's identity did not match. A + /// timeout or any other connection error indicates the endpoint could not be reached. + fn from_connection_error(error: quinn::ConnectionError) -> Self { + match error { + quinn::ConnectionError::TimedOut => ProbeOutcome::Timeout, + quinn::ConnectionError::TransportError(_) => ProbeOutcome::WrongIdentity, + _ => ProbeOutcome::Unreachable, + } + } +} + /// Handle to a network. /// /// This handle can be cheaply cloned and shared across many threads. @@ -321,6 +351,17 @@ impl Network { self.0.connect(addr.into(), Some(peer_id)).await } + /// Probe `addr` to verify it is reachable and that the server's cryptographic identity matches + /// `expected_peer_id`, without joining the peer set or disturbing any existing connection to + /// that peer. + pub async fn probe_address>( + &self, + addr: A, + expected_peer_id: PeerId, + ) -> ProbeOutcome { + self.0.probe_address(addr.into(), expected_peer_id).await + } + pub fn disconnect(&self, peer: PeerId) -> Result<()> { self.0.disconnect(peer) } @@ -410,6 +451,40 @@ impl NetworkInner { receiver.await? } + async fn probe_address(&self, addr: Address, expected_peer_id: PeerId) -> ProbeOutcome { + let socket_addr = match addr.resolve().await { + Ok(socket_addr) => socket_addr, + Err(e) => { + trace!(?addr, "probe: failed to resolve address: {e}"); + return ProbeOutcome::BadAddress; + } + }; + + let connecting = match self + .endpoint + .connect_for_probe(socket_addr, expected_peer_id) + { + Ok(connecting) => connecting, + Err(e) => { + trace!(%socket_addr, "probe: failed to initiate connection: {e}"); + return ProbeOutcome::BadAddress; + } + }; + + match tokio::time::timeout(self.config.connect_timeout(), connecting).await { + // TLS completed, including verifying the server's identity against `expected_peer_id`. + Ok(Ok(connection)) => { + connection.close(0u32.into(), b"probe complete"); + ProbeOutcome::Reachable + } + Ok(Err(e)) => { + trace!(%socket_addr, "probe: connection failed: {e}"); + ProbeOutcome::from_connection_error(e) + } + Err(_) => ProbeOutcome::Timeout, + } + } + fn disconnect(&self, peer_id: PeerId) -> Result<()> { let active_peers = self .active_peers diff --git a/crates/anemo/src/network/peer.rs b/crates/anemo/src/network/peer.rs index 5f0de06b..50516009 100644 --- a/crates/anemo/src/network/peer.rs +++ b/crates/anemo/src/network/peer.rs @@ -1,5 +1,7 @@ use super::{ - wire::{network_message_frame_codec, read_response, write_request}, + wire::{ + read_response, request_message_frame_codec, response_message_frame_codec, write_request, + }, OutboundRequestLayer, }; use crate::{connection::Connection, Config, PeerId, Request, Response, Result}; @@ -50,9 +52,9 @@ impl Peer { async fn do_rpc(&self, request: Request) -> Result> { let (send_stream, recv_stream) = self.connection.open_bi().await?; let mut send_stream = - FramedWrite::new(send_stream, network_message_frame_codec(&self.config)); + FramedWrite::new(send_stream, request_message_frame_codec(&self.config)); let mut recv_stream = - FramedRead::new(recv_stream, network_message_frame_codec(&self.config)); + FramedRead::new(recv_stream, response_message_frame_codec(&self.config)); // // Write Request diff --git a/crates/anemo/src/network/request_handler.rs b/crates/anemo/src/network/request_handler.rs index ecee741c..f5aa2cf3 100644 --- a/crates/anemo/src/network/request_handler.rs +++ b/crates/anemo/src/network/request_handler.rs @@ -1,5 +1,8 @@ +use super::wire::MessageFrameCodec; use super::{ - wire::{network_message_frame_codec, read_request, write_response}, + wire::{ + read_request, request_message_frame_codec, response_message_frame_codec, write_response, + }, ActivePeers, }; use crate::{ @@ -10,7 +13,7 @@ use bytes::Bytes; use quinn::RecvStream; use std::convert::Infallible; use std::sync::Arc; -use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec}; +use tokio_util::codec::{FramedRead, FramedWrite}; use tower::{util::BoxCloneService, ServiceExt}; use tracing::{debug, trace}; @@ -121,8 +124,8 @@ impl InboundRequestHandler { struct BiStreamRequestHandler { connection: Connection, service: BoxCloneService, Response, Infallible>, - send_stream: FramedWrite, - recv_stream: FramedRead, + send_stream: FramedWrite, + recv_stream: FramedRead, } impl BiStreamRequestHandler { @@ -136,8 +139,8 @@ impl BiStreamRequestHandler { Self { connection, service, - send_stream: FramedWrite::new(send_stream, network_message_frame_codec(config)), - recv_stream: FramedRead::new(recv_stream, network_message_frame_codec(config)), + send_stream: FramedWrite::new(send_stream, response_message_frame_codec(config)), + recv_stream: FramedRead::new(recv_stream, request_message_frame_codec(config)), } } diff --git a/crates/anemo/src/network/tests.rs b/crates/anemo/src/network/tests.rs index 67780433..81698964 100644 --- a/crates/anemo/src/network/tests.rs +++ b/crates/anemo/src/network/tests.rs @@ -1,4 +1,6 @@ -use crate::{types::PeerEvent, Network, NetworkRef, Request, Response, Result}; +use crate::{ + types::PeerEvent, Config, Network, NetworkRef, ProbeOutcome, Request, Response, Result, +}; use bytes::{Buf, BufMut, Bytes, BytesMut}; use futures::FutureExt; use std::{convert::Infallible, time::Duration}; @@ -120,6 +122,109 @@ async fn connect_with_invalid_peer_id_ensure_server_doesnt_succeed() -> Result<( Ok(()) } +#[tokio::test] +async fn probe_address_reachable_and_identity_match() -> Result<()> { + let _guard = crate::init_tracing_for_testing(); + + let prober = build_network()?; + let target = build_network()?; + + let outcome = prober + .probe_address(target.local_addr(), target.peer_id()) + .await; + assert_eq!(outcome, ProbeOutcome::Reachable); + + // A probe must never join the peer set on either side. + assert!(prober.peers().is_empty()); + assert!(target.peers().is_empty()); + + Ok(()) +} + +#[tokio::test] +async fn probe_address_identity_mismatch() -> Result<()> { + let _guard = crate::init_tracing_for_testing(); + + let prober = build_network()?; + let target = build_network()?; + let other = build_network()?; + + // Probe the target's address but expect a different peer's identity. + let outcome = prober + .probe_address(target.local_addr(), other.peer_id()) + .await; + assert_eq!(outcome, ProbeOutcome::WrongIdentity); + + assert!(target.peers().is_empty()); + + Ok(()) +} + +#[tokio::test] +async fn probe_address_unreachable() -> Result<()> { + let _guard = crate::init_tracing_for_testing(); + + // Use a short connect timeout so the probe to a dead address returns promptly. + let prober = build_network_with_config(Config { + connect_timeout_ms: Some(300), + ..Default::default() + })?; + + // A bound UDP socket that never speaks QUIC: packets are delivered and dropped, so the + // handshake never completes and the probe times out. + let dead_socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap(); + let dead_addr = dead_socket.local_addr().unwrap(); + + let outcome = prober.probe_address(dead_addr, prober.peer_id()).await; + assert_eq!(outcome, ProbeOutcome::Timeout); + + Ok(()) +} + +#[tokio::test] +async fn probe_does_not_disrupt_existing_connection() -> Result<()> { + let _guard = crate::init_tracing_for_testing(); + + let prober = build_network()?; + let target = build_network()?; + + // Establish a production connection from the prober to the target. + let (mut target_events, _) = target.subscribe().unwrap(); + let peer = prober.connect(target.local_addr()).await?; + assert_eq!(peer, target.peer_id()); + assert_eq!( + target_events.recv().await, + Ok(PeerEvent::NewPeer(prober.peer_id())) + ); + + // Probe the target from the same node that already holds a connection to it. This is the case + // the probe SNI marker must make safe: were the probe admitted as a peer, `add_peer` would + // tie-break against the existing connection and could drop it. + let outcome = prober + .probe_address(target.local_addr(), target.peer_id()) + .await; + assert_eq!(outcome, ProbeOutcome::Reachable); + + // Give the target a chance to (incorrectly) process the probe as a peer. + tokio::task::yield_now().await; + + // The production connection is intact: no LostPeer event, peer still present, RPC still works. + assert!(matches!( + target_events.try_recv(), + Err(tokio::sync::broadcast::error::TryRecvError::Empty) + )); + assert!(target.peers().contains(&prober.peer_id())); + assert!(prober.peers().contains(&target.peer_id())); + + let msg = b"still here"; + let response = prober + .rpc(target.peer_id(), Request::new(msg.as_ref().into())) + .await?; + assert_eq!(response.into_body(), msg.as_ref()); + + Ok(()) +} + #[tokio::test] async fn connect_with_hostname() -> Result<()> { let _guard = crate::init_tracing_for_testing(); @@ -868,3 +973,104 @@ async fn network_ref_via_extension() -> Result<()> { Ok(()) } + +fn build_network_with_config(config: Config) -> Result { + let network = Network::bind("localhost:0") + .random_private_key() + .server_name("test") + .config(config) + .start(echo_service())?; + + trace!( + address =% network.local_addr(), + peer_id =% network.peer_id(), + "starting network" + ); + + Ok(network) +} + +#[tokio::test] +async fn server_max_request_frame_size_rejects_large_requests() -> Result<()> { + let _guard = crate::init_tracing_for_testing(); + + let server = build_network_with_config(Config { + max_request_frame_size: Some(128), + ..Default::default() + })?; + let client = build_network()?; + + let peer = client.connect(server.local_addr()).await?; + + // Request body within the server's request limit succeeds. + let small = vec![0u8; 32]; + let response = client.rpc(peer, Request::new(small.clone().into())).await?; + assert_eq!(response.into_body().as_ref(), small.as_slice()); + + // Request body exceeding the server's request limit is rejected by the server, + // and the client sees the RPC fail. + let big = vec![0u8; 4096]; + client + .rpc(peer, Request::new(big.into())) + .await + .unwrap_err(); + + Ok(()) +} + +#[tokio::test] +async fn client_max_response_frame_size_rejects_large_responses() -> Result<()> { + let _guard = crate::init_tracing_for_testing(); + + let server = build_network()?; + let client = build_network_with_config(Config { + max_response_frame_size: Some(128), + ..Default::default() + })?; + + let peer = client.connect(server.local_addr()).await?; + + // Response body within the client's response limit succeeds. + let small = vec![0u8; 32]; + let response = client.rpc(peer, Request::new(small.clone().into())).await?; + assert_eq!(response.into_body().as_ref(), small.as_slice()); + + // The echoed response exceeds the client's response limit. The request itself + // is still within the (default) request limit, so it leaves the client and the + // server processes it; the failure happens when the client tries to decode the + // response. + let big = vec![0u8; 4096]; + client + .rpc(peer, Request::new(big.into())) + .await + .unwrap_err(); + + Ok(()) +} + +#[tokio::test] +async fn max_frame_size_falls_back_for_both_directions() -> Result<()> { + let _guard = crate::init_tracing_for_testing(); + + // Setting only the legacy `max_frame_size` should constrain both directions, + // matching pre-existing behavior. + let server = build_network_with_config(Config { + max_frame_size: Some(128), + ..Default::default() + })?; + let client = build_network()?; + + let peer = client.connect(server.local_addr()).await?; + + let small = vec![0u8; 32]; + let response = client.rpc(peer, Request::new(small.clone().into())).await?; + assert_eq!(response.into_body().as_ref(), small.as_slice()); + + let big = vec![0u8; 4096]; + client + .rpc(peer, Request::new(big.into())) + .await + .unwrap_err(); + + Ok(()) +} diff --git a/crates/anemo/src/network/wire.rs b/crates/anemo/src/network/wire.rs index 3a126600..8e66c440 100644 --- a/crates/anemo/src/network/wire.rs +++ b/crates/anemo/src/network/wire.rs @@ -9,23 +9,133 @@ use crate::{ Config, Request, Response, Result, }; use anyhow::{anyhow, bail}; -use bytes::{BufMut, Bytes, BytesMut}; +use bytes::{Buf, BufMut, Bytes, BytesMut}; use futures::{SinkExt, StreamExt}; +use std::io; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; -use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec}; +use tokio_util::codec::{Decoder, Encoder, FramedRead, FramedWrite}; const ANEMO: &[u8; 5] = b"anemo"; -/// Returns a fully configured length-delimited codec for writing/reading -/// serialized frames to/from a socket. -pub(crate) fn network_message_frame_codec(config: &Config) -> LengthDelimitedCodec { - let mut builder = LengthDelimitedCodec::builder(); +/// Maximum number of bytes to pre-allocate when decoding a frame. +const MAX_PREALLOCATION: usize = 1 << 20; // 1 MB - if let Some(max_frame_size) = config.max_frame_size() { - builder.max_frame_length(max_frame_size); +/// Default maximum frame length. +const DEFAULT_MAX_FRAME_LENGTH: usize = 8 * 1024 * 1024; // 8 MB + +/// A length-delimited codec that uses the same wire format as tokio-util's +/// `LengthDelimitedCodec` (4-byte big-endian length prefix + data). +pub(crate) struct MessageFrameCodec { + state: DecodeState, + max_frame_length: usize, +} + +#[derive(Debug, Clone, Copy)] +enum DecodeState { + /// Waiting for the 4-byte length prefix. + Head, + /// Accumulating body bytes; stores the total expected frame length. + Data(usize), +} + +impl MessageFrameCodec { + fn new(max_frame_length: usize) -> Self { + Self { + state: DecodeState::Head, + max_frame_length, + } + } +} + +impl Decoder for MessageFrameCodec { + type Item = BytesMut; + type Error = io::Error; + + fn decode(&mut self, src: &mut BytesMut) -> io::Result> { + loop { + match self.state { + DecodeState::Head => { + if src.len() < 4 { + src.reserve(4 - src.len()); + return Ok(None); + } + + let len = u32::from_be_bytes([src[0], src[1], src[2], src[3]]) as usize; + src.advance(4); + + if len > self.max_frame_length { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "frame of length {} exceeds max frame length of {}", + len, self.max_frame_length, + ), + )); + } + + if len == 0 { + return Ok(Some(BytesMut::new())); + } + + // Only pre-allocate up to MAX_PREALLOCATION. + let to_reserve = len.min(MAX_PREALLOCATION); + src.reserve(to_reserve); + + self.state = DecodeState::Data(len); + } + DecodeState::Data(frame_len) => { + if src.len() < frame_len { + // Reserve only up to MAX_PREALLOCATION beyond what we already have. + let remaining = frame_len - src.len(); + let to_reserve = remaining.min(MAX_PREALLOCATION); + src.reserve(to_reserve); + return Ok(None); + } + + self.state = DecodeState::Head; + return Ok(Some(src.split_to(frame_len))); + } + } + } + } +} + +impl Encoder for MessageFrameCodec { + type Error = io::Error; + + fn encode(&mut self, data: Bytes, dst: &mut BytesMut) -> io::Result<()> { + let len = data.len(); + if len > self.max_frame_length { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "frame of length {} exceeds max frame length of {}", + len, self.max_frame_length, + ), + )); + } + + dst.reserve(4 + len); + dst.put_u32(len as u32); + dst.extend_from_slice(&data); + Ok(()) } +} + +/// Returns a message frame codec sized for request messages. +pub(crate) fn request_message_frame_codec(config: &Config) -> MessageFrameCodec { + let max_frame_length = config + .request_frame_size() + .unwrap_or(DEFAULT_MAX_FRAME_LENGTH); + MessageFrameCodec::new(max_frame_length) +} - builder.length_field_length(4).big_endian().new_codec() +/// Returns a message frame codec sized for response messages. +pub(crate) fn response_message_frame_codec(config: &Config) -> MessageFrameCodec { + let max_frame_length = config + .response_frame_size() + .unwrap_or(DEFAULT_MAX_FRAME_LENGTH); + MessageFrameCodec::new(max_frame_length) } /// Anemo requires mTLS in order to ensure that both sides of the connections are authenticated by @@ -83,7 +193,7 @@ pub(crate) async fn write_version_frame( } pub(crate) async fn write_request( - send_stream: &mut FramedWrite, + send_stream: &mut FramedWrite, request: Request, ) -> Result<()> { // Write Version Frame @@ -105,7 +215,7 @@ pub(crate) async fn write_request( } pub(crate) async fn write_response( - send_stream: &mut FramedWrite, + send_stream: &mut FramedWrite, response: Response, ) -> Result<()> { // Write Version Frame @@ -129,7 +239,7 @@ pub(crate) async fn write_response( } pub(crate) async fn read_request( - recv_stream: &mut FramedRead, + recv_stream: &mut FramedRead, ) -> Result> { // Read Version Frame let version = read_version_frame(recv_stream.get_mut()).await?; @@ -154,7 +264,7 @@ pub(crate) async fn read_request( } pub(crate) async fn read_response( - recv_stream: &mut FramedRead, + recv_stream: &mut FramedRead, ) -> Result> { // Read Version Frame let version = read_version_frame(recv_stream.get_mut()).await?; @@ -225,3 +335,239 @@ mod test { assert_eq!(HEADER.as_ref(), buf); } } + +#[cfg(test)] +mod message_frame_codec_tests { + use super::{MessageFrameCodec, DEFAULT_MAX_FRAME_LENGTH, MAX_PREALLOCATION}; + use bytes::{BufMut, Bytes, BytesMut}; + use tokio_util::codec::{Decoder, Encoder, LengthDelimitedCodec}; + + fn new_legacy_codec() -> LengthDelimitedCodec { + LengthDelimitedCodec::builder() + .length_field_length(4) + .big_endian() + .new_codec() + } + + fn legacy_encode(codec: &mut LengthDelimitedCodec, data: &[u8]) -> BytesMut { + let mut buf = BytesMut::new(); + codec + .encode(Bytes::copy_from_slice(data), &mut buf) + .unwrap(); + buf + } + + fn custom_encode(codec: &mut MessageFrameCodec, data: &[u8]) -> BytesMut { + let mut buf = BytesMut::new(); + codec + .encode(Bytes::copy_from_slice(data), &mut buf) + .unwrap(); + buf + } + + #[test] + fn empty_frame_legacy_to_custom() { + let mut enc = new_legacy_codec(); + let wire = legacy_encode(&mut enc, &[]); + + let mut dec = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert!(frame.is_empty()); + } + + #[test] + fn empty_frame_custom_to_legacy() { + let mut enc = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let wire = custom_encode(&mut enc, &[]); + + let mut dec = new_legacy_codec(); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert!(frame.is_empty()); + } + + #[test] + fn small_frame_legacy_to_custom() { + let data: Vec = (0..64).collect(); + let mut enc = new_legacy_codec(); + let wire = legacy_encode(&mut enc, &data); + + let mut dec = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert_eq!(&frame[..], &data); + } + + #[test] + fn small_frame_custom_to_legacy() { + let data: Vec = (0..64).collect(); + let mut enc = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let wire = custom_encode(&mut enc, &data); + + let mut dec = new_legacy_codec(); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert_eq!(&frame[..], &data); + } + + #[test] + fn medium_frame_around_preallocation_boundary() { + let size = MAX_PREALLOCATION + 1; + let data: Vec = (0..size).map(|i| (i % 256) as u8).collect(); + + // Legacy encode -> custom decode + let mut enc = new_legacy_codec(); + let wire = legacy_encode(&mut enc, &data); + let mut dec = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert_eq!(&frame[..], &data); + + // Custom encode -> legacy decode + let mut enc = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let wire = custom_encode(&mut enc, &data); + let mut dec = new_legacy_codec(); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert_eq!(&frame[..], &data); + } + + #[test] + fn large_frame_above_preallocation_limit() { + let size = MAX_PREALLOCATION * 3 + 1; + let data: Vec = (0..size).map(|i| (i % 256) as u8).collect(); + + // Encode with custom codec, decode with custom codec (full buffer available) + let mut enc = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let wire = custom_encode(&mut enc, &data); + let mut dec = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert_eq!(frame.len(), size); + assert_eq!(&frame[..], &data); + + // Cross-codec: legacy encode -> custom decode + let mut enc = new_legacy_codec(); + let wire = legacy_encode(&mut enc, &data); + let mut dec = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert_eq!(&frame[..], &data); + + // Cross-codec: custom encode -> legacy decode + let mut enc = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let wire = custom_encode(&mut enc, &data); + let mut dec = new_legacy_codec(); + let frame = dec.decode(&mut wire.clone()).unwrap().unwrap(); + assert_eq!(&frame[..], &data); + } + + #[test] + fn max_frame_length_rejection_custom() { + let max = 128; + let data = vec![0u8; max + 1]; + + // Encode should fail + let mut enc = MessageFrameCodec::new(max); + let mut buf = BytesMut::new(); + assert!(enc.encode(Bytes::from(data.clone()), &mut buf).is_err()); + + // Decode should fail when length prefix exceeds max + let mut dec = MessageFrameCodec::new(max); + let mut wire = BytesMut::new(); + wire.put_u32((max + 1) as u32); + assert!(dec.decode(&mut wire).is_err()); + } + + #[test] + fn max_frame_length_rejection_matches_legacy() { + let max = 128; + let data = vec![0u8; max + 1]; + + let mut legacy = LengthDelimitedCodec::builder() + .length_field_length(4) + .big_endian() + .max_frame_length(max) + .new_codec(); + let mut buf = BytesMut::new(); + let legacy_result = legacy.encode(Bytes::from(data.clone()), &mut buf); + + let mut custom = MessageFrameCodec::new(max); + let mut buf = BytesMut::new(); + let custom_result = custom.encode(Bytes::from(data), &mut buf); + + // Both should reject oversized frames + assert!(legacy_result.is_err()); + assert!(custom_result.is_err()); + } + + #[test] + fn incremental_delivery_one_byte_at_a_time() { + let data: Vec = (0..256).map(|i| (i % 256) as u8).collect(); + let mut enc = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let wire = custom_encode(&mut enc, &data); + + let mut dec = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let mut src = BytesMut::new(); + + // Feed one byte at a time; should return None until complete + for i in 0..wire.len() - 1 { + src.extend_from_slice(&wire[i..i + 1]); + assert!( + dec.decode(&mut src).unwrap().is_none(), + "should not produce frame at byte {}", + i + ); + } + + // Feed the last byte + src.extend_from_slice(&wire[wire.len() - 1..]); + let frame = dec.decode(&mut src).unwrap().unwrap(); + assert_eq!(&frame[..], &data); + } + + #[test] + fn multiple_frames_in_sequence() { + let data1: Vec = (0..100).collect(); + let data2: Vec = (100..250).collect(); + + let mut enc = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let mut wire = BytesMut::new(); + enc.encode(Bytes::from(data1.clone()), &mut wire).unwrap(); + enc.encode(Bytes::from(data2.clone()), &mut wire).unwrap(); + + let mut dec = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let frame1 = dec.decode(&mut wire).unwrap().unwrap(); + assert_eq!(&frame1[..], &data1); + let frame2 = dec.decode(&mut wire).unwrap().unwrap(); + assert_eq!(&frame2[..], &data2); + } + + #[test] + fn preallocation_cap() { + // Create a frame with a large claimed length + let claimed_len: usize = 64 * 1024 * 1024; // 64 MB + let mut wire = BytesMut::new(); + wire.put_u32(claimed_len as u32); + // Only provide a few bytes of actual data — not the full frame + wire.extend_from_slice(&[0u8; 64]); + + let mut dec = MessageFrameCodec::new(claimed_len); + // Decode should return None (not enough data) + assert!(dec.decode(&mut wire).unwrap().is_none()); + + // The buffer capacity should be bounded to at most MAX_PREALLOCATION plus some overhead. + assert!( + wire.capacity() < MAX_PREALLOCATION * 2, + "buffer capacity {} should be bounded", + wire.capacity(), + ); + } + + #[test] + fn wire_format_compatibility() { + // Verify the wire format is identical between legacy and custom codecs + let data = b"hello world"; + + let mut legacy = new_legacy_codec(); + let legacy_wire = legacy_encode(&mut legacy, data); + + let mut custom = MessageFrameCodec::new(DEFAULT_MAX_FRAME_LENGTH); + let custom_wire = custom_encode(&mut custom, data); + + assert_eq!(legacy_wire, custom_wire); + } +} diff --git a/crates/anemo/src/routing/mod.rs b/crates/anemo/src/routing/mod.rs index 137f6058..3189eeeb 100644 --- a/crates/anemo/src/routing/mod.rs +++ b/crates/anemo/src/routing/mod.rs @@ -4,6 +4,7 @@ use std::{ collections::{BTreeMap, HashMap}, convert::Infallible, fmt, + marker::PhantomData, sync::Arc, }; use tower::{ @@ -32,21 +33,52 @@ impl RouteId { } } +/// Marker type for a [`Router`] that is still accepting service registrations. +/// +/// In this state the router exposes [`Router::route`], [`Router::add_rpc_service`], +/// and [`Router::merge`]. Calling [`Router::route_layer`] consumes the router and +/// returns one in the [`ServicesSealed`] state — no further services can be added +/// after that point. +pub struct ServicesOpen; + +/// Marker type for a [`Router`] that has been sealed by applying a layer. +/// +/// In this state the router exposes only [`Router::route_layer`] (to stack +/// additional layers); attempting to add services is a compile error. +pub struct ServicesSealed; + /// The router type for composing handlers and services. -#[derive(Clone)] -pub struct Router { +/// +/// `Router` carries a typestate parameter `S` that distinguishes a router which +/// is still accepting service registrations ([`ServicesOpen`], the default) +/// from one which has had a layer applied and can therefore no longer accept +/// new services ([`ServicesSealed`]). +pub struct Router { routes: BTreeMap, matcher: RouteMatcher, fallback: Route, + _state: PhantomData S>, +} + +impl Clone for Router { + fn clone(&self) -> Self { + Self { + routes: self.routes.clone(), + matcher: self.matcher.clone(), + fallback: self.fallback.clone(), + _state: PhantomData, + } + } } -impl Router { +impl Router { #[allow(clippy::new_without_default)] pub fn new() -> Self { Self { routes: Default::default(), matcher: Default::default(), fallback: Route::new(not_found::NotFound), + _state: PhantomData, } } @@ -64,7 +96,9 @@ impl Router { panic!("Paths must start with a `/`"); } - if ::downcast_ref::(&service).is_some() { + if ::downcast_ref::>(&service).is_some() + || ::downcast_ref::>(&service).is_some() + { panic!("Invalid route: `Router::route` cannot be used with `Router`s.") } @@ -95,18 +129,20 @@ impl Router { self.route(&path, service) } - /// Merge two routers into one. + /// Merge another router's routes into this one. /// /// This is useful for breaking apps into smaller pieces and combining them - /// into one. - pub fn merge(mut self, other: R) -> Self + /// into one. The other router may itself be in either typestate; any layers + /// it has already applied remain baked into the imported routes. + pub fn merge(mut self, other: R) -> Self where - R: Into, + R: Into>, { let Router { routes, matcher, fallback, + _state: _, } = other.into(); for (id, route) in routes { @@ -124,9 +160,10 @@ impl Router { self } - /// Apply a [`tower::Layer`] to the router that will only run if the request matches - /// a route. - pub fn route_layer(self, layer: L) -> Self + /// Apply a [`tower::Layer`] to all currently registered routes, returning a + /// [`Router`]. To stack additional layers, chain further + /// calls to [`Router::route_layer`] on the returned router. + pub fn route_layer(self, layer: L) -> Router where L: tower::Layer, L::Service: Service, Response = Response, Error = Infallible> @@ -139,25 +176,62 @@ impl Router { routes, matcher, fallback, + _state: _, } = self; - let routes = routes - .into_iter() - .map(|(id, route)| { - let route = Route::new(layer.layer(route)); - (id, route) - }) - .collect(); - Router { + routes: layer_routes(routes, layer), + matcher, + fallback, + _state: PhantomData, + } + } +} + +impl Router { + /// Apply an additional [`tower::Layer`] on top of the layers already applied + /// to this router's routes. + pub fn route_layer(self, layer: L) -> Self + where + L: tower::Layer, + L::Service: Service, Response = Response, Error = Infallible> + + Clone + + Send + + 'static, + >>::Future: Send + 'static, + { + let Router { routes, matcher, fallback, + _state: _, + } = self; + + Router { + routes: layer_routes(routes, layer), + matcher, + fallback, + _state: PhantomData, } } } -impl Service> for Router { +fn layer_routes(routes: BTreeMap, layer: L) -> BTreeMap +where + L: tower::Layer, + L::Service: Service, Response = Response, Error = Infallible> + + Clone + + Send + + 'static, + >>::Future: Send + 'static, +{ + routes + .into_iter() + .map(|(id, route)| (id, Route::new(layer.layer(route)))) + .collect() +} + +impl Service> for Router { type Response = Response; type Error = Infallible; type Future = @@ -415,19 +489,17 @@ mod test { } #[tokio::test] - async fn middleware_applies_to_routes_above() { - let pending = + async fn middleware_applies_to_all_routes_registered_before_seal() { + let pending_one = + tower::service_fn(|_request| async { std::future::pending::>().await }); + let pending_two = tower::service_fn(|_request| async { std::future::pending::>().await }); - let ready = tower::service_fn(|_request| async { - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - Ok(Response::new(Bytes::from_static(b"ready!"))) - }); let router = Router::new() - .route("/one", pending) + .route("/one", pending_one) + .route("/two", pending_two) .route_layer(crate::middleware::timeout::inbound::TimeoutLayer::new( Some(std::time::Duration::new(0, 0)), - )) - .route("/two", ready); + )); let request = Request::new(Bytes::new()).with_route("/one"); let response = router.clone().oneshot(request).await.unwrap(); @@ -435,6 +507,23 @@ mod test { let request = Request::new(Bytes::new()).with_route("/two"); let response = router.clone().oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::RequestTimeout); + } + + #[tokio::test] + async fn stacked_layers_compose_on_sealed_router() { + let ready = tower::service_fn(|_request| async { + Ok(Response::new(Bytes::from_static(b"ready!"))) + }); + // The first call to `route_layer` returns Router; the + // second is the impl on the sealed state, stacking another layer. + let router = Router::new() + .route("/one", ready) + .route_layer(tower::layer::util::Identity::new()) + .route_layer(tower::layer::util::Identity::new()); + + let request = Request::new(Bytes::new()).with_route("/one"); + let response = router.oneshot(request).await.unwrap(); assert_eq!(response.status(), StatusCode::Success); } diff --git a/crates/anemo/src/types/address.rs b/crates/anemo/src/types/address.rs index 16964e12..3c92cc72 100644 --- a/crates/anemo/src/types/address.rs +++ b/crates/anemo/src/types/address.rs @@ -1,5 +1,5 @@ /// Representation of a network address that is dial-able in Anemo -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub enum Address { /// A plain SocketAddr SocketAddr(std::net::SocketAddr), diff --git a/crates/anemo/src/types/mod.rs b/crates/anemo/src/types/mod.rs index ec3c21ab..570f27b9 100644 --- a/crates/anemo/src/types/mod.rs +++ b/crates/anemo/src/types/mod.rs @@ -9,9 +9,10 @@ pub use peer_id::{ConnectionOrigin, Direction, PeerId}; pub use http::Extensions; use quinn::ConnectionError; -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)] #[repr(u16)] pub enum Version { + #[default] V1 = 1, } @@ -28,12 +29,6 @@ impl Version { } } -impl Default for Version { - fn default() -> Self { - Self::V1 - } -} - pub type HeaderMap = std::collections::HashMap; pub mod header { @@ -43,7 +38,7 @@ pub mod header { pub const TIMEOUT: &str = "timeout"; } -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum PeerAffinity { /// Always attempt to maintain a connection with this Peer. High, @@ -54,7 +49,7 @@ pub enum PeerAffinity { Never, } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct PeerInfo { pub peer_id: PeerId, pub affinity: PeerAffinity, diff --git a/crates/anemo/src/types/response.rs b/crates/anemo/src/types/response.rs index 3be27102..7838655c 100644 --- a/crates/anemo/src/types/response.rs +++ b/crates/anemo/src/types/response.rs @@ -48,10 +48,11 @@ impl RawResponseHeader { } } -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)] #[repr(u16)] #[non_exhaustive] pub enum StatusCode { + #[default] Success = 200, BadRequest = 400, NotFound = 404, @@ -105,12 +106,6 @@ impl StatusCode { } } -impl Default for StatusCode { - fn default() -> Self { - Self::Success - } -} - impl std::fmt::Display for StatusCode { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { use StatusCode::*;