From 47ce3df3f2167f5543495c97dcfa44628a83e36b Mon Sep 17 00:00:00 2001
From: Eric Rodrigues Pires <eric@eric.dev.br>
Date: Tue, 9 Jun 2026 05:12:37 -0300
Subject: [PATCH] Add `--udp-timeout` CLI flag

---
 CHANGELOG.md                        |   3 +-
 book/src/cli.md                     |   9 +-
 src/config.rs                       |  16 ++-
 src/entrypoint.rs                   |   1 +
 src/udp.rs                          |  76 ++++++----
 tests/integration/main.rs           |   1 +
 tests/integration/udp_rate_limit.rs |   2 +-
 tests/integration/udp_timeout.rs    | 216 ++++++++++++++++++++++++++++
 8 files changed, 294 insertions(+), 30 deletions(-)
 create mode 100644 tests/integration/udp_timeout.rs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3af1131..98f9f3c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,10 +7,11 @@
 - Add experimental UDP support via `udp.sandhole` alias.
 - Add `tcp.sandhole` alias for TCP proxying.
 - Add `--completions` CLI flag.
+- Add `--udp-timeout` CLI flag.
 
 ### Changed
 
-- **BREAKING**: Use gauges for `system_used_memory` and `system_total_memory`.
+- **BREAKING**: Use gauges for `system_used_memory` and `system_total_memory` Prometheus metrics.
 - **BREAKING**: Reserve `.sandhole` aliases for Sandhole.
 - Update dependencies.
 
diff --git a/book/src/cli.md b/book/src/cli.md
index c017589..56758bd 100644
--- a/book/src/cli.md
+++ b/book/src/cli.md
@@ -7,7 +7,8 @@ Sandhole exposes several options, which you can see by running `sandhole --help`
 <pre class="terminal">
 Expose HTTP/SSH/TCP services through SSH port forwarding.
 
-<b><u>Usage:</u></b> <b>sandhole</b> [OPTIONS] &lt;--domain &lt;DOMAIN&gt;|--no-domain&gt;
+<b><u>Usage:</u></b> <b>sandhole</b> [OPTIONS] &lt;--domain &lt;DOMAIN&gt;|--no-domain|--completions &lt;COMPLETIONS&gt;&gt;
+
 
 <b><u>Options:</u></b>
       <b>--domain</b> &lt;DOMAIN&gt;
@@ -420,6 +421,12 @@ Expose HTTP/SSH/TCP services through SSH port forwarding.
 
           By default, these connections are not terminated by Sandhole.
 
+      <b>--udp-timeout</b> &lt;DURATION&gt;
+          How long until SSH channels from UDP sockets are automatically
+          garbage-collected
+
+          [default: 60s]
+
   <b>-h</b>, <b>--help</b>
           Print help (see a summary with &#39;-h&#39;)
 
diff --git a/src/config.rs b/src/config.rs
index 09d4a25..5dc8643 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -535,6 +535,15 @@ pub struct ApplicationConfig {
     /// By default, these connections are not terminated by Sandhole.
     #[arg(long, value_parser = validate_duration, value_name = "DURATION")]
     pub tcp_connection_timeout: Option<Duration>,
+
+    /// How long until SSH channels from UDP sockets are automatically garbage-collected.
+    #[arg(
+        long,
+        default_value = "60s",
+        value_parser = validate_duration,
+        value_name = "DURATION"
+    )]
+    pub udp_timeout: Duration,
 }
 
 fn validate_domain(value: &str) -> color_eyre::Result<String> {
@@ -655,7 +664,8 @@ mod application_config_tests {
                 unproxied_connection_timeout: None,
                 authentication_request_timeout: Duration::from_secs(5),
                 http_request_timeout: None,
-                tcp_connection_timeout: None
+                tcp_connection_timeout: None,
+                udp_timeout: Duration::from_secs(60),
             }
         )
     }
@@ -725,6 +735,7 @@ mod application_config_tests {
             "--authentication-request-timeout=6s",
             "--http-request-timeout=15s",
             "--tcp-connection-timeout=30s",
+            "--udp-timeout=30s",
         ]);
         assert_eq!(
             config,
@@ -796,7 +807,8 @@ mod application_config_tests {
                 unproxied_connection_timeout: Some(Duration::from_secs(4)),
                 authentication_request_timeout: Duration::from_secs(6),
                 http_request_timeout: Some(Duration::from_secs(15)),
-                tcp_connection_timeout: Some(Duration::from_secs(30))
+                tcp_connection_timeout: Some(Duration::from_secs(30)),
+                udp_timeout: Duration::from_secs(30),
             }
         )
     }
diff --git a/src/entrypoint.rs b/src/entrypoint.rs
index 6aa55be..829052c 100644
--- a/src/entrypoint.rs
+++ b/src/entrypoint.rs
@@ -341,6 +341,7 @@ pub async fn entrypoint(config: ApplicationConfig) -> color_eyre::Result<()> {
             .conn_manager(Arc::clone(&udp_connections))
             .ip_filter(Arc::clone(&ip_filter))
             .disable_udp_logs(config.disable_udp_logs)
+            .udp_timeout(config.udp_timeout)
             .build(),
     );
     // Add udp handler service as a listener for udp port updates.
diff --git a/src/udp.rs b/src/udp.rs
index c92a04e..6524b2a 100644
--- a/src/udp.rs
+++ b/src/udp.rs
@@ -3,6 +3,7 @@ use std::{
     mem::size_of,
     net::{IpAddr, SocketAddr},
     sync::Arc,
+    time::Duration,
 };
 
 use crate::{
@@ -19,10 +20,13 @@ use ahash::RandomState;
 use bon::Builder;
 use color_eyre::eyre::Context;
 use dashmap::DashMap;
+use futures_util::pin_mut;
 use metrics::counter;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt, WriteHalf},
+    select,
     sync::Mutex,
+    time::sleep,
 };
 
 pub const MAX_PACKET_SIZE: usize = size_of::<u16>() + u16::MAX as usize;
@@ -38,7 +42,10 @@ fn datagram_buffer() -> Box<[u8; MAX_PACKET_SIZE]> {
 }
 
 // Type for a UDP socket with the write and read halves.
-type UdpSocketHandler = (Arc<Mutex<WriteHalf<SshChannel>>>, DroppableHandle<()>);
+struct UdpSocketHandler {
+    write: Arc<Mutex<WriteHalf<SshChannel>>>,
+    _read_task: DroppableHandle<()>,
+}
 
 // Service that handles creating UDP sockets for reverse forwarding connections.
 #[derive(Builder)]
@@ -57,6 +64,8 @@ pub(crate) struct UdpHandler {
     ip_filter: Arc<IpFilter>,
     // Whether to send UDP logs to the SSH handles behind the forwarded connections.
     disable_udp_logs: bool,
+    // How long until SSH channels from UDP sockets are automatically garbage-collected.
+    udp_timeout: Duration,
 }
 
 pub(crate) trait UdpPortHandler {
@@ -99,8 +108,10 @@ impl UdpPortHandler for Arc<UdpHandler> {
                             .copy_from_slice(&buf[..len]);
 
                         // Check for an existing SSH channel
-                        if let Some(entry) = clone.sockets.get(&(port, address)) {
-                            let channel = Arc::clone(&entry.value().0);
+                        if let Some(mut entry) = clone.sockets.get_mut(&(port, address)) {
+                            let value = entry.value_mut();
+                            let channel = Arc::clone(&value.write);
+                            drop(entry);
                             if channel
                                 .lock()
                                 .await
@@ -145,39 +156,54 @@ impl UdpPortHandler for Arc<UdpHandler> {
                                 continue;
                             }
 
-                            let read_handle = DroppableHandle(tokio::spawn(async move {
+                            let clone_2 = Arc::clone(&clone);
+                            let _read_task = DroppableHandle(tokio::spawn(async move {
                                 let mut write_buf = datagram_buffer();
                                 loop {
-                                    match ssh_read.read_u16().await {
-                                        Ok(len) => {
-                                            if let Err(error) = ssh_read
-                                                .read_exact(&mut write_buf[..len as usize])
-                                                .await
-                                            {
-                                                #[cfg(not(coverage_nightly))]
-                                                tracing::warn!(%port, %error, "Error reading UDP datagram from SSH channel.");
-                                                break;
-                                            } else if let Err(error) = udp_write
-                                                .send_to(&write_buf[..len as usize], address)
-                                                .await
-                                            {
+                                    let sleep_fut = sleep(clone_2.udp_timeout);
+                                    let read_fut = async {
+                                        match ssh_read.read_u16().await {
+                                            Ok(len) => {
+                                                if let Err(error) = ssh_read
+                                                    .read_exact(&mut write_buf[..len as usize])
+                                                    .await
+                                                {
+                                                    #[cfg(not(coverage_nightly))]
+                                                    tracing::warn!(%port, %error, "Error reading UDP datagram from SSH channel.");
+                                                    return false;
+                                                } else if let Err(error) = udp_write
+                                                    .send_to(&write_buf[..len as usize], address)
+                                                    .await
+                                                {
+                                                    #[cfg(not(coverage_nightly))]
+                                                    tracing::warn!(%port, %error, "Error reading from SSH channel for UDP.");
+                                                    return false;
+                                                }
+                                                true
+                                            }
+                                            Err(error) => {
                                                 #[cfg(not(coverage_nightly))]
-                                                tracing::warn!(%port, %error, "Error reading from SSH channel for UDP.");
-                                                break;
+                                                tracing::warn!(%port, %error, "Error reading UDP datagram size from SSH channel.");
+                                                false
                                             }
                                         }
-                                        Err(error) => {
-                                            #[cfg(not(coverage_nightly))]
-                                            tracing::warn!(%port, %error, "Error reading UDP datagram size from SSH channel.");
-                                            break;
-                                        }
+                                    };
+                                    pin_mut!(sleep_fut);
+                                    pin_mut!(read_fut);
+                                    select! {
+                                        success = read_fut => if !success { break; },
+                                        _ = sleep_fut => break,
                                     }
                                 }
+                                clone_2.sockets.remove(&(port, address));
                             }));
 
                             clone.sockets.insert(
                                 (port, address),
-                                (Arc::new(Mutex::new(ssh_write)), read_handle),
+                                UdpSocketHandler {
+                                    write: Arc::new(Mutex::new(ssh_write)),
+                                    _read_task,
+                                },
                             );
                         }
                     }
diff --git a/tests/integration/main.rs b/tests/integration/main.rs
index 29002ff..4d356cf 100644
--- a/tests/integration/main.rs
+++ b/tests/integration/main.rs
@@ -105,5 +105,6 @@ mod udp_no_valid_forwarding;
 mod udp_rate_limit;
 mod udp_reject_low_ports;
 mod udp_reject_port_above_max;
+mod udp_timeout;
 mod websocket_connection;
 mod websocket_timeout;
diff --git a/tests/integration/udp_rate_limit.rs b/tests/integration/udp_rate_limit.rs
index b564679..5c0e813 100644
--- a/tests/integration/udp_rate_limit.rs
+++ b/tests/integration/udp_rate_limit.rs
@@ -132,7 +132,7 @@ async fn udp_rate_limit() {
     );
     let mut chunks = data.chunks_exact(chunk_size);
     assert!(
-        timeout(Duration::from_secs(5), async {
+        timeout(Duration::from_secs(10), async {
             for chunk in &mut chunks {
                 udp_socket.send(chunk).await.unwrap();
                 assert_eq!(udp_socket.recv(&mut buf).await.unwrap(), 2);
diff --git a/tests/integration/udp_timeout.rs b/tests/integration/udp_timeout.rs
new file mode 100644
index 0000000..7e1bdff
--- /dev/null
+++ b/tests/integration/udp_timeout.rs
@@ -0,0 +1,216 @@
+use std::{sync::Arc, time::Duration};
+
+use clap::Parser;
+use russh::keys::{key::PrivateKeyWithHashAlg, load_secret_key};
+use russh::{
+    Channel,
+    client::{Msg, Session},
+};
+use sandhole::{ApplicationConfig, entrypoint};
+use tokio::net::UdpSocket;
+use tokio::{
+    net::TcpStream,
+    time::{sleep, timeout},
+};
+
+use crate::common::SandholeHandle;
+
+/// This test ensures that a UDP socket times out after a certain time
+/// configured by the server.
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn udp_timeout() {
+    // 1. Initialize Sandhole
+    let config = ApplicationConfig::parse_from([
+        "sandhole",
+        "--domain=foobar.tld",
+        "--user-keys-directory",
+        &(format!(
+            "{}/tests/data/user_keys",
+            std::env::var("CARGO_MANIFEST_DIR").unwrap()
+        )),
+        "--admin-keys-directory",
+        &(format!(
+            "{}/tests/data/admin_keys",
+            std::env::var("CARGO_MANIFEST_DIR").unwrap()
+        )),
+        "--certificates-directory",
+        &(format!(
+            "{}/tests/data/certificates",
+            std::env::var("CARGO_MANIFEST_DIR").unwrap()
+        )),
+        "--private-key-file",
+        &(format!(
+            "{}/tests/data/server_keys/ssh",
+            std::env::var("CARGO_MANIFEST_DIR").unwrap()
+        )),
+        "--acme-cache-directory",
+        &(format!(
+            "{}/tests/data/acme_cache",
+            std::env::var("CARGO_MANIFEST_DIR").unwrap()
+        )),
+        "--disable-directory-creation",
+        "--listen-address=127.0.0.1",
+        "--ssh-port=18022",
+        "--http-port=18080",
+        "--https-port=18443",
+        "--acme-use-staging",
+        "--bind-hostnames=none",
+        "--allow-requested-ports",
+        "--idle-connection-timeout=1s",
+        "--authentication-request-timeout=5s",
+        "--http-request-timeout=5s",
+        "--udp-timeout=500ms",
+    ]);
+    let _sandhole_handle = SandholeHandle(tokio::spawn(async move { entrypoint(config).await }));
+    if timeout(Duration::from_secs(5), async {
+        while TcpStream::connect("127.0.0.1:18022").await.is_err() {
+            sleep(Duration::from_millis(100)).await;
+        }
+    })
+    .await
+    .is_err()
+    {
+        panic!("Timeout waiting for Sandhole to start.")
+    };
+
+    // 2. Start SSH client that will be proxied
+    let key = load_secret_key(
+        std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap())
+            .join("tests/data/private_keys/key1"),
+        None,
+    )
+    .expect("Missing file key1");
+    let ssh_client = SshClient;
+    let mut session_one = russh::client::connect(Default::default(), "127.0.0.1:18022", ssh_client)
+        .await
+        .expect("Failed to connect to SSH server");
+    assert!(
+        session_one
+            .authenticate_publickey(
+                "user",
+                PrivateKeyWithHashAlg::new(
+                    Arc::new(key),
+                    session_one
+                        .best_supported_rsa_hash()
+                        .await
+                        .unwrap()
+                        .flatten()
+                )
+            )
+            .await
+            .expect("SSH authentication failed")
+            .success(),
+        "authentication didn't succeed"
+    );
+    session_one
+        .tcpip_forward("udp.sandhole", 12345)
+        .await
+        .expect("tcpip_forward failed");
+
+    // 3. Connect to the UDP port of our proxy
+    let udp_socket = UdpSocket::bind("127.0.0.1:0")
+        .await
+        .expect("UDP connection failed");
+    udp_socket.connect("127.0.0.1:12345").await.unwrap();
+
+    // 4. Send messages without timeout
+    let mut buf = [0u8; 32];
+    udp_socket.send(b"Some message").await.unwrap();
+    if timeout(Duration::from_secs(5), async {
+        assert_eq!(udp_socket.recv(&mut buf).await.unwrap(), 1);
+    })
+    .await
+    .is_err()
+    {
+        panic!("Timeout waiting for UDP socket to reply.")
+    };
+    assert_eq!(&buf[..1], b"1");
+    udp_socket.send(b"Another message").await.unwrap();
+    let mut buf = [0u8; 32];
+    if timeout(Duration::from_secs(5), async {
+        assert_eq!(udp_socket.recv(&mut buf).await.unwrap(), 1);
+    })
+    .await
+    .is_err()
+    {
+        panic!("Timeout waiting for UDP socket to reply.")
+    };
+    assert_eq!(&buf[..1], b"2");
+
+    // 5. Wait for timeout then receive same messages again
+    sleep(Duration::from_secs(1)).await;
+    udp_socket.send(b"Some message again").await.unwrap();
+    if timeout(Duration::from_secs(5), async {
+        assert_eq!(udp_socket.recv(&mut buf).await.unwrap(), 1);
+    })
+    .await
+    .is_err()
+    {
+        panic!("Timeout waiting for UDP socket to reply.")
+    };
+    assert_eq!(&buf[..1], b"1");
+    udp_socket.send(b"Another message again").await.unwrap();
+    let mut buf = [0u8; 32];
+    if timeout(Duration::from_secs(5), async {
+        assert_eq!(udp_socket.recv(&mut buf).await.unwrap(), 1);
+    })
+    .await
+    .is_err()
+    {
+        panic!("Timeout waiting for UDP socket to reply.")
+    };
+    assert_eq!(&buf[..1], b"2");
+
+    // 6. Attempt to close UDP forwarding
+    session_one
+        .cancel_tcpip_forward("udp.sandhole", 12345)
+        .await
+        .expect("cancel_tcpip_forward failed");
+}
+
+struct SshClient;
+
+impl russh::client::Handler for SshClient {
+    type Error = color_eyre::eyre::Error;
+
+    async fn check_server_key(
+        &mut self,
+        _key: &russh::keys::PublicKey,
+    ) -> Result<bool, Self::Error> {
+        Ok(true)
+    }
+
+    async fn server_channel_open_forwarded_tcpip(
+        &mut self,
+        mut channel: Channel<Msg>,
+        _connected_address: &str,
+        _connected_port: u32,
+        _originator_address: &str,
+        _originator_port: u32,
+        _session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        tokio::spawn(async move {
+            let mut counter = 0;
+            loop {
+                while let Some(msg) = &mut channel.wait().await {
+                    match msg {
+                        russh::ChannelMsg::Data { .. } => {
+                            counter += 1;
+                            channel
+                                .data(match counter {
+                                    1 => &b"\x00\x011"[..],
+                                    2 => &b"\x00\x012"[..],
+                                    _ => &b"\x00\x00"[..],
+                                })
+                                .await
+                                .unwrap();
+                        }
+                        russh::ChannelMsg::Close => break,
+                        msg => panic!("Unexpected message {msg:?}"),
+                    }
+                }
+            }
+        });
+        Ok(())
+    }
+}
-- 
2.51.2