diff --git a/src/blockchain/blockchain.rs b/src/blockchain/blockchain.rs index 16dd3b5..235a5b6 100644 --- a/src/blockchain/blockchain.rs +++ b/src/blockchain/blockchain.rs @@ -65,6 +65,15 @@ impl Blockchain { } pub fn add_block(&mut self, block: Block) -> Result<(), &str> { + match &self.last_block { + None => {} + Some(last_block) => { + if last_block.index >= block.index && last_block.hash == block.hash { + info!("Ignoring block {}, we already have it", block.index); + return Err("Already have that block"); + } + } + } if !self.check_block(&block, &self.last_block) { warn!("Bad block found, ignoring:\n{:?}", &block); return Err("Bad block found, ignoring"); diff --git a/src/p2p/network.rs b/src/p2p/network.rs index db149c6..a83b4b2 100644 --- a/src/p2p/network.rs +++ b/src/p2p/network.rs @@ -273,7 +273,9 @@ fn handle_message(context: Arc>, message: Message, peers: &mut Pe return State::Error; } if ok { - if height > my_height { + let peer = peers.get_mut_peer(token).unwrap(); + peer.set_height(height); + if peer.is_higher(my_height) { State::message(Message::GetBlock { index: my_height }) } else { State::message(Message::GetPeers) @@ -284,14 +286,18 @@ fn handle_message(context: Arc>, message: Message, peers: &mut Pe } Message::Error => { State::Error } Message::Ping { height } => { - if height > my_height { + let peer = peers.get_mut_peer(token).unwrap(); + peer.set_height(height); + if peer.is_higher(my_height) { State::message(Message::GetBlock { index: my_height }) } else { State::message(Message::pong(my_height)) } } Message::Pong { height } => { - if height > my_height { + let peer = peers.get_mut_peer(token).unwrap(); + peer.set_height(height); + if peer.is_higher(my_height) { State::message(Message::GetBlock { index: my_height }) } else { State::idle() diff --git a/src/p2p/peer.rs b/src/p2p/peer.rs index af73fc1..18efaf8 100644 --- a/src/p2p/peer.rs +++ b/src/p2p/peer.rs @@ -7,13 +7,14 @@ pub struct Peer { addr: SocketAddr, stream: TcpStream, state: State, + height: u64, inbound: bool, public: bool, } impl Peer { pub fn new(addr: SocketAddr, stream: TcpStream, state: State, inbound: bool) -> Self { - Peer { addr, stream, state, inbound, public: false } + Peer { addr, stream, state, height: 0, inbound, public: false } } pub fn get_addr(&self) -> SocketAddr { @@ -36,6 +37,14 @@ impl Peer { self.state = state; } + pub fn set_height(&mut self, height: u64) { + self.height = height; + } + + pub fn is_higher(&self, height: u64) -> bool { + self.height > height + } + pub fn is_public(&self) -> bool { self.public } diff --git a/src/p2p/peers.rs b/src/p2p/peers.rs index f7ee712..ff5592a 100644 --- a/src/p2p/peers.rs +++ b/src/p2p/peers.rs @@ -6,6 +6,7 @@ use crate::p2p::{Peer, State, Message}; use crate::p2p::network::LISTEN_PORT; use crate::p2p::network::next; use rand::random; +use rand::seq::IteratorRandom; #[allow(unused_imports)] use log::{trace, debug, info, warn, error}; @@ -98,6 +99,7 @@ impl Peers { } pub fn send_pings(&mut self, registry: &Registry, height: u64) { + let mut ping_sent = false; for (token, peer) in self.peers.iter_mut() { match peer.get_state() { State::Idle { from } => { @@ -113,12 +115,28 @@ impl Peers { peer.set_state(State::message(message)); let stream = peer.get_stream(); registry.reregister(stream, token.clone(), Interest::WRITABLE).unwrap(); + ping_sent = true; } } _ => {} } } + if !ping_sent { + let mut rng = rand::thread_rng(); + match self.peers + .iter_mut() + .filter_map(|(token, peer)| if peer.get_state().is_idle() && peer.is_higher(height) { Some((token, peer)) } else { None }) + .choose(&mut rng) { + None => {} + Some((token, peer)) => { + debug!("Found some peer higher than we are, sending block request"); + registry.reregister(peer.get_stream(), token.clone(), Interest::WRITABLE).unwrap(); + peer.set_state(State::message(Message::GetBlock { index: height })); + } + } + } + for (token, peer) in self.peers.iter_mut() { if peer.get_state().need_reconnect() { let addr = peer.get_addr();