From 38f1752fda46a9a45f7e507d52923dc0d95f157d Mon Sep 17 00:00:00 2001
From: ospab <ospab@ospab.host>
Date: Sat, 30 May 2026 02:12:15 +0300
Subject: [PATCH] fix(client): stabilize UDP sessions - prevent crashes on
 transient recv errors in udp_nat and proxy

---
 ostp-client/src/tunnel/proxy.rs   | 14 +++++--
 ostp-client/src/tunnel/udp_nat.rs | 68 ++++++++++++++++---------------
 2 files changed, 47 insertions(+), 35 deletions(-)

diff --git a/ostp-client/src/tunnel/proxy.rs b/ostp-client/src/tunnel/proxy.rs
index 25df852..3041986 100644
--- a/ostp-client/src/tunnel/proxy.rs
+++ b/ostp-client/src/tunnel/proxy.rs
@@ -371,11 +371,19 @@ async fn handle_udp_associate(
     loop {
         tokio::select! {
             res = client_tcp.read(&mut tcp_buf) => {
-                let n = res?;
-                if n == 0 { break; }
+                match res {
+                    Ok(0) | Err(_) => break,
+                    Ok(_) => {}
+                }
             }
             res = sock_rx.recv_from(&mut buf) => {
-                let (len, addr) = res?;
+                let (len, addr) = match res {
+                    Ok(v) => v,
+                    Err(e) => {
+                        tracing::debug!("udp_associate recv_from error: {}", e);
+                        continue; // transient error, don't kill the session
+                    }
+                };
                 {
                     let mut guard = client_udp_addr.lock().unwrap();
                     if guard.is_none() {
diff --git a/ostp-client/src/tunnel/udp_nat.rs b/ostp-client/src/tunnel/udp_nat.rs
index a385119..a615d8c 100644
--- a/ostp-client/src/tunnel/udp_nat.rs
+++ b/ostp-client/src/tunnel/udp_nat.rs
@@ -96,13 +96,8 @@ async fn start_udp_session(
         }
     }
 
-    let udp = match relay_addr {
-        SocketAddr::V4(_) => UdpSocket::bind("127.0.0.1:0").await?,
-        SocketAddr::V6(_) => match UdpSocket::bind("[::1]:0").await {
-            Ok(sock) => sock,
-            Err(_) => UdpSocket::bind("[::]:0").await?,
-        },
-    };
+    // Local SOCKS5 proxy always returns 127.0.0.1 (IPv4), so always bind IPv4
+    let udp = UdpSocket::bind("127.0.0.1:0").await?;
     
     let mut buf = vec![0u8; 65536];
     
@@ -128,36 +123,45 @@ async fn start_udp_session(
                 }
             }
             res = udp.recv_from(&mut buf) => {
-                let (len, _peer) = res?;
-                if len < 10 { continue; } // At least 10 bytes for SOCKS5 header
-                let frag = buf[2];
-                if frag != 0 { continue; } // fragment not supported
-                let atyp = buf[3];
-                let (header_len, remote_dst) = match atyp {
-                    1 => {
-                        if len < 10 { continue; }
-                        let ip = std::net::Ipv4Addr::new(buf[4], buf[5], buf[6], buf[7]);
-                        let port = u16::from_be_bytes([buf[8], buf[9]]);
-                        (10, SocketAddr::new(std::net::IpAddr::V4(ip), port))
+                match res {
+                    Err(e) => {
+                        tracing::debug!("udp_nat recv_from error: {}", e);
+                        continue; // transient error, don't kill the session
                     }
-                    4 => {
-                        if len < 22 { continue; }
-                        let mut octets = [0u8; 16];
-                        octets.copy_from_slice(&buf[4..20]);
-                        let ip = std::net::Ipv6Addr::from(octets);
-                        let port = u16::from_be_bytes([buf[20], buf[21]]);
-                        (22, SocketAddr::new(std::net::IpAddr::V6(ip), port))
+                    Ok((len, _peer)) => {
+                        if len < 4 { continue; }
+                        let frag = buf[2];
+                        if frag != 0 { continue; } // fragment not supported
+                        let atyp = buf[3];
+                        let (header_len, remote_dst) = match atyp {
+                            1 => {
+                                if len < 10 { continue; }
+                                let ip = std::net::Ipv4Addr::new(buf[4], buf[5], buf[6], buf[7]);
+                                let port = u16::from_be_bytes([buf[8], buf[9]]);
+                                (10, SocketAddr::new(std::net::IpAddr::V4(ip), port))
+                            }
+                            4 => {
+                                if len < 22 { continue; }
+                                let mut octets = [0u8; 16];
+                                octets.copy_from_slice(&buf[4..20]);
+                                let ip = std::net::Ipv6Addr::from(octets);
+                                let port = u16::from_be_bytes([buf[20], buf[21]]);
+                                (22, SocketAddr::new(std::net::IpAddr::V6(ip), port))
+                            }
+                            _ => continue,
+                        };
+                        let payload = buf[header_len..len].to_vec();
+                        use futures::SinkExt;
+                        let _ = smoltcp_tx.lock().await.send((payload, remote_dst, client_src)).await;
                     }
-                    _ => continue, // Domain name not supported for incoming packets in typical UDP associate
-                };
-                let payload = buf[header_len..len].to_vec();
-                use futures::SinkExt;
-                let _ = smoltcp_tx.lock().await.send((payload, remote_dst, client_src)).await;
+                }
             }
             // If TCP drops, UDP association is over
             res = tcp.read(&mut tcp_buf) => {
-                let n = res?;
-                if n == 0 { break; }
+                match res {
+                    Ok(0) | Err(_) => break,
+                    Ok(_) => {}
+                }
             }
         }
     }