From 09858062a56d2316ab122ddebabecc039951425e Mon Sep 17 00:00:00 2001
From: Mahmoud Mazouz <hello@fuzzypixelz.com>
Date: Mon, 7 Oct 2024 14:14:52 +0200
Subject: [PATCH] Add timeout for TLS handshakes (#1514)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Add timeout for TLS handshakes

* Remove `dbg!` (quand même)

* Add `tls_handshake_timeout_ms` endpoint config option
---
 io/zenoh-links/zenoh-link-tls/src/lib.rs     |  4 ++++
 io/zenoh-links/zenoh-link-tls/src/unicast.rs | 25 +++++++++++++++++---
 io/zenoh-links/zenoh-link-tls/src/utils.rs   | 22 +++++++++++++----
 3 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/io/zenoh-links/zenoh-link-tls/src/lib.rs b/io/zenoh-links/zenoh-link-tls/src/lib.rs
index 2a18e3b5b3..a547c5d77f 100644
--- a/io/zenoh-links/zenoh-link-tls/src/lib.rs
+++ b/io/zenoh-links/zenoh-link-tls/src/lib.rs
@@ -108,4 +108,8 @@ pub mod config {
 
     pub const TLS_VERIFY_NAME_ON_CONNECT: &str = "verify_name_on_connect";
     pub const TLS_VERIFY_NAME_ON_CONNECT_DEFAULT: bool = true;
+
+    /// The time duration in milliseconds to wait for the TLS handshake to complete.
+    pub const TLS_HANDSHAKE_TIMEOUT_MS: &str = "tls_handshake_timeout_ms";
+    pub const TLS_HANDSHAKE_TIMEOUT_MS_DEFAULT: u64 = 10_000;
 }
diff --git a/io/zenoh-links/zenoh-link-tls/src/unicast.rs b/io/zenoh-links/zenoh-link-tls/src/unicast.rs
index 5a29ab29c0..60eb47b323 100644
--- a/io/zenoh-links/zenoh-link-tls/src/unicast.rs
+++ b/io/zenoh-links/zenoh-link-tls/src/unicast.rs
@@ -369,7 +369,16 @@ impl LinkManagerUnicastTrait for LinkManagerUnicastTls {
             let token = token.clone();
             let manager = self.manager.clone();
 
-            async move { accept_task(socket, acceptor, token, manager).await }
+            async move {
+                accept_task(
+                    socket,
+                    acceptor,
+                    token,
+                    manager,
+                    tls_server_config.tls_handshake_timeout,
+                )
+                .await
+            }
         };
 
         // Update the endpoint locator address
@@ -406,6 +415,7 @@ async fn accept_task(
     acceptor: TlsAcceptor,
     token: CancellationToken,
     manager: NewLinkChannelSender,
+    tls_handshake_timeout: Duration,
 ) -> ZResult<()> {
     async fn accept(socket: &TcpListener) -> ZResult<(TcpStream, SocketAddr)> {
         let res = socket.accept().await.map_err(|e| zerror!(e))?;
@@ -436,9 +446,18 @@ async fn accept_task(
                         };
 
                         // Accept the TLS connection
-                        let tls_stream = match acceptor.accept(tcp_stream).await {
-                            Ok(stream) => TlsStream::Server(stream),
+                        let tls_stream = match tokio::time::timeout(
+                            tls_handshake_timeout,
+                            acceptor.accept(tcp_stream),
+                        )
+                        .await
+                        {
+                            Ok(Ok(stream)) => TlsStream::Server(stream),
                             Err(e) => {
+                                tracing::warn!("TLS handshake timed out: {e}");
+                                continue;
+                            }
+                            Ok(Err(e)) => {
                                 let e = format!("Can not accept TLS connection: {e}");
                                 tracing::warn!("{}", e);
                                 continue;
diff --git a/io/zenoh-links/zenoh-link-tls/src/utils.rs b/io/zenoh-links/zenoh-link-tls/src/utils.rs
index 2894bcf337..b6e2c69578 100644
--- a/io/zenoh-links/zenoh-link-tls/src/utils.rs
+++ b/io/zenoh-links/zenoh-link-tls/src/utils.rs
@@ -14,10 +14,11 @@
 use std::{
     convert::TryFrom,
     fs::File,
-    io,
-    io::{BufReader, Cursor},
+    io::{self, BufReader, Cursor},
     net::SocketAddr,
+    str::FromStr,
     sync::Arc,
+    time::Duration,
 };
 
 use rustls::{
@@ -37,7 +38,7 @@ use zenoh_protocol::core::{
 };
 use zenoh_result::{bail, zerror, ZError, ZResult};
 
-use crate::config::*;
+use crate::config::{self, *};
 
 #[derive(Default, Clone, Copy, Debug)]
 pub struct TlsConfigurator;
@@ -149,6 +150,7 @@ impl ConfigurationInspector<ZenohConfig> for TlsConfigurator {
 
 pub(crate) struct TlsServerConfig {
     pub(crate) server_config: ServerConfig,
+    pub(crate) tls_handshake_timeout: Duration,
 }
 
 impl TlsServerConfig {
@@ -217,7 +219,19 @@ impl TlsServerConfig {
                 .with_single_cert(certs, keys.remove(0))
                 .map_err(|e| zerror!(e))?
         };
-        Ok(TlsServerConfig { server_config: sc })
+
+        let tls_handshake_timeout = Duration::from_millis(
+            config
+                .get(config::TLS_HANDSHAKE_TIMEOUT_MS)
+                .map(u64::from_str)
+                .transpose()?
+                .unwrap_or(config::TLS_HANDSHAKE_TIMEOUT_MS_DEFAULT),
+        );
+
+        Ok(TlsServerConfig {
+            server_config: sc,
+            tls_handshake_timeout,
+        })
     }
 
     async fn load_tls_private_key(config: &Config<'_>) -> ZResult<Vec<u8>> {