From 6dab6b185d15f0ca1c9e891030e43187c58a669e Mon Sep 17 00:00:00 2001 From: Matt Corallo Date: Sat, 8 Sep 2018 13:56:45 -0400 Subject: [PATCH] Disconnect duplicate node_id connections after Noise handshake --- src/ln/peer_handler.rs | 44 +++++++++++++++++++++++++++++------------- 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/src/ln/peer_handler.rs b/src/ln/peer_handler.rs index 3122e572e..d7a7553ed 100644 --- a/src/ln/peer_handler.rs +++ b/src/ln/peer_handler.rs @@ -7,7 +7,7 @@ use util::byte_utils; use util::events::{EventsProvider,Event}; use util::logger::Logger; -use std::collections::{HashMap,LinkedList}; +use std::collections::{HashMap,hash_map,LinkedList}; use std::sync::{Arc, Mutex}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::{cmp,error,mem,hash,fmt}; @@ -90,6 +90,18 @@ struct PeerHolder { /// Only add to this set when noise completes: node_id_to_descriptor: HashMap, } +struct MutPeerHolder<'a, Descriptor: SocketDescriptor + 'a> { + peers: &'a mut HashMap, + node_id_to_descriptor: &'a mut HashMap, +} +impl PeerHolder { + fn borrow_parts(&mut self) -> MutPeerHolder { + MutPeerHolder { + peers: &mut self.peers, + node_id_to_descriptor: &mut self.node_id_to_descriptor, + } + } +} pub struct PeerManager { message_handler: MessageHandler, @@ -100,7 +112,6 @@ pub struct PeerManager { logger: Arc, } - macro_rules! encode_msg { ($msg: expr, $msg_code: expr) => { { @@ -267,14 +278,14 @@ impl PeerManager { fn do_read_event(&self, peer_descriptor: &mut Descriptor, data: Vec) -> Result { let pause_read = { - let mut peers = self.peers.lock().unwrap(); - let (should_insert_node_id, pause_read) = match peers.peers.get_mut(peer_descriptor) { + let mut peers_lock = self.peers.lock().unwrap(); + let peers = peers_lock.borrow_parts(); + let pause_read = match peers.peers.get_mut(peer_descriptor) { None => panic!("Descriptor for read_event is not already known to PeerManager"), Some(peer) => { assert!(peer.pending_read_buffer.len() > 0); assert!(peer.pending_read_buffer.len() > peer.pending_read_buffer_pos); - let mut insert_node_id = None; let mut read_pos = 0; while read_pos < data.len() { { @@ -353,6 +364,18 @@ impl PeerManager { } } + macro_rules! insert_node_id { + () => { + match peers.node_id_to_descriptor.entry(peer.their_node_id.unwrap()) { + hash_map::Entry::Occupied(_) => { + peer.their_node_id = None; // Unset so that we don't generate a peer_disconnected event + return Err(PeerHandleError{ no_connection_possible: false }) + }, + hash_map::Entry::Vacant(entry) => entry.insert(peer_descriptor.clone()), + }; + } + } + let next_step = peer.channel_encryptor.get_noise_step(); match next_step { NextNoiseStep::ActOne => { @@ -366,7 +389,7 @@ impl PeerManager { peer.pending_read_buffer = [0; 18].to_vec(); // Message length header is 18 bytes peer.pending_read_is_header = true; - insert_node_id = Some(peer.their_node_id.unwrap()); + insert_node_id!(); let mut local_features = msgs::LocalFeatures::new(); if self.initial_syncs_sent.load(Ordering::Acquire) < INITIAL_SYNCS_TO_SEND { self.initial_syncs_sent.fetch_add(1, Ordering::AcqRel); @@ -382,7 +405,7 @@ impl PeerManager { peer.pending_read_buffer = [0; 18].to_vec(); // Message length header is 18 bytes peer.pending_read_is_header = true; peer.their_node_id = Some(their_node_id); - insert_node_id = Some(peer.their_node_id.unwrap()); + insert_node_id!(); }, NextNoiseStep::NoiseComplete => { if peer.pending_read_is_header { @@ -607,15 +630,10 @@ impl PeerManager { Self::do_attempt_write_data(peer_descriptor, peer); - (insert_node_id /* should_insert_node_id */, peer.pending_outbound_buffer.len() > 10) // pause_read + peer.pending_outbound_buffer.len() > 10 // pause_read } }; - match should_insert_node_id { - Some(node_id) => { peers.node_id_to_descriptor.insert(node_id, peer_descriptor.clone()); }, - None => {} - }; - pause_read }; -- 2.39.5