Time out reply socket after idle, not fixed, time

3 files changed, 19 insertions, 17 deletions

diff --git a/socket-relay/Cargo.toml b/socket-relay/Cargo.toml
index b7141bec8f..b40c72f6f2 100644
--- a/socket-relay/Cargo.toml
+++ b/socket-relay/Cargo.toml
@@ -13,3 +13,4 @@ error-chain = "0.11"
 futures = "0.1.15"
 log = "0.3"
 tokio-core = "0.1"
+tokio-timer = "0.1"
diff --git a/socket-relay/src/lib.rs b/socket-relay/src/lib.rs
index a342cce2dc..c0b930e2dc 100644
--- a/socket-relay/src/lib.rs
+++ b/socket-relay/src/lib.rs
@@ -2,6 +2,7 @@ extern crate futures;
 #[macro_use]
 extern crate log;
 extern crate tokio_core;
+extern crate tokio_timer;
 
 /// UDP based relaying.
 pub mod udp;
diff --git a/socket-relay/src/udp.rs b/socket-relay/src/udp.rs
index cc9a9139a1..00c5a95417 100644
--- a/socket-relay/src/udp.rs
+++ b/socket-relay/src/udp.rs
@@ -11,9 +11,10 @@ use std::result::Result as StdResult;
 use std::time::Duration;
 
 use tokio_core::net::{UdpCodec, UdpSocket};
-use tokio_core::reactor::{Handle, Timeout};
+use tokio_core::reactor::Handle;
+use tokio_timer::Timer;
 
-/// The amount of time the forwarding socket is open for replies.
+/// The amount of idle (no replies) time needed for the forwarding socket to close.
 pub static FORWARD_TIMEOUT_MS: u64 = 60000;
 
 /// Number of slots in internal channel transfering responses back to clients.
@@ -96,24 +97,23 @@ impl Relay {
             error!("Error while forwarding to destination addr: {}", e);
         });
 
-        let recv_future = forward_stream
-            .filter_map(move |(addr, data)| if addr == destination {
-                Some(data)
-            } else {
-                None
-            })
-            .map_err(|e| {
-                error!("Error reading datagrams from forward socket: {}", e)
-            })
+        let recv_future = Timer::default()
+            .timeout_stream(
+                forward_stream
+                    .filter_map(move |(addr, data)| if addr == destination {
+                        Some(data)
+                    } else {
+                        None
+                    })
+                    .map_err(|e| {
+                        error!("Error reading datagrams from forward socket: {}", e)
+                    }),
+                Duration::from_millis(FORWARD_TIMEOUT_MS),
+            )
             .forward(response_sink)
             .map(|_| ());
 
-        let network_future = send_future.and_then(|_| recv_future);
-
-        let timeout =
-            Timeout::new(Duration::from_millis(FORWARD_TIMEOUT_MS), &handle)?.map_err(|_| ());
-
-        handle.spawn(network_future.select(timeout).map(|_| ()).map_err(|_| ()));
+        handle.spawn(send_future.and_then(|_| recv_future));
         Ok(())
     }