From 2b056fdd6ec9bece1b7907a83d169b06dfc0ce5b Mon Sep 17 00:00:00 2001 From: Simon Wicky Date: Tue, 20 May 2025 16:08:05 +0200 Subject: [PATCH] change buffer allocation method and use connection timeout --- common/client-libs/mixnet-client/src/client.rs | 1 + common/nymnoise/src/config.rs | 9 ++++++++- common/nymnoise/src/error.rs | 3 +++ common/nymnoise/src/lib.rs | 8 ++++++-- common/nymnoise/src/stream.rs | 12 ++++++------ nym-node/src/node/mod.rs | 2 ++ 6 files changed, 26 insertions(+), 9 deletions(-) diff --git a/common/client-libs/mixnet-client/src/client.rs b/common/client-libs/mixnet-client/src/client.rs index 7a57680f06..f1762440d7 100644 --- a/common/client-libs/mixnet-client/src/client.rs +++ b/common/client-libs/mixnet-client/src/client.rs @@ -332,6 +332,7 @@ mod tests { NoiseConfig::new( Arc::new(x25519::KeyPair::new(&mut rng)), NoiseNetworkView::new_empty(), + Duration::from_millis(1_500), ), Default::default(), ) diff --git a/common/nymnoise/src/config.rs b/common/nymnoise/src/config.rs index 3be3280b8f..5075c1d112 100644 --- a/common/nymnoise/src/config.rs +++ b/common/nymnoise/src/config.rs @@ -5,6 +5,7 @@ use std::{ collections::HashMap, net::{IpAddr, SocketAddr}, sync::Arc, + time::Duration, }; use arc_swap::ArcSwap; @@ -91,16 +92,22 @@ pub struct NoiseConfig { pub(crate) local_key: Arc, pub(crate) pattern: NoisePattern, + pub(crate) timeout: Duration, pub(crate) unsafe_disabled: bool, // allows for nodes to not attempt to do a noise handshake, VERY UNSAFE, FOR DEBUG PURPOSE ONLY } impl NoiseConfig { - pub fn new(noise_key: Arc, network: NoiseNetworkView) -> Self { + pub fn new( + noise_key: Arc, + network: NoiseNetworkView, + timeout: Duration, + ) -> Self { NoiseConfig { network, local_key: noise_key, pattern: Default::default(), + timeout, unsafe_disabled: false, } } diff --git a/common/nymnoise/src/error.rs b/common/nymnoise/src/error.rs index 5f29d96687..5df1bda2f2 100644 --- a/common/nymnoise/src/error.rs +++ b/common/nymnoise/src/error.rs @@ -24,6 +24,9 @@ pub enum NoiseError { #[error("Unknown noise version")] UnknownVersion, + + #[error("Handshake timeout")] + HandshakeTimeout(#[from] tokio::time::error::Elapsed), } impl From for NoiseError { diff --git a/common/nymnoise/src/lib.rs b/common/nymnoise/src/lib.rs index 3276282dfe..f5ce5f4803 100644 --- a/common/nymnoise/src/lib.rs +++ b/common/nymnoise/src/lib.rs @@ -38,7 +38,9 @@ async fn upgrade_noise_initiator_v1( let secret_hash = generate_psk_v1(remote_pub_key); let noise_stream = NoiseStream::new_initiator(conn, config, remote_pub_key, &secret_hash)?; - Ok(Connection::Noise(noise_stream.perform_handshake().await?)) + Ok(Connection::Noise( + tokio::time::timeout(config.timeout, noise_stream.perform_handshake()).await??, + )) } pub async fn upgrade_noise_initiator( @@ -84,7 +86,9 @@ async fn upgrade_noise_responder_v1( let secret_hash = generate_psk_v1(config.local_key.public_key()); let noise_stream = NoiseStream::new_responder(conn, config, &secret_hash)?; - Ok(Connection::Noise(noise_stream.perform_handshake().await?)) + Ok(Connection::Noise( + tokio::time::timeout(config.timeout, noise_stream.perform_handshake()).await??, + )) } pub async fn upgrade_noise_responder( diff --git a/common/nymnoise/src/stream.rs b/common/nymnoise/src/stream.rs index 500bf64ee3..55f4f92fe6 100644 --- a/common/nymnoise/src/stream.rs +++ b/common/nymnoise/src/stream.rs @@ -17,8 +17,8 @@ use tokio::{ }; use tokio_util::codec::{Framed, LengthDelimitedCodec}; -const MAXMSGLEN: usize = 65535; const TAGLEN: usize = 16; +const HANDSHAKE_MAX_LEN: usize = 1024; // using this constant to limit the handshake's buffer size pub(crate) type Psk = [u8; 32]; @@ -66,7 +66,7 @@ impl NoiseStream { .new_framed(inner_stream), handshake: Some(handshake), noise: None, - dec_buffer: BytesMut::with_capacity(MAXMSGLEN), + dec_buffer: BytesMut::new(), } } @@ -92,7 +92,7 @@ impl NoiseStream { &mut self, handshake: &mut HandshakeState, ) -> Result<(), NoiseError> { - let mut buf = BytesMut::zeroed(MAXMSGLEN + TAGLEN); + let mut buf = BytesMut::zeroed(HANDSHAKE_MAX_LEN); // we're in the handshake, we can afford a smaller buffer let len = handshake.write_message(&[], &mut buf)?; buf.truncate(len); self.inner_stream.send(buf.into()).await?; @@ -105,7 +105,7 @@ impl NoiseStream { ) -> Result<(), NoiseError> { match self.inner_stream.next().await { Some(Ok(msg)) => { - let mut buf = vec![0u8; MAXMSGLEN]; + let mut buf = BytesMut::zeroed(HANDSHAKE_MAX_LEN); // we're in the handshake, we can afford a smaller buffer handshake.read_message(&msg, &mut buf)?; Ok(()) } @@ -136,7 +136,7 @@ impl AsyncRead for NoiseStream { Poll::Ready(Some(Ok(noise_msg))) => { // We have a new noise msg - let mut dec_msg = vec![0u8; MAXMSGLEN]; + let mut dec_msg = BytesMut::zeroed(noise_msg.len() - TAGLEN); let len = match projected_self.noise { Some(transport_state) => { match transport_state.read_message(&noise_msg, &mut dec_msg) { @@ -187,7 +187,7 @@ impl AsyncWrite for NoiseStream { ready!(projected_self.inner_stream.as_mut().poll_ready(cx))?; // Ready to send, encrypting message - let mut noise_buf = BytesMut::zeroed(MAXMSGLEN + TAGLEN); + let mut noise_buf = BytesMut::zeroed(buf.len() + TAGLEN); let Ok(len) = (match projected_self.noise { Some(transport_state) => transport_state.write_message(buf, &mut noise_buf), diff --git a/nym-node/src/node/mod.rs b/nym-node/src/node/mod.rs index 5a6b947c68..a566b3f2ff 100644 --- a/nym-node/src/node/mod.rs +++ b/nym-node/src/node/mod.rs @@ -1110,6 +1110,7 @@ impl NymNode { let noise_config = nym_noise::config::NoiseConfig::new( self.x25519_noise_keys.clone(), NoiseNetworkView::new_empty(), + self.config.mixnet.debug.initial_connection_timeout, ) .with_unsafe_disabled(true); @@ -1163,6 +1164,7 @@ impl NymNode { let noise_config = nym_noise::config::NoiseConfig::new( self.x25519_noise_keys.clone(), network_refresher.noise_view(), + self.config.mixnet.debug.initial_connection_timeout, ) .with_unsafe_disabled(self.config.mixnet.debug.unsafe_disable_noise);