use crate::handler::{self, Proto, Push};
use crate::protocol::{Info, ReplySubstream, UpgradeError};
use futures::prelude::*;
use libp2p_core::{
connection::ConnectionId, multiaddr::Protocol, ConnectedPoint, Multiaddr, PeerId, PublicKey,
};
use libp2p_swarm::behaviour::{ConnectionClosed, ConnectionEstablished, DialFailure, FromSwarm};
use libp2p_swarm::{
dial_opts::DialOpts, AddressScore, ConnectionHandler, ConnectionHandlerUpgrErr, DialError,
IntoConnectionHandler, NegotiatedSubstream, NetworkBehaviour, NetworkBehaviourAction,
NotifyHandler, PollParameters,
};
use lru::LruCache;
use std::num::NonZeroUsize;
use std::{
collections::{HashMap, HashSet, VecDeque},
iter::FromIterator,
pin::Pin,
task::Context,
task::Poll,
time::Duration,
};
pub struct Behaviour {
config: Config,
connected: HashMap<PeerId, HashMap<ConnectionId, Multiaddr>>,
pending_replies: VecDeque<Reply>,
events: VecDeque<NetworkBehaviourAction<Event, Proto>>,
pending_push: HashSet<PeerId>,
discovered_peers: PeerCache,
}
enum Reply {
Queued {
peer: PeerId,
io: ReplySubstream<NegotiatedSubstream>,
observed: Multiaddr,
},
Sending {
peer: PeerId,
io: Pin<Box<dyn Future<Output = Result<(), UpgradeError>> + Send>>,
},
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct Config {
pub protocol_version: String,
pub local_public_key: PublicKey,
pub agent_version: String,
pub initial_delay: Duration,
pub interval: Duration,
pub push_listen_addr_updates: bool,
pub cache_size: usize,
}
impl Config {
pub fn new(protocol_version: String, local_public_key: PublicKey) -> Self {
Self {
protocol_version,
agent_version: format!("rust-libp2p/{}", env!("CARGO_PKG_VERSION")),
local_public_key,
initial_delay: Duration::from_millis(500),
interval: Duration::from_secs(5 * 60),
push_listen_addr_updates: false,
cache_size: 100,
}
}
pub fn with_agent_version(mut self, v: String) -> Self {
self.agent_version = v;
self
}
pub fn with_initial_delay(mut self, d: Duration) -> Self {
self.initial_delay = d;
self
}
pub fn with_interval(mut self, d: Duration) -> Self {
self.interval = d;
self
}
pub fn with_push_listen_addr_updates(mut self, b: bool) -> Self {
self.push_listen_addr_updates = b;
self
}
pub fn with_cache_size(mut self, cache_size: usize) -> Self {
self.cache_size = cache_size;
self
}
}
impl Behaviour {
pub fn new(config: Config) -> Self {
let discovered_peers = match NonZeroUsize::new(config.cache_size) {
None => PeerCache::disabled(),
Some(size) => PeerCache::enabled(size),
};
Self {
config,
connected: HashMap::new(),
pending_replies: VecDeque::new(),
events: VecDeque::new(),
pending_push: HashSet::new(),
discovered_peers,
}
}
pub fn push<I>(&mut self, peers: I)
where
I: IntoIterator<Item = PeerId>,
{
for p in peers {
if self.pending_push.insert(p) && !self.connected.contains_key(&p) {
let handler = self.new_handler();
self.events.push_back(NetworkBehaviourAction::Dial {
opts: DialOpts::peer_id(p).build(),
handler,
});
}
}
}
fn on_connection_established(
&mut self,
ConnectionEstablished {
peer_id,
connection_id: conn,
endpoint,
failed_addresses,
..
}: ConnectionEstablished,
) {
let addr = match endpoint {
ConnectedPoint::Dialer { address, .. } => address.clone(),
ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr.clone(),
};
self.connected
.entry(peer_id)
.or_default()
.insert(conn, addr);
if let Some(entry) = self.discovered_peers.get_mut(&peer_id) {
for addr in failed_addresses {
entry.remove(addr);
}
}
}
}
impl NetworkBehaviour for Behaviour {
type ConnectionHandler = Proto;
type OutEvent = Event;
fn new_handler(&mut self) -> Self::ConnectionHandler {
Proto::new(self.config.initial_delay, self.config.interval)
}
fn on_connection_handler_event(
&mut self,
peer_id: PeerId,
connection: ConnectionId,
event: <<Self::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::OutEvent,
) {
match event {
handler::Event::Identified(mut info) => {
info.listen_addrs
.retain(|addr| multiaddr_matches_peer_id(addr, &peer_id));
self.discovered_peers
.put(peer_id, info.listen_addrs.iter().cloned());
let observed = info.observed_addr.clone();
self.events
.push_back(NetworkBehaviourAction::GenerateEvent(Event::Received {
peer_id,
info,
}));
self.events
.push_back(NetworkBehaviourAction::ReportObservedAddr {
address: observed,
score: AddressScore::Finite(1),
});
}
handler::Event::IdentificationPushed => {
self.events
.push_back(NetworkBehaviourAction::GenerateEvent(Event::Pushed {
peer_id,
}));
}
handler::Event::Identify(sender) => {
let observed = self
.connected
.get(&peer_id)
.and_then(|addrs| addrs.get(&connection))
.expect(
"`on_connection_handler_event` is only called \
with an established connection and calling `NetworkBehaviour::on_event` \
with `FromSwarm::ConnectionEstablished ensures there is an entry; qed",
);
self.pending_replies.push_back(Reply::Queued {
peer: peer_id,
io: sender,
observed: observed.clone(),
});
}
handler::Event::IdentificationError(error) => {
self.events
.push_back(NetworkBehaviourAction::GenerateEvent(Event::Error {
peer_id,
error,
}));
}
}
}
fn poll(
&mut self,
cx: &mut Context<'_>,
params: &mut impl PollParameters,
) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ConnectionHandler>> {
if let Some(event) = self.events.pop_front() {
return Poll::Ready(event);
}
let peer_push = self.pending_push.iter().find_map(|peer| {
self.connected.get(peer).map(|conns| {
let observed_addr = conns
.values()
.next()
.expect("connected peer has a connection")
.clone();
let listen_addrs = listen_addrs(params);
let protocols = supported_protocols(params);
let info = Info {
public_key: self.config.local_public_key.clone(),
protocol_version: self.config.protocol_version.clone(),
agent_version: self.config.agent_version.clone(),
listen_addrs,
protocols,
observed_addr,
};
(*peer, Push(info))
})
});
if let Some((peer_id, push)) = peer_push {
self.pending_push.remove(&peer_id);
return Poll::Ready(NetworkBehaviourAction::NotifyHandler {
peer_id,
event: push,
handler: NotifyHandler::Any,
});
}
if let Some(r) = self.pending_replies.pop_front() {
let mut sending = 0;
let to_send = self.pending_replies.len() + 1;
let mut reply = Some(r);
loop {
match reply {
Some(Reply::Queued { peer, io, observed }) => {
let info = Info {
listen_addrs: listen_addrs(params),
protocols: supported_protocols(params),
public_key: self.config.local_public_key.clone(),
protocol_version: self.config.protocol_version.clone(),
agent_version: self.config.agent_version.clone(),
observed_addr: observed,
};
let io = Box::pin(io.send(info));
reply = Some(Reply::Sending { peer, io });
}
Some(Reply::Sending { peer, mut io }) => {
sending += 1;
match Future::poll(Pin::new(&mut io), cx) {
Poll::Ready(Ok(())) => {
let event = Event::Sent { peer_id: peer };
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event));
}
Poll::Pending => {
self.pending_replies.push_back(Reply::Sending { peer, io });
if sending == to_send {
break;
} else {
reply = self.pending_replies.pop_front();
}
}
Poll::Ready(Err(err)) => {
let event = Event::Error {
peer_id: peer,
error: ConnectionHandlerUpgrErr::Upgrade(
libp2p_core::upgrade::UpgradeError::Apply(err),
),
};
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event));
}
}
}
None => unreachable!(),
}
}
}
Poll::Pending
}
fn addresses_of_peer(&mut self, peer: &PeerId) -> Vec<Multiaddr> {
self.discovered_peers.get(peer)
}
fn on_swarm_event(&mut self, event: FromSwarm<Self::ConnectionHandler>) {
match event {
FromSwarm::ConnectionEstablished(connection_established) => {
self.on_connection_established(connection_established)
}
FromSwarm::ConnectionClosed(ConnectionClosed {
peer_id,
connection_id,
remaining_established,
..
}) => {
if remaining_established == 0 {
self.connected.remove(&peer_id);
self.pending_push.remove(&peer_id);
} else if let Some(addrs) = self.connected.get_mut(&peer_id) {
addrs.remove(&connection_id);
}
}
FromSwarm::DialFailure(DialFailure { peer_id, error, .. }) => {
if let Some(peer_id) = peer_id {
if !self.connected.contains_key(&peer_id) {
self.pending_push.remove(&peer_id);
}
}
if let Some(entry) = peer_id.and_then(|id| self.discovered_peers.get_mut(&id)) {
if let DialError::Transport(errors) = error {
for (addr, _error) in errors {
entry.remove(addr);
}
}
}
}
FromSwarm::NewListenAddr(_) => {
if self.config.push_listen_addr_updates {
self.pending_push.extend(self.connected.keys());
}
}
FromSwarm::ExpiredListenAddr(_) => {
if self.config.push_listen_addr_updates {
self.pending_push.extend(self.connected.keys());
}
}
FromSwarm::AddressChange(_)
| FromSwarm::ListenFailure(_)
| FromSwarm::NewListener(_)
| FromSwarm::ListenerError(_)
| FromSwarm::ListenerClosed(_)
| FromSwarm::NewExternalAddr(_)
| FromSwarm::ExpiredExternalAddr(_) => {}
}
}
}
#[allow(clippy::large_enum_variant)]
#[derive(Debug)]
pub enum Event {
Received {
peer_id: PeerId,
info: Info,
},
Sent {
peer_id: PeerId,
},
Pushed {
peer_id: PeerId,
},
Error {
peer_id: PeerId,
error: ConnectionHandlerUpgrErr<UpgradeError>,
},
}
fn supported_protocols(params: &impl PollParameters) -> Vec<String> {
params
.supported_protocols()
.map(|p| String::from_utf8_lossy(&p).to_string())
.collect()
}
fn listen_addrs(params: &impl PollParameters) -> Vec<Multiaddr> {
let mut listen_addrs: Vec<_> = params.external_addresses().map(|r| r.addr).collect();
listen_addrs.extend(params.listened_addresses());
listen_addrs
}
fn multiaddr_matches_peer_id(addr: &Multiaddr, peer_id: &PeerId) -> bool {
let last_component = addr.iter().last();
if let Some(Protocol::P2p(multi_addr_peer_id)) = last_component {
return multi_addr_peer_id == *peer_id.as_ref();
}
true
}
struct PeerCache(Option<LruCache<PeerId, HashSet<Multiaddr>>>);
impl PeerCache {
fn disabled() -> Self {
Self(None)
}
fn enabled(size: NonZeroUsize) -> Self {
Self(Some(LruCache::new(size)))
}
fn get_mut(&mut self, peer: &PeerId) -> Option<&mut HashSet<Multiaddr>> {
self.0.as_mut()?.get_mut(peer)
}
fn put(&mut self, peer: PeerId, addresses: impl Iterator<Item = Multiaddr>) {
let cache = match self.0.as_mut() {
None => return,
Some(cache) => cache,
};
cache.put(peer, HashSet::from_iter(addresses));
}
fn get(&mut self, peer: &PeerId) -> Vec<Multiaddr> {
let cache = match self.0.as_mut() {
None => return Vec::new(),
Some(cache) => cache,
};
cache
.get(peer)
.cloned()
.map(Vec::from_iter)
.unwrap_or_default()
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::pin_mut;
use libp2p::mplex::MplexConfig;
use libp2p::noise;
use libp2p::tcp;
use libp2p_core::{identity, muxing::StreamMuxerBox, transport, upgrade, PeerId, Transport};
use libp2p_swarm::{Swarm, SwarmEvent};
use std::time::Duration;
fn transport() -> (
identity::PublicKey,
transport::Boxed<(PeerId, StreamMuxerBox)>,
) {
let id_keys = identity::Keypair::generate_ed25519();
let noise_keys = noise::Keypair::<noise::X25519Spec>::new()
.into_authentic(&id_keys)
.unwrap();
let pubkey = id_keys.public();
let transport = tcp::async_io::Transport::new(tcp::Config::default().nodelay(true))
.upgrade(upgrade::Version::V1)
.authenticate(noise::NoiseConfig::xx(noise_keys).into_authenticated())
.multiplex(MplexConfig::new())
.boxed();
(pubkey, transport)
}
#[test]
fn periodic_identify() {
let (mut swarm1, pubkey1) = {
let (pubkey, transport) = transport();
let protocol = Behaviour::new(
Config::new("a".to_string(), pubkey.clone()).with_agent_version("b".to_string()),
);
let swarm = Swarm::with_async_std_executor(transport, protocol, pubkey.to_peer_id());
(swarm, pubkey)
};
let (mut swarm2, pubkey2) = {
let (pubkey, transport) = transport();
let protocol = Behaviour::new(
Config::new("c".to_string(), pubkey.clone()).with_agent_version("d".to_string()),
);
let swarm = Swarm::with_async_std_executor(transport, protocol, pubkey.to_peer_id());
(swarm, pubkey)
};
swarm1
.listen_on("/ip4/127.0.0.1/tcp/0".parse().unwrap())
.unwrap();
let listen_addr = async_std::task::block_on(async {
loop {
let swarm1_fut = swarm1.select_next_some();
pin_mut!(swarm1_fut);
if let SwarmEvent::NewListenAddr { address, .. } = swarm1_fut.await {
return address;
}
}
});
swarm2.dial(listen_addr).unwrap();
async_std::task::block_on(async move {
loop {
let swarm1_fut = swarm1.select_next_some();
pin_mut!(swarm1_fut);
let swarm2_fut = swarm2.select_next_some();
pin_mut!(swarm2_fut);
match future::select(swarm1_fut, swarm2_fut)
.await
.factor_second()
.0
{
future::Either::Left(SwarmEvent::Behaviour(Event::Received {
info, ..
})) => {
assert_eq!(info.public_key, pubkey2);
assert_eq!(info.protocol_version, "c");
assert_eq!(info.agent_version, "d");
assert!(!info.protocols.is_empty());
assert!(info.listen_addrs.is_empty());
return;
}
future::Either::Right(SwarmEvent::Behaviour(Event::Received {
info, ..
})) => {
assert_eq!(info.public_key, pubkey1);
assert_eq!(info.protocol_version, "a");
assert_eq!(info.agent_version, "b");
assert!(!info.protocols.is_empty());
assert_eq!(info.listen_addrs.len(), 1);
return;
}
_ => {}
}
}
})
}
#[test]
fn identify_push() {
let _ = env_logger::try_init();
let (mut swarm1, pubkey1) = {
let (pubkey, transport) = transport();
let protocol = Behaviour::new(Config::new("a".to_string(), pubkey.clone()));
let swarm = Swarm::with_async_std_executor(transport, protocol, pubkey.to_peer_id());
(swarm, pubkey)
};
let (mut swarm2, pubkey2) = {
let (pubkey, transport) = transport();
let protocol = Behaviour::new(
Config::new("a".to_string(), pubkey.clone()).with_agent_version("b".to_string()),
);
let swarm = Swarm::with_async_std_executor(transport, protocol, pubkey.to_peer_id());
(swarm, pubkey)
};
Swarm::listen_on(&mut swarm1, "/ip4/127.0.0.1/tcp/0".parse().unwrap()).unwrap();
let listen_addr = async_std::task::block_on(async {
loop {
let swarm1_fut = swarm1.select_next_some();
pin_mut!(swarm1_fut);
if let SwarmEvent::NewListenAddr { address, .. } = swarm1_fut.await {
return address;
}
}
});
swarm2.dial(listen_addr).unwrap();
async_std::task::block_on(async move {
loop {
let swarm1_fut = swarm1.select_next_some();
let swarm2_fut = swarm2.select_next_some();
{
pin_mut!(swarm1_fut);
pin_mut!(swarm2_fut);
match future::select(swarm1_fut, swarm2_fut)
.await
.factor_second()
.0
{
future::Either::Left(SwarmEvent::Behaviour(Event::Received {
info,
..
})) => {
assert_eq!(info.public_key, pubkey2);
assert_eq!(info.protocol_version, "a");
assert_eq!(info.agent_version, "b");
assert!(!info.protocols.is_empty());
assert!(info.listen_addrs.is_empty());
return;
}
future::Either::Right(SwarmEvent::ConnectionEstablished { .. }) => {
}
_ => continue,
}
}
swarm2
.behaviour_mut()
.push(std::iter::once(pubkey1.to_peer_id()));
}
})
}
#[test]
fn discover_peer_after_disconnect() {
let _ = env_logger::try_init();
let mut swarm1 = {
let (pubkey, transport) = transport();
let protocol = Behaviour::new(
Config::new("a".to_string(), pubkey.clone())
.with_initial_delay(Duration::from_secs(10)),
);
Swarm::with_async_std_executor(transport, protocol, pubkey.to_peer_id())
};
let mut swarm2 = {
let (pubkey, transport) = transport();
let protocol = Behaviour::new(
Config::new("a".to_string(), pubkey.clone()).with_agent_version("b".to_string()),
);
Swarm::with_async_std_executor(transport, protocol, pubkey.to_peer_id())
};
let swarm1_peer_id = *swarm1.local_peer_id();
let listener = swarm1
.listen_on("/ip4/127.0.0.1/tcp/0".parse().unwrap())
.unwrap();
let listen_addr = async_std::task::block_on(async {
loop {
match swarm1.select_next_some().await {
SwarmEvent::NewListenAddr {
address,
listener_id,
} if listener_id == listener => return address,
_ => {}
}
}
});
async_std::task::spawn(async move {
loop {
swarm1.next().await;
}
});
swarm2.dial(listen_addr).unwrap();
async_std::task::block_on(async {
loop {
if let SwarmEvent::Behaviour(Event::Received { .. }) =
swarm2.select_next_some().await
{
break;
}
}
});
swarm2.disconnect_peer_id(swarm1_peer_id).unwrap();
async_std::task::block_on(async {
loop {
if let SwarmEvent::ConnectionClosed { peer_id, .. } =
swarm2.select_next_some().await
{
break peer_id;
}
}
});
swarm2.dial(swarm1_peer_id).unwrap();
let connected_peer = async_std::task::block_on(async {
loop {
if let SwarmEvent::ConnectionEstablished { peer_id, .. } =
swarm2.select_next_some().await
{
break peer_id;
}
}
});
assert_eq!(connected_peer, swarm1_peer_id);
}
#[test]
fn check_multiaddr_matches_peer_id() {
let peer_id = PeerId::random();
let other_peer_id = PeerId::random();
let mut addr: Multiaddr = "/ip4/147.75.69.143/tcp/4001"
.parse()
.expect("failed to parse multiaddr");
let addr_without_peer_id: Multiaddr = addr.clone();
let mut addr_with_other_peer_id = addr.clone();
addr.push(Protocol::P2p(peer_id.into()));
addr_with_other_peer_id.push(Protocol::P2p(other_peer_id.into()));
assert!(multiaddr_matches_peer_id(&addr, &peer_id));
assert!(!multiaddr_matches_peer_id(
&addr_with_other_peer_id,
&peer_id
));
assert!(multiaddr_matches_peer_id(&addr_without_peer_id, &peer_id));
}
}