use crate::{
config::{self, NonReservedPeerMode},
error,
types::ProtocolName,
};
use bytes::Bytes;
use codec::{DecodeAll, Encode};
use libp2p::{
core::connection::ConnectionId,
swarm::{
behaviour::FromSwarm, ConnectionHandler, IntoConnectionHandler, NetworkBehaviour,
NetworkBehaviourAction, PollParameters,
},
Multiaddr, PeerId,
};
use log::{debug, error, warn};
use sc_network_common::{role::Roles, sync::message::BlockAnnouncesHandshake};
use sp_runtime::traits::Block as BlockT;
use std::{
collections::{HashMap, HashSet, VecDeque},
iter,
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: sc_peerset::SetId = sc_peerset::SetId::from(0);
const NUM_HARDCODED_PEERSETS: usize = 1;
mod rep {
use sc_peerset::ReputationChange as Rep;
pub const BAD_MESSAGE: Rep = Rep::new(-(1 << 12), "Bad message");
}
pub struct Protocol<B: BlockT> {
pending_messages: VecDeque<CustomMessageOutcome>,
peerset_handle: sc_peerset::PeersetHandle,
behaviour: Notifications,
notification_protocols: Vec<ProtocolName>,
bad_handshake_substreams: HashSet<(PeerId, sc_peerset::SetId)>,
peers: HashMap<PeerId, Roles>,
_marker: std::marker::PhantomData<B>,
}
impl<B: BlockT> Protocol<B> {
pub fn new(
roles: Roles,
network_config: &config::NetworkConfiguration,
block_announces_protocol: config::NonDefaultSetConfig,
) -> error::Result<(Self, sc_peerset::PeersetHandle, Vec<(PeerId, Multiaddr)>)> {
let mut known_addresses = Vec::new();
let (peerset, peerset_handle) = {
let mut sets =
Vec::with_capacity(NUM_HARDCODED_PEERSETS + network_config.extra_sets.len());
let mut default_sets_reserved = HashSet::new();
for reserved in network_config.default_peers_set.reserved_nodes.iter() {
default_sets_reserved.insert(reserved.peer_id);
if !reserved.multiaddr.is_empty() {
known_addresses.push((reserved.peer_id, reserved.multiaddr.clone()));
}
}
let mut bootnodes = Vec::with_capacity(network_config.boot_nodes.len());
for bootnode in network_config.boot_nodes.iter() {
bootnodes.push(bootnode.peer_id);
}
sets.push(sc_peerset::SetConfig {
in_peers: network_config.default_peers_set.in_peers,
out_peers: network_config.default_peers_set.out_peers,
bootnodes,
reserved_nodes: default_sets_reserved.clone(),
reserved_only: network_config.default_peers_set.non_reserved_mode ==
NonReservedPeerMode::Deny,
});
for set_cfg in &network_config.extra_sets {
let mut reserved_nodes = HashSet::new();
for reserved in set_cfg.set_config.reserved_nodes.iter() {
reserved_nodes.insert(reserved.peer_id);
known_addresses.push((reserved.peer_id, reserved.multiaddr.clone()));
}
let reserved_only =
set_cfg.set_config.non_reserved_mode == NonReservedPeerMode::Deny;
sets.push(sc_peerset::SetConfig {
in_peers: set_cfg.set_config.in_peers,
out_peers: set_cfg.set_config.out_peers,
bootnodes: Vec::new(),
reserved_nodes,
reserved_only,
});
}
sc_peerset::Peerset::from_config(sc_peerset::PeersetConfig { sets })
};
let behaviour = {
Notifications::new(
peerset,
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(network_config.extra_sets.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 {
pending_messages: VecDeque::new(),
peerset_handle: peerset_handle.clone(),
behaviour,
notification_protocols: iter::once(block_announces_protocol.notifications_protocol)
.chain(network_config.extra_sets.iter().map(|s| s.notifications_protocol.clone()))
.collect(),
bad_handshake_substreams: Default::default(),
peers: HashMap::new(),
_marker: Default::default(),
};
Ok((protocol, peerset_handle, known_addresses))
}
pub fn open_peers(&self) -> impl Iterator<Item = &PeerId> {
self.behaviour.open_peers()
}
pub fn num_discovered_peers(&self) -> usize {
self.behaviour.num_discovered_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, sc_peerset::SetId::from(position));
self.peers.remove(peer_id);
} else {
warn!(target: "sub-libp2p", "disconnect_peer() with invalid protocol name")
}
}
pub fn peerset_debug_info(&mut self) -> serde_json::Value {
self.behaviour.peerset_debug_info()
}
pub fn num_connected_peers(&self) -> usize {
self.peers.len()
}
pub fn report_peer(&self, who: PeerId, reputation: sc_peerset::ReputationChange) {
self.peerset_handle.report_peer(who, reputation)
}
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(sc_peerset::SetId::from(index), handshake);
} else {
error!(
target: "sub-libp2p",
"set_notification_handshake with unknown protocol: {}",
protocol
);
}
}
pub fn set_reserved_only(&self, reserved_only: bool) {
self.peerset_handle.set_reserved_only(HARDCODED_PEERSETS_SYNC, reserved_only);
}
pub fn remove_reserved_peer(&self, peer: PeerId) {
self.peerset_handle.remove_reserved_peer(HARDCODED_PEERSETS_SYNC, peer);
}
pub fn reserved_peers(&self) -> impl Iterator<Item = &PeerId> {
self.behaviour.reserved_peers(HARDCODED_PEERSETS_SYNC)
}
pub fn add_reserved_peer(&self, peer: PeerId) {
self.peerset_handle.add_reserved_peer(HARDCODED_PEERSETS_SYNC, peer);
}
pub fn set_reserved_peers(&self, peers: HashSet<PeerId>) {
self.peerset_handle.set_reserved_peers(HARDCODED_PEERSETS_SYNC, peers);
}
pub fn set_reserved_peerset_peers(&self, protocol: ProtocolName, peers: HashSet<PeerId>) {
if let Some(index) = self.notification_protocols.iter().position(|p| *p == protocol) {
self.peerset_handle.set_reserved_peers(sc_peerset::SetId::from(index), peers);
} else {
error!(
target: "sub-libp2p",
"set_reserved_peerset_peers with unknown protocol: {}",
protocol
);
}
}
pub fn remove_set_reserved_peer(&self, protocol: ProtocolName, peer: PeerId) {
if let Some(index) = self.notification_protocols.iter().position(|p| *p == protocol) {
self.peerset_handle.remove_reserved_peer(sc_peerset::SetId::from(index), peer);
} else {
error!(
target: "sub-libp2p",
"remove_set_reserved_peer with unknown protocol: {}",
protocol
);
}
}
pub fn add_set_reserved_peer(&self, protocol: ProtocolName, peer: PeerId) {
if let Some(index) = self.notification_protocols.iter().position(|p| *p == protocol) {
self.peerset_handle.add_reserved_peer(sc_peerset::SetId::from(index), peer);
} else {
error!(
target: "sub-libp2p",
"add_set_reserved_peer with unknown protocol: {}",
protocol
);
}
}
pub fn add_default_set_discovered_nodes(&mut self, peer_ids: impl Iterator<Item = PeerId>) {
for peer_id in peer_ids {
self.peerset_handle.add_to_peers_set(HARDCODED_PEERSETS_SYNC, peer_id);
}
}
pub fn add_to_peers_set(&self, protocol: ProtocolName, peer: PeerId) {
if let Some(index) = self.notification_protocols.iter().position(|p| *p == protocol) {
self.peerset_handle.add_to_peers_set(sc_peerset::SetId::from(index), peer);
} else {
error!(
target: "sub-libp2p",
"add_to_peers_set with unknown protocol: {}",
protocol
);
}
}
pub fn remove_from_peers_set(&self, protocol: ProtocolName, peer: PeerId) {
if let Some(index) = self.notification_protocols.iter().position(|p| *p == protocol) {
self.peerset_handle.remove_from_peers_set(sc_peerset::SetId::from(index), peer);
} else {
error!(
target: "sub-libp2p",
"remove_from_peers_set 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 new_handler(&mut self) -> Self::ConnectionHandler {
self.behaviour.new_handler()
}
fn addresses_of_peer(&mut self, _: &PeerId) -> Vec<Multiaddr> {
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: <<Self::ConnectionHandler as IntoConnectionHandler>::Handler as
ConnectionHandler>::OutEvent,
) {
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<NetworkBehaviourAction<Self::OutEvent, Self::ConnectionHandler>> {
if let Some(message) = self.pending_messages.pop_front() {
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(message))
}
let event = match self.behaviour.poll(cx, params) {
Poll::Pending => return Poll::Pending,
Poll::Ready(NetworkBehaviourAction::GenerateEvent(ev)) => ev,
Poll::Ready(NetworkBehaviourAction::Dial { opts, handler }) =>
return Poll::Ready(NetworkBehaviourAction::Dial { opts, handler }),
Poll::Ready(NetworkBehaviourAction::NotifyHandler { peer_id, handler, event }) =>
return Poll::Ready(NetworkBehaviourAction::NotifyHandler {
peer_id,
handler,
event,
}),
Poll::Ready(NetworkBehaviourAction::ReportObservedAddr { address, score }) =>
return Poll::Ready(NetworkBehaviourAction::ReportObservedAddr { address, score }),
Poll::Ready(NetworkBehaviourAction::CloseConnection { peer_id, connection }) =>
return Poll::Ready(NetworkBehaviourAction::CloseConnection { peer_id, connection }),
};
let outcome = match event {
NotificationsOut::CustomProtocolOpen {
peer_id,
set_id,
received_handshake,
notifications_sink,
negotiated_fallback,
} => {
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,
};
self.peers.insert(peer_id, roles);
CustomMessageOutcome::NotificationStreamOpened {
remote: peer_id,
protocol: self.notification_protocols[usize::from(set_id)].clone(),
negotiated_fallback,
received_handshake: handshake.encode(),
roles,
notifications_sink,
}
},
Ok(msg) => {
debug!(
target: "sync",
"Expected Status message from {}, but got {:?}",
peer_id,
msg,
);
self.peerset_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;
self.peers.insert(peer_id, roles);
CustomMessageOutcome::NotificationStreamOpened {
remote: peer_id,
protocol: self.notification_protocols[usize::from(set_id)]
.clone(),
negotiated_fallback,
received_handshake,
roles,
notifications_sink,
}
},
Err(err2) => {
log::debug!(
target: "sync",
"Couldn't decode handshake sent by {}: {:?}: {} & {}",
peer_id,
received_handshake,
err,
err2,
);
self.peerset_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.peerset_handle.report_peer(peer_id, rep::BAD_MESSAGE);
self.peers.remove(&peer_id);
CustomMessageOutcome::None
},
}
}
},
NotificationsOut::CustomProtocolReplaced { peer_id, notifications_sink, set_id } =>
if self.bad_handshake_substreams.contains(&(peer_id, set_id)) {
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 {
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 {
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(NetworkBehaviourAction::GenerateEvent(outcome))
}
if let Some(message) = self.pending_messages.pop_front() {
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(message))
}
cx.waker().wake_by_ref();
Poll::Pending
}
}