Improved peer reconnection code.

src/p2p/network.rs

@@ -88,28 +88,8 @@ impl Network {
                             poll.registry().reregister(&mut server, SERVER, Interest::READABLE).expect("Error reregistering server");
                         }
                         token => {
-                            match peers.get_mut_peer(&token) {
-                                Some(_peer) => {
-                                    match handle_connection_event(context.clone(), &mut peers, &poll.registry(), &event) {
-                                        Ok(result) => {
-                                            if !result {
-                                                match peers.remove_peer(&token) {
-                                                    None => {}
-                                                    Some(mut peer) => {
-                                                        let stream = peer.get_stream();
-                                                        let _ = poll.registry().deregister(stream);
-                                                        let _ = stream.shutdown(Shutdown::Both);
-                                                        info!("Peer connection {:?} has shut down", &peer.get_addr());
-                                                    }
-                                                }
-                                            }
-                                        }
-                                        Err(_err) => {
-                                            peers.remove_peer(&token);
-                                        }
-                                    }
-                                }
-                                None => { warn!("Odd event from poll"); }
+                            if !handle_connection_event(context.clone(), &mut peers, &poll.registry(), &event) {
+                                let _ = peers.close_peer(poll.registry(), &token);
                             }
                         }
                     }
@@ -126,9 +106,9 @@ impl Network {
     }
 }
 
-fn handle_connection_event(context: Arc<Mutex<Context>>, peers: &mut Peers, registry: &Registry, event: &Event) -> io::Result<bool> {
+fn handle_connection_event(context: Arc<Mutex<Context>>, peers: &mut Peers, registry: &Registry, event: &Event) -> bool {
     if event.is_error() || (event.is_read_closed() && event.is_write_closed()) {
-        return Ok(false);
+        return false;
     }
 
     if event.is_readable() {
@@ -168,14 +148,10 @@ fn handle_connection_event(context: Arc<Mutex<Context>>, peers: &mut Peers, regi
                         }
                     }
                 }
-                Err(_) => { return Ok(false); }
+                Err(_) => { return false; }
             }
         } else {
-            // Try to reregister connection
-            let peer = peers.get_mut_peer(&event.token()).expect("Error getting peer for connection");
-            let mut stream = peer.get_stream();
-            registry.reregister(stream, event.token(), Interest::READABLE);
-            return Ok(true);
+            return false;
         }
     }
 
@@ -220,7 +196,7 @@ fn handle_connection_event(context: Arc<Mutex<Context>>, peers: &mut Peers, regi
         }
     }
 
-    Ok(true)
+    true
 }
 
 fn read_message(stream: &mut TcpStream) -> Result<Vec<u8>, ()> {
@@ -233,7 +209,7 @@ fn read_message(stream: &mut TcpStream) -> Result<Vec<u8>, ()> {
         }
     };
     trace!("Payload size is {}", data_size);
-    if data_size > MAX_PACKET_SIZE {
+    if data_size > MAX_PACKET_SIZE || data_size == 0 {
         return Err(());
     }
 

src/p2p/peer.rs

@@ -24,6 +24,10 @@ impl Peer {
         &mut self.stream
     }
 
+    pub fn set_stream(&mut self, stream: TcpStream) {
+        self.stream = stream;
+    }
+
     pub fn get_state(&self) -> &State {
         &self.state
     }
@@ -48,6 +52,10 @@ impl Peer {
         self.state.disabled()
     }
 
+    pub fn is_inbound(&self) -> bool {
+        self.inbound
+    }
+
     /// If loopback address then we care about ip and port.
     /// If regular address then we only care about the ip and ignore the port.
     pub fn equals(&self, addr: &SocketAddr) -> bool {

src/p2p/peers.rs

@@ -1,5 +1,5 @@
 use std::collections::HashMap;
-use std::net::SocketAddr;
+use std::net::{SocketAddr, Shutdown};
 use mio::{Token, Interest, Registry};
 use mio::net::TcpStream;
 use crate::p2p::{Peer, State, Message};
@@ -32,8 +32,23 @@ impl Peers {
         self.peers.get_mut(token)
     }
 
-    pub fn remove_peer(&mut self, token: &Token) -> Option<Peer> {
-        self.peers.remove(token)
+    pub fn close_peer(&mut self, registry: &Registry, token: &Token) {
+        let mut peer = self.peers.get_mut(token);
+        match peer {
+            Some(mut peer) => {
+                let mut stream = peer.get_stream();
+                let _ = stream.shutdown(Shutdown::Both);
+                registry.deregister(stream);
+                info!("Peer connection {:?} has shut down", &peer.get_addr());
+
+                if !peer.disabled() && !peer.is_inbound() {
+                    peer.set_state(State::offline());
+                } else {
+                    self.peers.remove(token);
+                }
+            }
+            None => {}
+        }
     }
 
     pub fn add_peers_from_exchange(&mut self, peers: Vec<String>) {
@@ -41,7 +56,8 @@ impl Peers {
         // TODO make it return error if these peers are wrong and seem like an attack
         for peer in peers.iter() {
             let addr: SocketAddr = peer.parse().expect(&format!("Error parsing peer {}", peer));
-            if addr.ip().is_loopback() {
+            if skip_addr(&addr) {
+                info!("Skipping address from exchange: {}", &addr);
                 continue; // Return error in future
             }
             let mut found = false;
@@ -101,6 +117,23 @@ impl Peers {
                 _ => {}
             }
         }
+
+        for (token, peer) in self.peers.iter_mut() {
+            if peer.get_state().need_reconnect() {
+                let addr = peer.get_addr();
+                match TcpStream::connect(addr.clone()) {
+                    Ok(mut stream) => {
+                        registry.register(&mut stream, token.clone(), Interest::WRITABLE).unwrap();
+                        peer.set_state(State::Connecting);
+                        peer.set_stream(stream);
+                        info!("Created connection to peer {}", &addr);
+                    }
+                    Err(e) => {
+                        error!("Error connecting to peer {}: {}", &addr, e);
+                    }
+                }
+            }
+        }
     }
 
     pub fn connect_new_peers(&mut self, registry: &Registry, unique_token: &mut Token) {
@@ -124,4 +157,25 @@ impl Peers {
         }
         self.new_peers.clear();
     }
+}
+
+fn skip_addr(addr: &SocketAddr) -> bool {
+    if addr.ip().is_loopback() {
+        return true;
+    }
+    match addr {
+        SocketAddr::V4(addr) => {
+            if addr.ip().is_private() {
+                return true;
+            }
+        }
+        SocketAddr::V6(addr) => {
+            // TODO uncomment when stabilized
+            // if addr.ip().is_unique_local() {
+            //     return true;
+            // }
+        }
+    }
+
+    false
 }

src/p2p/state.rs

@@ -36,6 +36,13 @@ impl State {
         }
     }
 
+    pub fn is_idle(&self) -> bool {
+        match self {
+            State::Idle { .. } => { true }
+            _ => { false }
+        }
+    }
+
     pub fn disabled(&self) -> bool {
         match self {
             State::Error => { true }
@@ -46,4 +53,11 @@ impl State {
             _ => { false }
         }
     }
+
+    pub fn need_reconnect(&self) -> bool {
+        match self {
+            State::Offline { from } => { from.elapsed().as_secs() > 60 }
+            _ => { false }
+        }
+    }
 }