use crate::{
config::{MultiaddrWithPeerId, NonReservedPeerMode, NotificationHandshake, SetConfig},
litep2p::shim::notification::{
peerset::{Peerset, PeersetCommand},
NotificationProtocol,
},
peer_store::PeerStoreProvider,
service::{metrics::NotificationMetrics, traits::NotificationConfig},
NotificationService, ProtocolName,
};
use litep2p::protocol::notification::{Config, ConfigBuilder};
use sc_utils::mpsc::TracingUnboundedSender;
use std::sync::{atomic::AtomicUsize, Arc};
#[derive(Debug, Clone)]
pub struct ProtocolControlHandle {
pub tx: TracingUnboundedSender<PeersetCommand>,
pub connected_peers: Arc<AtomicUsize>,
}
impl ProtocolControlHandle {
pub fn new(
tx: TracingUnboundedSender<PeersetCommand>,
connected_peers: Arc<AtomicUsize>,
) -> Self {
Self { tx, connected_peers }
}
}
#[derive(Debug)]
pub struct NotificationProtocolConfig {
pub protocol_name: ProtocolName,
max_notification_size: usize,
set_config: SetConfig,
pub config: Config,
pub handle: ProtocolControlHandle,
}
impl NotificationProtocolConfig {
pub fn new(
protocol_name: ProtocolName,
fallback_names: Vec<ProtocolName>,
max_notification_size: usize,
handshake: Option<NotificationHandshake>,
set_config: SetConfig,
metrics: NotificationMetrics,
peerstore_handle: Arc<dyn PeerStoreProvider>,
) -> (Self, Box<dyn NotificationService>) {
let connected_peers = Arc::new(Default::default());
let (peerset, peerset_tx) = Peerset::new(
protocol_name.clone(),
set_config.out_peers as usize,
set_config.in_peers as usize,
set_config.non_reserved_mode == NonReservedPeerMode::Deny,
set_config.reserved_nodes.iter().map(|address| address.peer_id).collect(),
Arc::clone(&connected_peers),
peerstore_handle,
);
let (config, handle) = ConfigBuilder::new(protocol_name.clone().into())
.with_handshake(handshake.map_or(vec![1], |handshake| (*handshake).to_vec()))
.with_max_size(max_notification_size as usize)
.with_auto_accept_inbound(true)
.with_fallback_names(fallback_names.into_iter().map(From::from).collect())
.build();
let protocol = NotificationProtocol::new(protocol_name.clone(), handle, peerset, metrics);
(
Self {
protocol_name,
max_notification_size,
set_config,
config,
handle: ProtocolControlHandle::new(peerset_tx, connected_peers),
},
Box::new(protocol),
)
}
pub fn protocol_name(&self) -> &ProtocolName {
&self.protocol_name
}
pub fn set_config(&self) -> &SetConfig {
&self.set_config
}
pub fn allow_non_reserved(&mut self, in_peers: u32, out_peers: u32) {
self.set_config.in_peers = in_peers;
self.set_config.out_peers = out_peers;
self.set_config.non_reserved_mode = NonReservedPeerMode::Accept;
}
pub fn add_reserved(&mut self, peer: MultiaddrWithPeerId) {
self.set_config.reserved_nodes.push(peer);
}
pub fn max_notification_size(&self) -> usize {
self.max_notification_size
}
}
impl NotificationConfig for NotificationProtocolConfig {
fn set_config(&self) -> &SetConfig {
&self.set_config
}
fn protocol_name(&self) -> &ProtocolName {
&self.protocol_name
}
}