From 9a779a639f7022a00a973183a0edcaf27369d046 Mon Sep 17 00:00:00 2001
From: Henry de Valence <hdevalence@hdevalence.ca>
Date: Mon, 21 Oct 2019 22:17:57 -0700
Subject: [PATCH] Add a timeout for peer handshakes.

---
 zebra-network/src/config.rs   | 3 +++
 zebra-network/src/peer_set.rs | 6 +++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/zebra-network/src/config.rs b/zebra-network/src/config.rs
index f5d8f4ce..8f91a2b4 100644
--- a/zebra-network/src/config.rs
+++ b/zebra-network/src/config.rs
@@ -22,6 +22,8 @@ pub struct Config {
     pub ewma_decay_time: Duration,
     /// The outgoing request buffer size for the peer set.
     pub peerset_request_buffer_size: usize,
+    /// The timeout for peer handshakes.
+    pub handshake_timeout: Duration,
 }
 
 impl Default for Config {
@@ -36,6 +38,7 @@ impl Default for Config {
             ewma_default_rtt: Duration::from_secs(1),
             ewma_decay_time: Duration::from_secs(60),
             peerset_request_buffer_size: 1,
+            handshake_timeout: Duration::from_secs(4),
         }
     }
 }
diff --git a/zebra-network/src/peer_set.rs b/zebra-network/src/peer_set.rs
index 99465cdd..bcccc7be 100644
--- a/zebra-network/src/peer_set.rs
+++ b/zebra-network/src/peer_set.rs
@@ -19,6 +19,7 @@ use tokio::net::{TcpListener, TcpStream};
 use tower::{
     buffer::Buffer,
     discover::{Change, ServiceStream},
+    timeout::Timeout,
     Service, ServiceExt,
 };
 use tower_load::{peak_ewma::PeakEwmaDiscover, NoInstrument};
@@ -74,7 +75,10 @@ where
 {
     let (address_book, timestamp_collector) = TimestampCollector::spawn();
     let peer_connector = Buffer::new(
-        PeerConnector::new(config.clone(), inbound_service, timestamp_collector),
+        Timeout::new(
+            PeerConnector::new(config.clone(), inbound_service, timestamp_collector),
+            config.handshake_timeout,
+        ),
         1,
     );