anemo: allow separate max_frame_size for requests and responses

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.
     ///
@@ -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<usize>,
 
+    /// 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<usize>,
+
+    /// 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<usize>,
+
     /// 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<usize> {
-        self.max_frame_size
+    pub(crate) fn max_request_frame_size(&self) -> Option<usize> {
+        self.max_request_frame_size.or(self.max_frame_size)
+    }
+
+    pub(crate) fn max_response_frame_size(&self) -> Option<usize> {
+        self.max_response_frame_size.or(self.max_frame_size)
     }
 
     pub(crate) fn inbound_request_timeout(&self) -> Option<Duration> {

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 => {

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<Item = &'a PeerId>,
+        to_insert: impl Iterator<Item = PeerInfo>,
+    ) -> (Vec<Option<PeerInfo>>, Vec<Option<PeerInfo>>) {
+        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<PeerId, PeerInfo>> {
         self.0.read().unwrap()
     }

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<Bytes>) -> Result<Response<Bytes>> {
         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

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<Request<Bytes>, Response<Bytes>, Infallible>,
-    send_stream: FramedWrite<SendStream, LengthDelimitedCodec>,
-    recv_stream: FramedRead<RecvStream, LengthDelimitedCodec>,
+    send_stream: FramedWrite<SendStream, MessageFrameCodec>,
+    recv_stream: FramedRead<RecvStream, MessageFrameCodec>,
 }
 
 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)),
         }
     }
 

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<Network> {
+    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(())
+}