use crate::{
config, error,
peer_store::{PeerStoreHandle, PeerStoreProvider},
protocol_controller::{self, SetId},
types::ProtocolName,
};
use bytes::Bytes;
use codec::{DecodeAll, Encode};
use futures::{channel::oneshot, stream::FuturesUnordered, StreamExt};
use libp2p::{
core::Endpoint,
swarm::{
behaviour::FromSwarm, ConnectionDenied, ConnectionId, NetworkBehaviour, PollParameters,
THandler, THandlerInEvent, THandlerOutEvent, ToSwarm,
},
Multiaddr, PeerId,
};
use log::{debug, error, warn};
use sc_network_common::{role::Roles, sync::message::BlockAnnouncesHandshake};
use sc_utils::mpsc::{TracingUnboundedReceiver, TracingUnboundedSender};
use sp_runtime::traits::Block as BlockT;
use std::{
collections::{HashMap, HashSet},
future::Future,
iter,
pin::Pin,
task::Poll,
};
use message::{generic::Message as GenericMessage, Message};
use notifications::{Notifications, NotificationsOut};
pub use notifications::{NotificationsSink, NotifsHandlerError, Ready};
mod notifications;
pub mod message;
pub(crate) const BLOCK_ANNOUNCES_TRANSACTIONS_SUBSTREAM_SIZE: u64 = 16 * 1024 * 1024;
const HARDCODED_PEERSETS_SYNC: SetId = SetId::from(0);
mod rep {
use crate::ReputationChange as Rep;
pub const BAD_MESSAGE: Rep = Rep::new(-(1 << 12), "Bad message");
}
type PendingSyncSubstreamValidation =
Pin<Box<dyn Future<Output = Result<(PeerId, Roles), PeerId>> + Send>>;
pub struct Protocol<B: BlockT> {
peer_store_handle: PeerStoreHandle,
behaviour: Notifications,
notification_protocols: Vec<ProtocolName>,
bad_handshake_substreams: HashSet<(PeerId, SetId)>,
peers: HashMap<PeerId, Roles>,
sync_substream_validations: FuturesUnordered<PendingSyncSubstreamValidation>,
tx: TracingUnboundedSender<crate::event::SyncEvent<B>>,
_marker: std::marker::PhantomData<B>,
}
impl<B: BlockT> Protocol<B> {
pub fn new(
roles: Roles,
notification_protocols: Vec<config::NonDefaultSetConfig>,
block_announces_protocol: config::NonDefaultSetConfig,
peer_store_handle: PeerStoreHandle,
protocol_controller_handles: Vec<protocol_controller::ProtocolHandle>,
from_protocol_controllers: TracingUnboundedReceiver<protocol_controller::Message>,
tx: TracingUnboundedSender<crate::event::SyncEvent<B>>,
) -> error::Result<Self> {
let behaviour = {
Notifications::new(
protocol_controller_handles,
from_protocol_controllers,
iter::once(notifications::ProtocolConfig {
name: block_announces_protocol.notifications_protocol.clone(),
fallback_names: block_announces_protocol.fallback_names.clone(),
handshake: block_announces_protocol.handshake.as_ref().unwrap().to_vec(),
max_notification_size: block_announces_protocol.max_notification_size,
})
.chain(notification_protocols.iter().map(|s| notifications::ProtocolConfig {
name: s.notifications_protocol.clone(),
fallback_names: s.fallback_names.clone(),
handshake: s.handshake.as_ref().map_or(roles.encode(), |h| (*h).to_vec()),
max_notification_size: s.max_notification_size,
})),
)
};
let protocol = Self {
peer_store_handle,
behaviour,
notification_protocols: iter::once(block_announces_protocol.notifications_protocol)
.chain(notification_protocols.iter().map(|s| s.notifications_protocol.clone()))
.collect(),
bad_handshake_substreams: Default::default(),
peers: HashMap::new(),
sync_substream_validations: FuturesUnordered::new(),
tx,
_marker: Default::default(),
};
Ok(protocol)
}
pub fn open_peers(&self) -> impl Iterator<Item = &PeerId> {
self.behaviour.open_peers()
}
pub fn disconnect_peer(&mut self, peer_id: &PeerId, protocol_name: ProtocolName) {
if let Some(position) = self.notification_protocols.iter().position(|p| *p == protocol_name)
{
self.behaviour.disconnect_peer(peer_id, SetId::from(position));
} else {
warn!(target: "sub-libp2p", "disconnect_peer() with invalid protocol name")
}
}
pub fn num_connected_peers(&self) -> usize {
self.peers.len()
}
pub fn set_notification_handshake(&mut self, protocol: ProtocolName, handshake: Vec<u8>) {
if let Some(index) = self.notification_protocols.iter().position(|p| *p == protocol) {
self.behaviour.set_notif_protocol_handshake(SetId::from(index), handshake);
} else {
error!(
target: "sub-libp2p",
"set_notification_handshake with unknown protocol: {}",
protocol
);
}
}
}
#[derive(Debug)]
#[must_use]
pub enum CustomMessageOutcome {
NotificationStreamOpened {
remote: PeerId,
protocol: ProtocolName,
negotiated_fallback: Option<ProtocolName>,
roles: Roles,
received_handshake: Vec<u8>,
notifications_sink: NotificationsSink,
},
NotificationStreamReplaced {
remote: PeerId,
protocol: ProtocolName,
notifications_sink: NotificationsSink,
},
NotificationStreamClosed { remote: PeerId, protocol: ProtocolName },
NotificationsReceived { remote: PeerId, messages: Vec<(ProtocolName, Bytes)> },
None,
}
impl<B: BlockT> NetworkBehaviour for Protocol<B> {
type ConnectionHandler = <Notifications as NetworkBehaviour>::ConnectionHandler;
type OutEvent = CustomMessageOutcome;
fn handle_established_inbound_connection(
&mut self,
connection_id: ConnectionId,
peer: PeerId,
local_addr: &Multiaddr,
remote_addr: &Multiaddr,
) -> Result<THandler<Self>, ConnectionDenied> {
self.behaviour.handle_established_inbound_connection(
connection_id,
peer,
local_addr,
remote_addr,
)
}
fn handle_established_outbound_connection(
&mut self,
connection_id: ConnectionId,
peer: PeerId,
addr: &Multiaddr,
role_override: Endpoint,
) -> Result<THandler<Self>, ConnectionDenied> {
self.behaviour.handle_established_outbound_connection(
connection_id,
peer,
addr,
role_override,
)
}
fn handle_pending_outbound_connection(
&mut self,
_connection_id: ConnectionId,
_maybe_peer: Option<PeerId>,
_addresses: &[Multiaddr],
_effective_role: Endpoint,
) -> Result<Vec<Multiaddr>, ConnectionDenied> {
Ok(Vec::new())
}
fn on_swarm_event(&mut self, event: FromSwarm<Self::ConnectionHandler>) {
self.behaviour.on_swarm_event(event);
}
fn on_connection_handler_event(
&mut self,
peer_id: PeerId,
connection_id: ConnectionId,
event: THandlerOutEvent<Self>,
) {
self.behaviour.on_connection_handler_event(peer_id, connection_id, event);
}
fn poll(
&mut self,
cx: &mut std::task::Context,
params: &mut impl PollParameters,
) -> Poll<ToSwarm<Self::OutEvent, THandlerInEvent<Self>>> {
while let Poll::Ready(Some(validation_result)) =
self.sync_substream_validations.poll_next_unpin(cx)
{
match validation_result {
Ok((peer, roles)) => {
self.peers.insert(peer, roles);
},
Err(peer) => {
log::debug!(
target: "sub-libp2p",
"`SyncingEngine` rejected stream"
);
self.behaviour.disconnect_peer(&peer, HARDCODED_PEERSETS_SYNC);
},
}
}
let event = match self.behaviour.poll(cx, params) {
Poll::Pending => return Poll::Pending,
Poll::Ready(ToSwarm::GenerateEvent(ev)) => ev,
Poll::Ready(ToSwarm::Dial { opts }) => return Poll::Ready(ToSwarm::Dial { opts }),
Poll::Ready(ToSwarm::NotifyHandler { peer_id, handler, event }) =>
return Poll::Ready(ToSwarm::NotifyHandler { peer_id, handler, event }),
Poll::Ready(ToSwarm::ReportObservedAddr { address, score }) =>
return Poll::Ready(ToSwarm::ReportObservedAddr { address, score }),
Poll::Ready(ToSwarm::CloseConnection { peer_id, connection }) =>
return Poll::Ready(ToSwarm::CloseConnection { peer_id, connection }),
};
let outcome = match event {
NotificationsOut::CustomProtocolOpen {
peer_id,
set_id,
received_handshake,
notifications_sink,
negotiated_fallback,
inbound,
} => {
if set_id == HARDCODED_PEERSETS_SYNC {
match <Message<B> as DecodeAll>::decode_all(&mut &received_handshake[..]) {
Ok(GenericMessage::Status(handshake)) => {
let roles = handshake.roles;
let handshake = BlockAnnouncesHandshake::<B> {
roles: handshake.roles,
best_number: handshake.best_number,
best_hash: handshake.best_hash,
genesis_hash: handshake.genesis_hash,
};
let (tx, rx) = oneshot::channel();
let _ = self.tx.unbounded_send(
crate::SyncEvent::NotificationStreamOpened {
inbound,
remote: peer_id,
received_handshake: handshake,
sink: notifications_sink,
tx,
},
);
self.sync_substream_validations.push(Box::pin(async move {
match rx.await {
Ok(accepted) =>
if accepted {
Ok((peer_id, roles))
} else {
Err(peer_id)
},
Err(_) => Err(peer_id),
}
}));
CustomMessageOutcome::None
},
Ok(msg) => {
debug!(
target: "sync",
"Expected Status message from {}, but got {:?}",
peer_id,
msg,
);
self.peer_store_handle.report_peer(peer_id, rep::BAD_MESSAGE);
CustomMessageOutcome::None
},
Err(err) => {
match <BlockAnnouncesHandshake<B> as DecodeAll>::decode_all(
&mut &received_handshake[..],
) {
Ok(handshake) => {
let roles = handshake.roles;
let (tx, rx) = oneshot::channel();
let _ = self.tx.unbounded_send(
crate::SyncEvent::NotificationStreamOpened {
inbound,
remote: peer_id,
received_handshake: handshake,
sink: notifications_sink,
tx,
},
);
self.sync_substream_validations.push(Box::pin(async move {
match rx.await {
Ok(accepted) =>
if accepted {
Ok((peer_id, roles))
} else {
Err(peer_id)
},
Err(_) => Err(peer_id),
}
}));
CustomMessageOutcome::None
},
Err(err2) => {
log::debug!(
target: "sync",
"Couldn't decode handshake sent by {}: {:?}: {} & {}",
peer_id,
received_handshake,
err,
err2,
);
self.peer_store_handle.report_peer(peer_id, rep::BAD_MESSAGE);
CustomMessageOutcome::None
},
}
},
}
} else {
match (
Roles::decode_all(&mut &received_handshake[..]),
self.peers.get(&peer_id),
) {
(Ok(roles), _) => CustomMessageOutcome::NotificationStreamOpened {
remote: peer_id,
protocol: self.notification_protocols[usize::from(set_id)].clone(),
negotiated_fallback,
roles,
received_handshake,
notifications_sink,
},
(Err(_), Some(roles)) if received_handshake.is_empty() => {
CustomMessageOutcome::NotificationStreamOpened {
remote: peer_id,
protocol: self.notification_protocols[usize::from(set_id)].clone(),
negotiated_fallback,
roles: *roles,
received_handshake,
notifications_sink,
}
},
(Err(err), _) => {
debug!(target: "sync", "Failed to parse remote handshake: {}", err);
self.bad_handshake_substreams.insert((peer_id, set_id));
self.behaviour.disconnect_peer(&peer_id, set_id);
self.peer_store_handle.report_peer(peer_id, rep::BAD_MESSAGE);
CustomMessageOutcome::None
},
}
}
},
NotificationsOut::CustomProtocolReplaced { peer_id, notifications_sink, set_id } =>
if self.bad_handshake_substreams.contains(&(peer_id, set_id)) {
CustomMessageOutcome::None
} else if set_id == HARDCODED_PEERSETS_SYNC {
let _ = self.tx.unbounded_send(crate::SyncEvent::NotificationSinkReplaced {
remote: peer_id,
sink: notifications_sink,
});
CustomMessageOutcome::None
} else {
CustomMessageOutcome::NotificationStreamReplaced {
remote: peer_id,
protocol: self.notification_protocols[usize::from(set_id)].clone(),
notifications_sink,
}
},
NotificationsOut::CustomProtocolClosed { peer_id, set_id } => {
if self.bad_handshake_substreams.remove(&(peer_id, set_id)) {
CustomMessageOutcome::None
} else if set_id == HARDCODED_PEERSETS_SYNC {
let _ = self.tx.unbounded_send(crate::SyncEvent::NotificationStreamClosed {
remote: peer_id,
});
self.peers.remove(&peer_id);
CustomMessageOutcome::None
} else {
CustomMessageOutcome::NotificationStreamClosed {
remote: peer_id,
protocol: self.notification_protocols[usize::from(set_id)].clone(),
}
}
},
NotificationsOut::Notification { peer_id, set_id, message } => {
if self.bad_handshake_substreams.contains(&(peer_id, set_id)) {
CustomMessageOutcome::None
} else if set_id == HARDCODED_PEERSETS_SYNC {
let _ = self.tx.unbounded_send(crate::SyncEvent::NotificationsReceived {
remote: peer_id,
messages: vec![message.freeze()],
});
CustomMessageOutcome::None
} else {
let protocol_name = self.notification_protocols[usize::from(set_id)].clone();
CustomMessageOutcome::NotificationsReceived {
remote: peer_id,
messages: vec![(protocol_name, message.freeze())],
}
}
},
};
if !matches!(outcome, CustomMessageOutcome::None) {
return Poll::Ready(ToSwarm::GenerateEvent(outcome))
}
cx.waker().wake_by_ref();
Poll::Pending
}
}