Split HTLC tracking into separate Inbound/Outbound types

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

diff --git a/src/ln/peer_handler.rs b/src/ln/peer_handler.rs
index 3122e572e57fafcbd90e1c08c6c10b146d152c29..8a4228851b7a7e3981e698ccc6cd05f3cb3e5d3b 100644 (file)
--- 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<Descriptor: SocketDescriptor> {
        /// Only add to this set when noise completes:
        node_id_to_descriptor: HashMap<PublicKey, Descriptor>,
 }
+struct MutPeerHolder<'a, Descriptor: SocketDescriptor + 'a> {
+       peers: &'a mut HashMap<Descriptor, Peer>,
+       node_id_to_descriptor: &'a mut HashMap<PublicKey, Descriptor>,
+}
+impl<Descriptor: SocketDescriptor> PeerHolder<Descriptor> {
+       fn borrow_parts(&mut self) -> MutPeerHolder<Descriptor> {
+               MutPeerHolder {
+                       peers: &mut self.peers,
+                       node_id_to_descriptor: &mut self.node_id_to_descriptor,
+               }
+       }
+}
 
 pub struct PeerManager<Descriptor: SocketDescriptor> {
        message_handler: MessageHandler,
@@ -100,7 +112,6 @@ pub struct PeerManager<Descriptor: SocketDescriptor> {
        logger: Arc<Logger>,
 }
 
-
 macro_rules! encode_msg {
        ($msg: expr, $msg_code: expr) => {
                {
@@ -136,7 +147,12 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
        /// completed and we are sure the remote peer has the private key for the given node_id.
        pub fn get_peer_node_ids(&self) -> Vec<PublicKey> {
                let peers = self.peers.lock().unwrap();
-               peers.peers.values().filter_map(|p| p.their_node_id).collect()
+               peers.peers.values().filter_map(|p| {
+                       if !p.channel_encryptor.is_ready_for_encryption() || p.their_global_features.is_none() {
+                               return None;
+                       }
+                       p.their_node_id
+               }).collect()
        }
 
        /// Indicates a new outbound connection has been established to a node with the given node_id.
@@ -267,14 +283,14 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
 
        fn do_read_event(&self, peer_descriptor: &mut Descriptor, data: Vec<u8>) -> Result<bool, PeerHandleError> {
                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 +369,18 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
                                                                }
                                                        }
 
+                                                       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 +394,7 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
                                                                        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 +410,7 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
                                                                        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 {
@@ -417,6 +445,9 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
                                                                                                if msg.local_features.requires_unknown_bits() {
                                                                                                        return Err(PeerHandleError{ no_connection_possible: true });
                                                                                                }
+                                                                                               if peer.their_global_features.is_some() {
+                                                                                                       return Err(PeerHandleError{ no_connection_possible: false });
+                                                                                               }
                                                                                                peer.their_global_features = Some(msg.global_features);
                                                                                                peer.their_local_features = Some(msg.local_features);
 
@@ -607,15 +638,10 @@ impl<Descriptor: SocketDescriptor> PeerManager<Descriptor> {
 
                                        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
                };