use crate::util::RwLock;
use std::collections::VecDeque;
use std::net::{SocketAddr, TcpStream};
use std::sync::Arc;
use chrono::prelude::*;
use rand::{thread_rng, Rng};
use crate::core::core::hash::Hash;
use crate::core::pow::Difficulty;
use crate::msg::{
read_message, write_message, Hand, Shake, SockAddr, Type, PROTOCOL_VERSION, USER_AGENT,
};
use crate::peer::Peer;
use crate::types::{Capabilities, Direction, Error, P2PConfig, PeerInfo, PeerLiveInfo};
const NONCES_CAP: usize = 100;
const ADDRS_CAP: usize = 10;
pub struct Handshake {
nonces: Arc<RwLock<VecDeque<u64>>>,
pub addrs: Arc<RwLock<VecDeque<SocketAddr>>>,
genesis: Hash,
config: P2PConfig,
}
impl Handshake {
pub fn new(genesis: Hash, config: P2PConfig) -> Handshake {
Handshake {
nonces: Arc::new(RwLock::new(VecDeque::with_capacity(NONCES_CAP))),
addrs: Arc::new(RwLock::new(VecDeque::with_capacity(ADDRS_CAP))),
genesis,
config,
}
}
pub fn initiate(
&self,
capab: Capabilities,
total_difficulty: Difficulty,
self_addr: SocketAddr,
conn: &mut TcpStream,
) -> Result<PeerInfo, Error> {
let nonce = self.next_nonce();
let peer_addr = match conn.peer_addr() {
Ok(pa) => pa,
Err(e) => return Err(Error::Connection(e)),
};
let hand = Hand {
version: PROTOCOL_VERSION,
capabilities: capab,
nonce: nonce,
genesis: self.genesis,
total_difficulty: total_difficulty,
sender_addr: SockAddr(self_addr),
receiver_addr: SockAddr(peer_addr),
user_agent: USER_AGENT.to_string(),
};
write_message(conn, hand, Type::Hand)?;
let shake: Shake = read_message(conn, Type::Shake)?;
if shake.version != PROTOCOL_VERSION {
return Err(Error::ProtocolMismatch {
us: PROTOCOL_VERSION,
peer: shake.version,
});
} else if shake.genesis != self.genesis {
return Err(Error::GenesisMismatch {
us: self.genesis,
peer: shake.genesis,
});
}
let peer_info = PeerInfo {
capabilities: shake.capabilities,
user_agent: shake.user_agent,
addr: peer_addr,
version: shake.version,
live_info: Arc::new(RwLock::new(PeerLiveInfo {
total_difficulty: shake.total_difficulty,
height: 0,
last_seen: Utc::now(),
stuck_detector: Utc::now(),
})),
direction: Direction::Outbound,
};
if Peer::is_denied(&self.config, &peer_info.addr) {
return Err(Error::ConnectionClose);
}
debug!(
"Connected! Cumulative {} offered from {:?} {:?} {:?}",
shake.total_difficulty.to_num(),
peer_info.addr,
peer_info.user_agent,
peer_info.capabilities
);
Ok(peer_info)
}
pub fn accept(
&self,
capab: Capabilities,
total_difficulty: Difficulty,
conn: &mut TcpStream,
) -> Result<PeerInfo, Error> {
let hand: Hand = read_message(conn, Type::Hand)?;
if hand.version != PROTOCOL_VERSION {
return Err(Error::ProtocolMismatch {
us: PROTOCOL_VERSION,
peer: hand.version,
});
} else if hand.genesis != self.genesis {
return Err(Error::GenesisMismatch {
us: self.genesis,
peer: hand.genesis,
});
} else {
let nonces = self.nonces.read();
let addr = extract_ip(&hand.sender_addr.0, &conn);
if nonces.contains(&hand.nonce) {
let mut addrs = self.addrs.write();
addrs.push_back(addr);
if addrs.len() >= ADDRS_CAP {
addrs.pop_front();
}
return Err(Error::PeerWithSelf);
}
}
let peer_info = PeerInfo {
capabilities: hand.capabilities,
user_agent: hand.user_agent,
addr: extract_ip(&hand.sender_addr.0, &conn),
version: hand.version,
live_info: Arc::new(RwLock::new(PeerLiveInfo {
total_difficulty: hand.total_difficulty,
height: 0,
last_seen: Utc::now(),
stuck_detector: Utc::now(),
})),
direction: Direction::Inbound,
};
if Peer::is_denied(&self.config, &peer_info.addr) {
return Err(Error::ConnectionClose);
}
let shake = Shake {
version: PROTOCOL_VERSION,
capabilities: capab,
genesis: self.genesis,
total_difficulty: total_difficulty,
user_agent: USER_AGENT.to_string(),
};
write_message(conn, shake, Type::Shake)?;
trace!("Success handshake with {}.", peer_info.addr);
Ok(peer_info)
}
fn next_nonce(&self) -> u64 {
let nonce = thread_rng().gen();
let mut nonces = self.nonces.write();
nonces.push_back(nonce);
if nonces.len() >= NONCES_CAP {
nonces.pop_front();
}
nonce
}
}
fn extract_ip(advertised: &SocketAddr, conn: &TcpStream) -> SocketAddr {
match advertised {
&SocketAddr::V4(v4sock) => {
let ip = v4sock.ip();
if ip.is_loopback() || ip.is_unspecified() {
if let Ok(addr) = conn.peer_addr() {
return SocketAddr::new(addr.ip(), advertised.port());
}
}
}
&SocketAddr::V6(v6sock) => {
let ip = v6sock.ip();
if ip.is_loopback() || ip.is_unspecified() {
if let Ok(addr) = conn.peer_addr() {
return SocketAddr::new(addr.ip(), advertised.port());
}
}
}
}
advertised.clone()
}