use anyhow::{Context, Result};
use bitcoin_hashes::hex::ToHex;
use tokio::{io::AsyncWriteExt, net::TcpStream, select, sync::mpsc::Sender};
use crate::{
chaindef::{net_magic, NetworkMessage, RawNetworkMessage, Transaction},
def::ROSTRUM_VERSION,
errors::ConnectionError,
metrics,
signal::{NetworkNotifier, Waiter},
util::Channel,
};
use bitcoincash::consensus::encode;
use bitcoincash::{
network::{address, message_blockdata::Inventory, message_network},
Network, Txid,
};
use rand::prelude::*;
use std::net::{Ipv4Addr, SocketAddr};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use std::{collections::HashSet, io::ErrorKind, sync::atomic::AtomicBool, sync::atomic::Ordering};
use prometheus::HistogramVec;
use super::p2pmessage::P2PNetworkMessage;
pub struct P2PConnection {
is_broken: Arc<AtomicBool>,
netmsg_sender: Sender<Option<NetworkMessage>>,
}
impl P2PConnection {
pub(crate) async fn connect(
network: Network,
address: SocketAddr,
network_signals: Arc<NetworkNotifier>,
send_duration: Arc<HistogramVec>,
recv_duration: Arc<HistogramVec>,
parse_duration: Arc<HistogramVec>,
recv_size: Arc<HistogramVec>,
) -> Result<Self, ConnectionError> {
let conn = match TcpStream::connect(address).await {
Ok(c) => c,
Err(err) => {
return Err(ConnectionError {
msg: format!("p2p failed to connect to {:?} {}", address, err),
})
}
};
let (mut stream_read, mut stream_write) = conn.into_split();
let netmsg_channel = Channel::<Option<NetworkMessage>>::bounded(1000);
let netmsg_sender_self = netmsg_channel.sender();
let netmsg_sender = netmsg_channel.sender();
let rawmsg_channel = Channel::<P2PNetworkMessage>::bounded(1);
let rawmsg_sender = rawmsg_channel.sender();
let is_broken: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
let is_broken_copy = Arc::clone(&is_broken);
crate::thread::spawn_task("p2p_send", async move {
loop {
let mut netmsg_recv = netmsg_channel.receiver().await;
if Waiter::shutdown_check().is_err() {
is_broken_copy.store(true, Ordering::Relaxed);
let _ = stream_write.shutdown().await;
return Ok(());
}
let start = Instant::now();
let msg = match netmsg_recv.recv().await {
Some(Some(msg)) => msg,
None | Some(None) => {
is_broken_copy.store(true, Ordering::Relaxed);
let _ = stream_write.shutdown().await;
info!("p2p: disconnecting");
return Ok(());
}
};
metrics::observe(&send_duration, "wait", duration_to_seconds(start.elapsed()));
if let Err(e) = {
let start = Instant::now();
let raw_msg = RawNetworkMessage {
magic: net_magic(network),
payload: msg,
};
let res = stream_write
.write_all(encode::serialize(&raw_msg).as_slice())
.await
.context("p2p failed to send");
let duration = duration_to_seconds(start.elapsed());
metrics::observe(&send_duration, "send", duration);
res
} {
is_broken_copy.store(true, Ordering::Relaxed);
let _ = stream_write.shutdown().await;
return Err(e);
}
}
});
let is_broken_copy = Arc::clone(&is_broken);
crate::thread::spawn_task("p2p_recv", async move {
loop {
if Waiter::shutdown_check().is_err() {
return Ok(());
}
let start = Instant::now();
let p2p_msg = P2PNetworkMessage::async_consensus_decode(&mut stream_read).await;
{
let duration = duration_to_seconds(start.elapsed());
let label = format!(
"recv_{}",
p2p_msg
.as_ref()
.map(|msg| msg.cmd.as_ref())
.unwrap_or("err")
);
metrics::observe(&recv_duration, &label, duration);
}
let p2p_msg = match p2p_msg {
Ok(p2p_msg) => {
metrics::observe(
&recv_size,
p2p_msg.cmd.as_ref(),
p2p_msg.raw.len() as f64,
);
if p2p_msg.magic != net_magic(network) {
return Err(anyhow!(
"unexpected magic {} (instead of {})",
p2p_msg.magic,
net_magic(network)
));
}
p2p_msg
}
Err(encode::Error::Io(e)) if e.kind() == ErrorKind::UnexpectedEof => {
debug!("closing p2p_recv thread: connection closed");
is_broken_copy.store(true, Ordering::Relaxed);
return Ok(());
}
Err(e) => {
debug!("failed to recv a message from peer: {}", e);
is_broken_copy.store(true, Ordering::Relaxed);
return Ok(());
}
};
let start = Instant::now();
rawmsg_sender.send(p2p_msg).await?;
metrics::observe(&recv_duration, "send", duration_to_seconds(start.elapsed()));
}
});
let init_channel = Channel::<()>::bounded(1);
let init_channel_sender = init_channel.sender();
match netmsg_sender.send(Some(build_version_message())).await {
Ok(_) => {}
Err(err) => {
is_broken.store(true, Ordering::Relaxed);
return Err(ConnectionError {
msg: format!("p2p failed to send msg: {}", err),
});
}
}
let is_broken_copy = Arc::clone(&is_broken);
crate::thread::spawn_task("p2p_loop", async move {
let disconnect = || {
is_broken_copy.store(true, Ordering::Relaxed);
let _ = netmsg_sender.try_send(None);
};
loop {
if Waiter::shutdown_check().is_err() {
return Ok(());
}
let block_sender = network_signals.block_channel.sender();
let tx_sender = network_signals.tx_channel.sender();
let mut rawmsg_recv = rawmsg_channel.receiver().await;
select! {
raw_msg = rawmsg_recv.recv() => {
let raw_msg = match raw_msg {
Some(r) => r,
None => {
disconnect();
return Ok(())
}
};
let label = format!("parse_{}", raw_msg.cmd.as_ref());
let msg = match metrics::observe_duration(&parse_duration, &label, || raw_msg.parse()) {
Ok(msg) => msg,
Err(err) => {
info!("p2p: failed to parse '{}({:?})': {}", raw_msg.cmd, raw_msg.raw.to_hex(), err);
disconnect();
return Ok(());
}
};
match msg {
NetworkMessage::GetHeaders(_) => {
match netmsg_sender.send(Some(NetworkMessage::Headers(vec![]))).await {
Ok(_) => {},
Err(err) => {
debug!("failed to reply to getheaders: {}", err);
disconnect();
return Ok(())
}
}
}
NetworkMessage::Version(version) => {
debug!("p2p: peer version: {:?}", version);
match netmsg_sender.send(Some(NetworkMessage::Verack)).await {
Ok(_) => {},
Err(err) => {
debug!("failed to reply to version: {}", err);
disconnect();
return Ok(())
}
}
}
NetworkMessage::Inv(inventory) => {
let mut new_txs: HashSet<Txid> = HashSet::new();
for inv in inventory {
match inv {
Inventory::Transaction(txid) => {
new_txs.insert(txid);
}
Inventory::Block(blockhash) | Inventory::CompactBlock(blockhash) => {
#[cfg(nexa)]
{
let blockhash = crate::nexa::hash_types::BlockHash::from(blockhash);
if let Err(e) = block_sender.try_send(blockhash) {
trace!("Failed to notify of block: {}", e);
}
}
#[cfg(bch)]
{
let _ = block_sender.try_send(blockhash);
}
}
_ => { }
}
}
if let Err(e) = tx_sender.try_send(new_txs) {
trace!("p2p: Failed to notify of inv: {}", e);
}
},
NetworkMessage::Ping(nonce) => {
match netmsg_sender.send(Some(NetworkMessage::Pong(nonce))).await {
Ok(_) => {},
Err(err) => {
debug!("p2p: Failed to respond to ping: {}", err);
disconnect();
return Ok(())
}
};
}
NetworkMessage::Verack => {
match init_channel_sender.send(()).await {
Ok(_) => {},
Err(err) => {
debug!("p2p: Failed to end init version handshake");
disconnect();
return Err(err.into());
}
};
}
NetworkMessage::Block(_) => (),
NetworkMessage::Headers(_) => (),
NetworkMessage::Addr(_) => (),
msg => debug!("p2p: unexpected message: {:?}", msg),
}
}
}
}
});
let recv_ack = {
let sender_cpy = netmsg_sender_self.clone();
let broken_cpy = is_broken.clone();
let disconnect = move || {
let _ = sender_cpy.try_send(None);
broken_cpy.store(true, Ordering::Relaxed);
};
tokio::time::timeout(Duration::from_secs(5), async {
if init_channel.receiver().await.recv().await.is_none() {
disconnect();
return Err(ConnectionError {
msg: "p2p failed to receive verack".to_string(),
});
}
Ok(())
})
.await
};
match recv_ack {
Ok(r) => r?,
Err(e) => {
is_broken.store(true, Ordering::Relaxed);
let _ = netmsg_sender_self.try_send(None);
return Err(ConnectionError {
msg: format!("p2p failed to respond verack within {}", e),
});
}
}
Ok(P2PConnection {
is_broken,
netmsg_sender: netmsg_sender_self,
})
}
pub fn is_broken(&self) -> bool {
Waiter::shutdown_check().is_err() || self.is_broken.load(Ordering::Relaxed)
}
pub async fn send_tx(&self, tx: Transaction) -> Result<()> {
let msg = NetworkMessage::Tx(tx);
self.netmsg_sender
.send(Some(msg))
.await
.context("Failed to p2p send tx {}")?;
Ok(())
}
}
impl Drop for P2PConnection {
fn drop(&mut self) {
let _ = self.netmsg_sender.try_send(None);
}
}
fn build_version_message() -> NetworkMessage {
let addr = SocketAddr::new(Ipv4Addr::new(127, 0, 0, 1).into(), 0);
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("Time error")
.as_secs() as i64;
let services = bitcoincash::network::constants::ServiceFlags::NONE;
#[cfg(bch)]
let version = bitcoincash::network::constants::PROTOCOL_VERSION;
#[cfg(nexa)]
let version = 80003;
NetworkMessage::Version(message_network::VersionMessage {
version,
services,
timestamp,
receiver: address::Address::new(&addr, services),
sender: address::Address::new(&addr, services),
nonce: rand::rng().next_u64(),
user_agent: format!("/rostrum:{}/", ROSTRUM_VERSION),
start_height: 0,
relay: true,
})
}
#[inline]
pub(crate) fn duration_to_seconds(d: Duration) -> f64 {
let nanos = f64::from(d.subsec_nanos()) / 1e9;
d.as_secs() as f64 + nanos
}