From ca85f75c72d7b360bd71e60cd466de1e84e600f7 Mon Sep 17 00:00:00 2001 From: "Y. T. Chung" Date: Thu, 10 May 2018 01:26:55 +0800 Subject: [PATCH] [#106] Fixed bug, should not hold UDP connections forever --- src/relay/udprelay/local.rs | 27 +++++++++++++++++++++++++++ src/relay/udprelay/server.rs | 30 +++++++++++++++++++++++++++++- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/src/relay/udprelay/local.rs b/src/relay/udprelay/local.rs index c84fa33d..baaa6b00 100644 --- a/src/relay/udprelay/local.rs +++ b/src/relay/udprelay/local.rs @@ -4,11 +4,13 @@ use std::io::{self, Cursor, ErrorKind, Read}; use std::net::SocketAddr; use std::net::{IpAddr, Ipv4Addr}; use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; use futures::{self, Future, Stream}; use tokio; use tokio::net::UdpSocket; +use tokio::util::FutureExt; use tokio_io::IoFuture; use config::{Config, ServerAddr, ServerConfig}; @@ -47,6 +49,7 @@ fn listen(config: Arc, l: UdpSocket) -> IoFuture<()> { let svr_cfg_cloned_cloned = svr_cfg.clone(); let socket = socket.clone(); let config = config.clone(); + let timeout = *svr_cfg.timeout(); let rel = futures::lazy(|| UdpAssociateHeader::read_from(Cursor::new(pkt))) .map_err(From::from) @@ -82,12 +85,36 @@ fn listen(config: Arc, l: UdpSocket) -> IoFuture<()> { }) .and_then(move |(remote_udp, remote_addr, payload, addr)| { info!("UDP ASSOCIATE {} -> {}, payload length {} bytes", src, addr, payload.len()); + let to = timeout.unwrap_or(Duration::from_secs(5)); + let caddr = addr.clone(); remote_udp.send_dgram(payload, &remote_addr) + .deadline(Instant::now() + to) + .map_err(move |err| { + match err.into_inner() { + Some(e) => e, + None => { + error!("Udp associate sending datagram {} -> {} timed out in {:?}", src, caddr, to); + io::Error::new(io::ErrorKind::TimedOut, "udp send timed out") + } + } + }) .map(|(remote_udp, _)| (remote_udp, addr)) }) .and_then(move |(remote_udp, addr)| { let buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; + let to = timeout.unwrap_or(Duration::from_secs(5)); + let caddr = addr.clone(); remote_udp.recv_dgram(buf) + .deadline(Instant::now() + to) + .map_err(move |err| { + match err.into_inner() { + Some(e) => e, + None => { + error!("Udp associate waiting datagram {} <- {} timed out in {:?}", src, caddr, to); + io::Error::new(io::ErrorKind::TimedOut, "udp recv timed out") + } + } + }) .and_then(move |(_remote_udp, buf, n, _from)| { let svr_cfg = svr_cfg_cloned; decrypt_payload(svr_cfg.method(), svr_cfg.key(), &buf[..n]) diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs index 8fb0639d..fb92a99b 100644 --- a/src/relay/udprelay/server.rs +++ b/src/relay/udprelay/server.rs @@ -3,11 +3,13 @@ use std::io::{self, Cursor, ErrorKind}; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; use futures::{self, Future, Stream}; use tokio; use tokio::net::UdpSocket; +use tokio::util::FutureExt; use tokio_io::IoFuture; use config::{Config, ServerConfig}; @@ -50,6 +52,7 @@ fn listen(config: Arc, svr_cfg: Arc) -> IoFuture<()> { let svr_cfg_cloned = svr_cfg.clone(); let socket = socket.clone(); let config = config.clone(); + let timeout = *svr_cfg.timeout(); let rel = futures::lazy(move || decrypt_payload(svr_cfg.method(), svr_cfg.key(), &pkt)) .and_then(move |payload| { // Read Address in the front (ShadowSocks protocol) @@ -75,9 +78,21 @@ fn listen(config: Arc, svr_cfg: Arc) -> IoFuture<()> { .map(|(remote_udp, _)| (remote_udp, addr)) }) }) - .and_then(|(remote_udp, addr)| { + .and_then(move |(remote_udp, addr)| { let buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE]; + let to = timeout.unwrap_or(Duration::from_secs(5)); + let caddr = addr.clone(); remote_udp.recv_dgram(buf) + .deadline(Instant::now() + to) + .map_err(move |err| { + match err.into_inner() { + Some(e) => e, + None => { + error!("Udp associate waiting datagram {} -> {} timed out in {:?}", src, caddr, to); + io::Error::new(io::ErrorKind::TimedOut, "udp recv timed out") + } + } + }) .and_then(|(_remote_udp, buf, n, _from)| { let svr_cfg = svr_cfg_cloned; @@ -89,7 +104,20 @@ fn listen(config: Arc, svr_cfg: Arc) -> IoFuture<()> { }) .and_then(move |(buf, addr)| { info!("UDP ASSOCIATE {} <- {}, payload length {} bytes", src, addr, buf.len()); + + let to = timeout.unwrap_or(Duration::from_secs(5)); + let caddr = addr.clone(); SendDgramRc::new(socket, buf, src) + .deadline(Instant::now() + to) + .map_err(move |err| { + match err.into_inner() { + Some(e) => e, + None => { + error!("Udp associate sending datagram {} <- {} timed out in {:?}", src, caddr, to); + io::Error::new(io::ErrorKind::TimedOut, "udp recv timed out") + } + } + }) }) .map(|_| ());