From c4e7e4cb4b624d7738c2016d7b7885c6774ff9c2 Mon Sep 17 00:00:00 2001 From: Vlad Semenov Date: Tue, 18 Mar 2025 22:19:44 +0300 Subject: [PATCH 1/4] anemo: add KnownPeers batch_update (#66) --- crates/anemo/src/network/connection_manager.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/crates/anemo/src/network/connection_manager.rs b/crates/anemo/src/network/connection_manager.rs index eaec5aca..ce01fb7d 100644 --- a/crates/anemo/src/network/connection_manager.rs +++ b/crates/anemo/src/network/connection_manager.rs @@ -743,6 +743,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() } From 4b5f0f1d06a31c8ef78ec2e5b446bc633e4e2f77 Mon Sep 17 00:00:00 2001 From: Andrew Schran Date: Fri, 19 Dec 2025 11:52:49 -0500 Subject: [PATCH 2/4] add PartialEq, Eq to some types (#69) --- crates/anemo/src/types/address.rs | 2 +- crates/anemo/src/types/mod.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) 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..ec2a1d63 100644 --- a/crates/anemo/src/types/mod.rs +++ b/crates/anemo/src/types/mod.rs @@ -43,7 +43,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 +54,7 @@ pub enum PeerAffinity { Never, } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct PeerInfo { pub peer_id: PeerId, pub affinity: PeerAffinity, From 9248f8d3951df750360308df22618e5978ad6dc0 Mon Sep 17 00:00:00 2001 From: Andrew Schran Date: Mon, 2 Mar 2026 13:13:00 -0600 Subject: [PATCH 3/4] replace LengthDelimitedCodec with custom MessageFrameCodec (#70) --- crates/anemo/src/config.rs | 10 +- crates/anemo/src/crypto.rs | 5 +- crates/anemo/src/network/request_handler.rs | 7 +- crates/anemo/src/network/wire.rs | 363 +++++++++++++++++++- crates/anemo/src/types/mod.rs | 9 +- crates/anemo/src/types/response.rs | 9 +- 6 files changed, 367 insertions(+), 36 deletions(-) diff --git a/crates/anemo/src/config.rs b/crates/anemo/src/config.rs index 4bfceee5..2ffd4127 100644 --- a/crates/anemo/src/config.rs +++ b/crates/anemo/src/config.rs @@ -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. /// diff --git a/crates/anemo/src/crypto.rs b/crates/anemo/src/crypto.rs index 50d21588..e753e4af 100644 --- a/crates/anemo/src/crypto.rs +++ b/crates/anemo/src/crypto.rs @@ -326,9 +326,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/network/request_handler.rs b/crates/anemo/src/network/request_handler.rs index ecee741c..9d5aaf87 100644 --- a/crates/anemo/src/network/request_handler.rs +++ b/crates/anemo/src/network/request_handler.rs @@ -1,3 +1,4 @@ +use super::wire::MessageFrameCodec; use super::{ wire::{network_message_frame_codec, read_request, write_response}, ActivePeers, @@ -10,7 +11,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 +122,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 { diff --git a/crates/anemo/src/network/wire.rs b/crates/anemo/src/network/wire.rs index 3a126600..03b5f5b0 100644 --- a/crates/anemo/src/network/wire.rs +++ b/crates/anemo/src/network/wire.rs @@ -9,23 +9,124 @@ 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 + +/// 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, + ), + )); + } - if let Some(max_frame_size) = config.max_frame_size() { - builder.max_frame_length(max_frame_size); + dst.reserve(4 + len); + dst.put_u32(len as u32); + dst.extend_from_slice(&data); + Ok(()) } +} - builder.length_field_length(4).big_endian().new_codec() +/// Returns a fully configured message frame codec for writing/reading +/// serialized frames to/from a socket. +pub(crate) fn network_message_frame_codec(config: &Config) -> MessageFrameCodec { + let max_frame_length = config.max_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 +184,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 +206,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 +230,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 +255,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 +326,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/types/mod.rs b/crates/anemo/src/types/mod.rs index ec2a1d63..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 { 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::*; From b4c4cc7cdc636ea877760a04b3f8fd9d2415e066 Mon Sep 17 00:00:00 2001 From: Andrew Schran Date: Tue, 28 Apr 2026 12:07:03 -0400 Subject: [PATCH 4/4] anemo: allow separate max_frame_size for requests and responses Adds `max_request_frame_size` and `max_response_frame_size` Config options. Each falls back to `max_frame_size` when unset, preserving existing behavior. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/anemo/src/config.rs | 24 ++++- crates/anemo/src/network/peer.rs | 8 +- crates/anemo/src/network/request_handler.rs | 8 +- crates/anemo/src/network/tests.rs | 103 +++++++++++++++++++- crates/anemo/src/network/wire.rs | 17 +++- 5 files changed, 146 insertions(+), 14 deletions(-) diff --git a/crates/anemo/src/config.rs b/crates/anemo/src/config.rs index 2ffd4127..0767263e 100644 --- a/crates/anemo/src/config.rs +++ b/crates/anemo/src/config.rs @@ -91,12 +91,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 @@ -271,8 +285,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 max_request_frame_size(&self) -> Option { + self.max_request_frame_size.or(self.max_frame_size) + } + + pub(crate) fn max_response_frame_size(&self) -> Option { + self.max_response_frame_size.or(self.max_frame_size) } pub(crate) fn inbound_request_timeout(&self) -> Option { 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 9d5aaf87..f5aa2cf3 100644 --- a/crates/anemo/src/network/request_handler.rs +++ b/crates/anemo/src/network/request_handler.rs @@ -1,6 +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::{ @@ -137,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..e40ab51e 100644 --- a/crates/anemo/src/network/tests.rs +++ b/crates/anemo/src/network/tests.rs @@ -1,4 +1,4 @@ -use crate::{types::PeerEvent, Network, NetworkRef, Request, Response, Result}; +use crate::{types::PeerEvent, Config, Network, NetworkRef, Request, Response, Result}; use bytes::{Buf, BufMut, Bytes, BytesMut}; use futures::FutureExt; use std::{convert::Infallible, time::Duration}; @@ -868,3 +868,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 03b5f5b0..efb4cdfb 100644 --- a/crates/anemo/src/network/wire.rs +++ b/crates/anemo/src/network/wire.rs @@ -122,10 +122,19 @@ impl Encoder for MessageFrameCodec { } } -/// Returns a fully configured message frame codec for writing/reading -/// serialized frames to/from a socket. -pub(crate) fn network_message_frame_codec(config: &Config) -> MessageFrameCodec { - let max_frame_length = config.max_frame_size().unwrap_or(DEFAULT_MAX_FRAME_LENGTH); +/// Returns a message frame codec sized for request messages. +pub(crate) fn request_message_frame_codec(config: &Config) -> MessageFrameCodec { + let max_frame_length = config + .max_request_frame_size() + .unwrap_or(DEFAULT_MAX_FRAME_LENGTH); + MessageFrameCodec::new(max_frame_length) +} + +/// Returns a message frame codec sized for response messages. +pub(crate) fn response_message_frame_codec(config: &Config) -> MessageFrameCodec { + let max_frame_length = config + .max_response_frame_size() + .unwrap_or(DEFAULT_MAX_FRAME_LENGTH); MessageFrameCodec::new(max_frame_length) }