crypto/noise: Set timeout limits for the noise handshake (#373)

This PR sets timeout limits for the crypto/noise handshake negotiation.

By default, the timeout is bounded similar to protocol negotiation to
the substream open timeout.

To allow this change, the WebSocket internal API has been opened to
accept the open timeout. Before this change, we had no granular control
and used the connection timeout as a fallback.

cc @paritytech/networking

---------

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>

Author: lexnv

src/crypto/noise/mod.rs

@@ -709,62 +709,71 @@ pub async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
     role: Role,
     max_read_ahead_factor: usize,
     max_write_buffer_size: usize,
+    timeout: std::time::Duration,
 ) -> Result<(NoiseSocket<S>, PeerId), NegotiationError> {
-    tracing::debug!(target: LOG_TARGET, ?role, "start noise handshake");
-
-    let mut noise = NoiseContext::new(keypair, role)?;
-    let payload = match role {
-        Role::Dialer => {
-            // write initial message
-            let first_message = noise.first_message(Role::Dialer)?;
-            let _ = io.write(&first_message).await?;
-            io.flush().await?;
-
-            // read back response which contains the remote peer id
-            let message = noise.read_handshake_message(&mut io).await?;
-            // Decode the remote identity message.
-            let payload = handshake_schema::NoiseHandshakePayload::decode(message)
+    let handle_handshake = async move {
+        tracing::debug!(target: LOG_TARGET, ?role, "start noise handshake");
+
+        let mut noise = NoiseContext::new(keypair, role)?;
+        let payload = match role {
+            Role::Dialer => {
+                // write initial message
+                let first_message = noise.first_message(Role::Dialer)?;
+                let _ = io.write(&first_message).await?;
+                io.flush().await?;
+
+                // read back response which contains the remote peer id
+                let message = noise.read_handshake_message(&mut io).await?;
+                // Decode the remote identity message.
+                let payload = handshake_schema::NoiseHandshakePayload::decode(message)
                 .map_err(ParseError::from)
                 .map_err(|err| {
                     tracing::error!(target: LOG_TARGET, ?err, "failed to decode remote identity message");
                     err
                 })?;
 
-            // send the final message which contains local peer id
-            let second_message = noise.second_message()?;
-            let _ = io.write(&second_message).await?;
-            io.flush().await?;
+                // send the final message which contains local peer id
+                let second_message = noise.second_message()?;
+                let _ = io.write(&second_message).await?;
+                io.flush().await?;
 
-            payload
-        }
-        Role::Listener => {
-            // read remote's first message
-            let _ = noise.read_handshake_message(&mut io).await?;
-
-            // send local peer id.
-            let second_message = noise.second_message()?;
-            let _ = io.write(&second_message).await?;
-            io.flush().await?;
-
-            // read remote's second message which contains their peer id
-            let message = noise.read_handshake_message(&mut io).await?;
-            // Decode the remote identity message.
-            handshake_schema::NoiseHandshakePayload::decode(message).map_err(ParseError::from)?
-        }
-    };
+                payload
+            }
+            Role::Listener => {
+                // read remote's first message
+                let _ = noise.read_handshake_message(&mut io).await?;
+
+                // send local peer id.
+                let second_message = noise.second_message()?;
+                let _ = io.write(&second_message).await?;
+                io.flush().await?;
+
+                // read remote's second message which contains their peer id
+                let message = noise.read_handshake_message(&mut io).await?;
+                // Decode the remote identity message.
+                handshake_schema::NoiseHandshakePayload::decode(message)
+                    .map_err(ParseError::from)?
+            }
+        };
 
-    let dh_remote_pubkey = noise.get_handshake_dh_remote_pubkey()?;
-    let peer = parse_and_verify_peer_id(payload, dh_remote_pubkey)?;
+        let dh_remote_pubkey = noise.get_handshake_dh_remote_pubkey()?;
+        let peer = parse_and_verify_peer_id(payload, dh_remote_pubkey)?;
 
-    Ok((
-        NoiseSocket::new(
-            io,
-            noise.into_transport()?,
-            max_read_ahead_factor,
-            max_write_buffer_size,
-        ),
-        peer,
-    ))
+        Ok((
+            NoiseSocket::new(
+                io,
+                noise.into_transport()?,
+                max_read_ahead_factor,
+                max_write_buffer_size,
+            ),
+            peer,
+        ))
+    };
+
+    match tokio::time::timeout(timeout, handle_handshake).await {
+        Err(_) => Err(NegotiationError::Timeout),
+        Ok(result) => result,
+    }
 }
 
 // TODO: https://github.com/paritytech/litep2p/issues/125 add more tests
@@ -808,14 +817,16 @@ mod tests {
                 &keypair1,
                 Role::Dialer,
                 MAX_READ_AHEAD_FACTOR,
-                MAX_WRITE_BUFFER_SIZE
+                MAX_WRITE_BUFFER_SIZE,
+                std::time::Duration::from_secs(10),
             ),
             handshake(
                 io2,
                 &keypair2,
                 Role::Listener,
                 MAX_READ_AHEAD_FACTOR,
-                MAX_WRITE_BUFFER_SIZE
+                MAX_WRITE_BUFFER_SIZE,
+                std::time::Duration::from_secs(10),
             )
         );
         let (mut res1, mut res2) = (res1.unwrap(), res2.unwrap());

src/transport/tcp/connection.rs

@@ -437,6 +437,7 @@ impl TcpConnection {
             role,
             max_read_ahead_factor,
             max_write_buffer_size,
+            substream_open_timeout,
         )
         .await?;
 
@@ -1146,8 +1147,16 @@ mod tests {
             let keypair = Keypair::generate();
 
             // do a noise handshake
-            let (stream, _peer) =
-                noise::handshake(stream.inner(), &keypair, Role::Dialer, 5, 2).await.unwrap();
+            let (stream, _peer) = noise::handshake(
+                stream.inner(),
+                &keypair,
+                Role::Dialer,
+                5,
+                2,
+                std::time::Duration::from_secs(10),
+            )
+            .await
+            .unwrap();
             let stream: NoiseSocket<Compat<TcpStream>> = stream;
 
             // after the handshake, try to negotiate some random protocol instead of yamux
@@ -1196,8 +1205,16 @@ mod tests {
 
             // do a noise handshake
             let keypair = Keypair::generate();
-            let (stream, _peer) =
-                noise::handshake(stream.inner(), &keypair, Role::Listener, 5, 2).await.unwrap();
+            let (stream, _peer) = noise::handshake(
+                stream.inner(),
+                &keypair,
+                Role::Listener,
+                5,
+                2,
+                std::time::Duration::from_secs(10),
+            )
+            .await
+            .unwrap();
             let stream: NoiseSocket<Compat<TcpStream>> = stream;
 
             // after the handshake, try to negotiate some random protocol instead of yamux
@@ -1262,8 +1279,16 @@ mod tests {
 
             // do a noise handshake
             let keypair = Keypair::generate();
-            let (stream, _peer) =
-                noise::handshake(stream.inner(), &keypair, Role::Dialer, 5, 2).await.unwrap();
+            let (stream, _peer) = noise::handshake(
+                stream.inner(),
+                &keypair,
+                Role::Dialer,
+                5,
+                2,
+                std::time::Duration::from_secs(10),
+            )
+            .await
+            .unwrap();
             let _stream: NoiseSocket<Compat<TcpStream>> = stream;
 
             tokio::time::sleep(std::time::Duration::from_secs(60)).await;
@@ -1307,8 +1332,16 @@ mod tests {
 
             // do a noise handshake
             let keypair = Keypair::generate();
-            let (stream, _peer) =
-                noise::handshake(stream.inner(), &keypair, Role::Listener, 5, 2).await.unwrap();
+            let (stream, _peer) = noise::handshake(
+                stream.inner(),
+                &keypair,
+                Role::Listener,
+                5,
+                2,
+                std::time::Duration::from_secs(10),
+            )
+            .await
+            .unwrap();
             let _stream: NoiseSocket<Compat<TcpStream>> = stream;
 
             tokio::time::sleep(std::time::Duration::from_secs(60)).await;