Give peers one timer tick to finish handshake before disconnecting

[rust-lightning] / lightning / src / ln / peer_handler.rs

diff --git a/lightning/src/ln/peer_handler.rs b/lightning/src/ln/peer_handler.rs
index fb560941d56672476d68762ee44b6542e73d5681..b082bdbe6c6bf7d4e7e50851118f8340e5aea0e7 100644 (file)
--- a/lightning/src/ln/peer_handler.rs
+++ b/lightning/src/ln/peer_handler.rs
@@ -860,6 +860,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
                                                                        let features = InitFeatures::known();
                                                                        let resp = msgs::Init { features };
                                                                        self.enqueue_message(peer, &resp);
+                                                                       peer.awaiting_pong_tick_intervals = 0;
                                                                },
                                                                NextNoiseStep::ActThree => {
                                                                        let their_node_id = try_potential_handleerror!(peer.channel_encryptor.process_act_three(&peer.pending_read_buffer[..]));
@@ -870,6 +871,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
                                                                        let features = InitFeatures::known();
                                                                        let resp = msgs::Init { features };
                                                                        self.enqueue_message(peer, &resp);
+                                                                       peer.awaiting_pong_tick_intervals = 0;
                                                                },
                                                                NextNoiseStep::NoiseComplete => {
                                                                        if peer.pending_read_is_header {
@@ -1530,13 +1532,21 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
                        let peer_count = peers.len();
 
                        peers.retain(|descriptor, peer| {
-                               if !peer.channel_encryptor.is_ready_for_encryption() {
-                                       // The peer needs to complete its handshake before we can exchange messages
-                                       return true;
+                               let mut do_disconnect_peer = false;
+                               if !peer.channel_encryptor.is_ready_for_encryption() || peer.their_node_id.is_none() {
+                                       // The peer needs to complete its handshake before we can exchange messages. We
+                                       // give peers one timer tick to complet handshake.
+                                       if peer.awaiting_pong_tick_intervals != 0 {
+                                               do_disconnect_peer = true;
+                                       } else {
+                                               peer.awaiting_pong_tick_intervals = 1;
+                                               return true;
+                                       }
                                }
 
-                               if (peer.awaiting_pong_timer_tick_intervals > 0 && !peer.received_message_since_timer_tick)
-                                       || peer.awaiting_pong_timer_tick_intervals as u64 >
+                               if do_disconnect_peer
+                                       || (peer.awaiting_pong_timertick_intervals > 0 && !peer.received_message_since_timer_tick)
+                                       || peer.awaiting_pong_timertick_intervals as u64 >
                                                MAX_BUFFER_DRAIN_TICK_INTERVALS_PER_PEER as u64 * peer_count as u64
                                {
                                        descriptors_needing_disconnect.push(descriptor.clone());
@@ -1546,11 +1556,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
                                                        node_id_to_descriptor.remove(&node_id);
                                                        self.message_handler.chan_handler.peer_disconnected(&node_id, false);
                                                }
-                                               None => {
-                                                       // This can't actually happen as we should have hit
-                                                       // is_ready_for_encryption() previously on this same peer.
-                                                       unreachable!();
-                                               },
+                                               None => {},
                                        }
                                        return false;
                                }