change buffer allocation method and use connection timeout

Author: simonwicky

common/client-libs/mixnet-client/src/client.rs

@@ -343,6 +343,7 @@ mod tests {
             NoiseConfig::new(
                 Arc::new(x25519::KeyPair::new(&mut rng)),
                 NoiseNetworkView::new_empty(),
+                Duration::from_millis(1_500),
             ),
             Default::default(),
         )

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<x25519::KeyPair>,
     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<x25519::KeyPair>, network: NoiseNetworkView) -> Self {
+    pub fn new(
+        noise_key: Arc<x25519::KeyPair>,
+        network: NoiseNetworkView,
+        timeout: Duration,
+    ) -> Self {
         NoiseConfig {
             network,
             local_key: noise_key,
             pattern: Default::default(),
+            timeout,
             unsafe_disabled: false,
         }
     }

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<Error> for NoiseError {

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(

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),

nym-node/src/node/mod.rs

@@ -1079,6 +1079,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);
 
@@ -1128,6 +1129,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);