use std::{
collections::{HashMap, HashSet, VecDeque, hash_map::Entry},
num::NonZeroUsize,
sync::Arc,
task::{Context, Poll},
time::Duration,
};
use libp2p_core::{
ConnectedPoint, Endpoint, Multiaddr,
multiaddr::{self, Protocol},
transport::PortUse,
};
use libp2p_identity::{Keypair, PeerId, PublicKey};
use libp2p_swarm::{
_address_translation, ConnectionDenied, ConnectionId, DialError, ExternalAddresses,
ListenAddresses, NetworkBehaviour, NotifyHandler, PeerAddresses, StreamUpgradeError, THandler,
THandlerInEvent, THandlerOutEvent, ToSwarm,
behaviour::{ConnectionClosed, ConnectionEstablished, DialFailure, FromSwarm},
};
use crate::{
handler::{self, Handler, InEvent},
protocol::{Info, UpgradeError},
};
fn is_quic_addr(addr: &Multiaddr, v1: bool) -> bool {
use Protocol::*;
let mut iter = addr.iter();
let Some(first) = iter.next() else {
return false;
};
let Some(second) = iter.next() else {
return false;
};
let Some(third) = iter.next() else {
return false;
};
let fourth = iter.next();
let fifth = iter.next();
matches!(first, Ip4(_) | Ip6(_) | Dns(_) | Dns4(_) | Dns6(_))
&& matches!(second, Udp(_))
&& if v1 {
matches!(third, QuicV1)
} else {
matches!(third, Quic)
}
&& matches!(fourth, Some(P2p(_)) | None)
&& fifth.is_none()
}
fn is_tcp_addr(addr: &Multiaddr) -> bool {
use Protocol::*;
let mut iter = addr.iter();
let Some(first) = iter.next() else {
return false;
};
let Some(second) = iter.next() else {
return false;
};
matches!(first, Ip4(_) | Ip6(_) | Dns(_) | Dns4(_) | Dns6(_)) && matches!(second, Tcp(_))
}
fn extract_port(addr: &Multiaddr) -> Option<Protocol<'_>> {
addr.iter()
.find(|p| matches!(p, Protocol::Tcp(_) | Protocol::Udp(_)))
}
fn observed_port_matches_listening_port(
observed: &Multiaddr,
listen_addresses: &ListenAddresses,
) -> bool {
let Some(observed_port) = extract_port(observed) else {
return false;
};
listen_addresses
.iter()
.any(|addr| extract_port(addr).as_ref() == Some(&observed_port))
}
pub struct Behaviour {
config: Config,
connected: HashMap<PeerId, HashMap<ConnectionId, Multiaddr>>,
our_observed_addresses: HashMap<ConnectionId, Multiaddr>,
events: VecDeque<ToSwarm<Event, InEvent>>,
discovered_peers: PeerCache,
listen_addresses: ListenAddresses,
external_addresses: ExternalAddresses,
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct Config {
protocol_version: String,
local_key: Arc<KeyType>,
agent_version: String,
interval: Duration,
push_listen_addr_updates: bool,
cache_size: usize,
hide_listen_addrs: bool,
}
impl Config {
pub fn new(protocol_version: String, local_public_key: PublicKey) -> Self {
Self::new_with_key(protocol_version, local_public_key)
}
pub fn new_with_signed_peer_record(protocol_version: String, local_keypair: &Keypair) -> Self {
Self::new_with_key(protocol_version, local_keypair)
}
fn new_with_key(protocol_version: String, key: impl Into<KeyType>) -> Self {
Self {
protocol_version,
agent_version: format!("rust-libp2p/{}", env!("CARGO_PKG_VERSION")),
local_key: Arc::new(key.into()),
interval: Duration::from_secs(5 * 60),
push_listen_addr_updates: false,
cache_size: 100,
hide_listen_addrs: false,
}
}
pub fn with_agent_version(mut self, v: String) -> Self {
self.agent_version = v;
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
}
pub fn with_hide_listen_addrs(mut self, b: bool) -> Self {
self.hide_listen_addrs = b;
self
}
pub fn protocol_version(&self) -> &str {
&self.protocol_version
}
pub fn local_public_key(&self) -> &PublicKey {
self.local_key.public_key()
}
pub fn agent_version(&self) -> &str {
&self.agent_version
}
pub fn interval(&self) -> Duration {
self.interval
}
pub fn push_listen_addr_updates(&self) -> bool {
self.push_listen_addr_updates
}
pub fn cache_size(&self) -> usize {
self.cache_size
}
pub fn hide_listen_addrs(&self) -> bool {
self.hide_listen_addrs
}
}
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(),
our_observed_addresses: Default::default(),
events: VecDeque::new(),
discovered_peers,
listen_addresses: Default::default(),
external_addresses: Default::default(),
}
}
pub fn push<I>(&mut self, peers: I)
where
I: IntoIterator<Item = PeerId>,
{
for p in peers {
if !self.connected.contains_key(&p) {
tracing::debug!(peer=%p, "Not pushing to peer because we are not connected");
continue;
}
self.events.push_back(ToSwarm::NotifyHandler {
peer_id: p,
handler: NotifyHandler::Any,
event: InEvent::Push,
});
}
}
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(cache) = self.discovered_peers.0.as_mut() {
for addr in failed_addresses {
cache.remove(&peer_id, addr);
}
}
}
fn all_addresses(&self) -> HashSet<Multiaddr> {
let mut addrs = HashSet::from_iter(self.external_addresses.iter().cloned());
if !self.config.hide_listen_addrs {
addrs.extend(self.listen_addresses.iter().cloned());
};
addrs
}
fn emit_new_external_addr_candidate_event(
&mut self,
_connection_id: ConnectionId,
observed: &Multiaddr,
) {
if observed_port_matches_listening_port(observed, &self.listen_addresses) {
if !self.external_addresses.iter().any(|a| a == observed) {
self.events
.push_back(ToSwarm::NewExternalAddrCandidate(observed.clone()));
}
} else {
let translated_addresses = {
let mut addrs: Vec<_> = self
.listen_addresses
.iter()
.filter_map(|server| {
if (is_tcp_addr(server) && is_tcp_addr(observed))
|| (is_quic_addr(server, true) && is_quic_addr(observed, true))
|| (is_quic_addr(server, false) && is_quic_addr(observed, false))
{
_address_translation(server, observed)
} else {
None
}
})
.collect();
addrs.sort_unstable();
addrs.dedup();
addrs
};
if translated_addresses.is_empty() {
self.events
.push_back(ToSwarm::NewExternalAddrCandidate(observed.clone()));
} else {
for addr in translated_addresses {
self.events
.push_back(ToSwarm::NewExternalAddrCandidate(addr));
}
}
}
}
}
impl NetworkBehaviour for Behaviour {
type ConnectionHandler = Handler;
type ToSwarm = Event;
fn handle_established_inbound_connection(
&mut self,
_: ConnectionId,
peer: PeerId,
_: &Multiaddr,
remote_addr: &Multiaddr,
) -> Result<THandler<Self>, ConnectionDenied> {
Ok(Handler::new(
self.config.interval,
peer,
self.config.local_key.clone(),
self.config.protocol_version.clone(),
self.config.agent_version.clone(),
remote_addr.clone(),
self.all_addresses(),
))
}
fn handle_established_outbound_connection(
&mut self,
_connection_id: ConnectionId,
peer: PeerId,
addr: &Multiaddr,
_: Endpoint,
_port_use: PortUse,
) -> Result<THandler<Self>, ConnectionDenied> {
let mut addr = addr.clone();
if matches!(addr.iter().last(), Some(multiaddr::Protocol::P2p(_))) {
addr.pop();
}
Ok(Handler::new(
self.config.interval,
peer,
self.config.local_key.clone(),
self.config.protocol_version.clone(),
self.config.agent_version.clone(),
addr.clone(),
self.all_addresses(),
))
}
fn on_connection_handler_event(
&mut self,
peer_id: PeerId,
connection_id: ConnectionId,
event: THandlerOutEvent<Self>,
) {
match event {
handler::Event::Identified(mut info) => {
info.listen_addrs
.retain(|addr| multiaddr_matches_peer_id(addr, &peer_id));
let observed = info.observed_addr.clone();
self.events
.push_back(ToSwarm::GenerateEvent(Event::Received {
connection_id,
peer_id,
info: info.clone(),
}));
if let Some(ref mut discovered_peers) = self.discovered_peers.0 {
for address in &info.listen_addrs {
if discovered_peers.add(peer_id, address.clone()) {
self.events.push_back(ToSwarm::NewExternalAddrOfPeer {
peer_id,
address: address.clone(),
});
}
}
}
match self.our_observed_addresses.entry(connection_id) {
Entry::Vacant(not_yet_observed) => {
not_yet_observed.insert(observed.clone());
self.emit_new_external_addr_candidate_event(connection_id, &observed);
}
Entry::Occupied(already_observed) if already_observed.get() == &observed => {
}
Entry::Occupied(mut already_observed) => {
tracing::info!(
old_address=%already_observed.get(),
new_address=%observed,
"Our observed address on connection {connection_id} changed",
);
*already_observed.get_mut() = observed.clone();
self.emit_new_external_addr_candidate_event(connection_id, &observed);
}
}
}
handler::Event::Identification => {
self.events.push_back(ToSwarm::GenerateEvent(Event::Sent {
connection_id,
peer_id,
}));
}
handler::Event::IdentificationPushed(info) => {
self.events.push_back(ToSwarm::GenerateEvent(Event::Pushed {
connection_id,
peer_id,
info,
}));
}
handler::Event::IdentificationError(error) => {
self.events.push_back(ToSwarm::GenerateEvent(Event::Error {
connection_id,
peer_id,
error,
}));
}
}
}
#[tracing::instrument(level = "trace", name = "NetworkBehaviour::poll", skip(self))]
fn poll(&mut self, _: &mut Context<'_>) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
if let Some(event) = self.events.pop_front() {
return Poll::Ready(event);
}
Poll::Pending
}
fn handle_pending_outbound_connection(
&mut self,
_connection_id: ConnectionId,
maybe_peer: Option<PeerId>,
_addresses: &[Multiaddr],
_effective_role: Endpoint,
) -> Result<Vec<Multiaddr>, ConnectionDenied> {
let Some(peer) = maybe_peer else {
return Ok(vec![]);
};
Ok(self.discovered_peers.get(&peer))
}
fn on_swarm_event(&mut self, event: FromSwarm) {
let listen_addr_changed = self.listen_addresses.on_swarm_event(&event);
let external_addr_changed = self.external_addresses.on_swarm_event(&event);
if listen_addr_changed || external_addr_changed {
let change_events = self
.connected
.iter()
.flat_map(|(peer, map)| map.keys().map(|id| (*peer, id)))
.map(|(peer_id, connection_id)| ToSwarm::NotifyHandler {
peer_id,
handler: NotifyHandler::One(*connection_id),
event: InEvent::AddressesChanged(self.all_addresses()),
})
.collect::<Vec<_>>();
self.events.extend(change_events)
}
if listen_addr_changed && self.config.push_listen_addr_updates {
let push_events = self.connected.keys().map(|peer| ToSwarm::NotifyHandler {
peer_id: *peer,
handler: NotifyHandler::Any,
event: InEvent::Push,
});
self.events.extend(push_events);
}
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);
} else if let Some(addrs) = self.connected.get_mut(&peer_id) {
addrs.remove(&connection_id);
}
self.our_observed_addresses.remove(&connection_id);
}
FromSwarm::DialFailure(DialFailure {
peer_id: Some(peer_id),
error,
..
}) => {
if let Some(cache) = self.discovered_peers.0.as_mut() {
match error {
DialError::Transport(errors) => {
for (addr, _error) in errors {
cache.remove(&peer_id, addr);
}
}
DialError::WrongPeerId { address, .. }
| DialError::LocalPeerId { address } => {
cache.remove(&peer_id, address);
}
_ => (),
};
}
}
_ => {}
}
}
}
#[allow(clippy::large_enum_variant)]
#[derive(Debug)]
pub enum Event {
Received {
connection_id: ConnectionId,
peer_id: PeerId,
info: Info,
},
Sent {
connection_id: ConnectionId,
peer_id: PeerId,
},
Pushed {
connection_id: ConnectionId,
peer_id: PeerId,
info: Info,
},
Error {
connection_id: ConnectionId,
peer_id: PeerId,
error: StreamUpgradeError<UpgradeError>,
},
}
impl Event {
pub fn connection_id(&self) -> ConnectionId {
match self {
Event::Received { connection_id, .. }
| Event::Sent { connection_id, .. }
| Event::Pushed { connection_id, .. }
| Event::Error { connection_id, .. } => *connection_id,
}
}
}
fn multiaddr_matches_peer_id(addr: &Multiaddr, peer_id: &PeerId) -> bool {
let last_component = addr.iter().last();
if let Some(multiaddr::Protocol::P2p(multi_addr_peer_id)) = last_component {
return multi_addr_peer_id == *peer_id;
}
true
}
struct PeerCache(Option<PeerAddresses>);
impl PeerCache {
fn disabled() -> Self {
Self(None)
}
fn enabled(size: NonZeroUsize) -> Self {
Self(Some(PeerAddresses::new(size)))
}
fn get(&mut self, peer: &PeerId) -> Vec<Multiaddr> {
if let Some(cache) = self.0.as_mut() {
cache.get(peer).collect()
} else {
Vec::new()
}
}
}
#[derive(Debug, Clone)]
#[allow(clippy::large_enum_variant)]
pub(crate) enum KeyType {
PublicKey(PublicKey),
Keypair {
keypair: Keypair,
public_key: PublicKey,
},
}
impl From<PublicKey> for KeyType {
fn from(value: PublicKey) -> Self {
Self::PublicKey(value)
}
}
impl From<&Keypair> for KeyType {
fn from(value: &Keypair) -> Self {
Self::Keypair {
public_key: value.public(),
keypair: value.clone(),
}
}
}
impl KeyType {
pub(crate) fn public_key(&self) -> &PublicKey {
match &self {
KeyType::PublicKey(pubkey) => pubkey,
KeyType::Keypair { public_key, .. } => public_key,
}
}
}
#[cfg(test)]
mod tests {
use libp2p_swarm::behaviour::ListenAddresses;
use super::*;
#[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(multiaddr::Protocol::P2p(peer_id));
addr_with_other_peer_id.push(multiaddr::Protocol::P2p(other_peer_id));
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));
}
#[test]
fn test_extract_port() {
let tcp_addr: Multiaddr = "/ip4/127.0.0.1/tcp/8080".parse().unwrap();
assert_eq!(extract_port(&tcp_addr), Some(Protocol::Tcp(8080)));
let udp_addr: Multiaddr = "/ip4/127.0.0.1/udp/9090".parse().unwrap();
assert_eq!(extract_port(&udp_addr), Some(Protocol::Udp(9090)));
let no_port_addr: Multiaddr = "/ip4/127.0.0.1".parse().unwrap();
assert_eq!(extract_port(&no_port_addr), None);
}
#[test]
fn test_observed_port_matches_listening_port() {
use libp2p_swarm::behaviour::FromSwarm;
let mut listen_addresses = ListenAddresses::default();
let listen_addr: Multiaddr = "/ip4/0.0.0.0/tcp/8080".parse().unwrap();
listen_addresses.on_swarm_event(&FromSwarm::NewListenAddr(
libp2p_swarm::behaviour::NewListenAddr {
listener_id: libp2p_core::transport::ListenerId::next(),
addr: &listen_addr,
},
));
let observed_match: Multiaddr = "/ip4/192.168.1.100/tcp/8080".parse().unwrap();
assert!(observed_port_matches_listening_port(
&observed_match,
&listen_addresses
));
let observed_no_match: Multiaddr = "/ip4/192.168.1.100/tcp/8888".parse().unwrap();
assert!(!observed_port_matches_listening_port(
&observed_no_match,
&listen_addresses
));
let empty_listen_addresses = ListenAddresses::default();
assert!(!observed_port_matches_listening_port(
&observed_match,
&empty_listen_addresses
));
}
#[test]
fn test_address_translation_when_port_matches() {
use libp2p_identity::Keypair;
use libp2p_swarm::behaviour::FromSwarm;
let keypair = Keypair::generate_ed25519();
let config = Config::new("test/1.0.0".to_string(), keypair.public());
let mut behaviour = Behaviour::new(config);
let listen_addr: Multiaddr = "/ip4/0.0.0.0/tcp/8080".parse().unwrap();
behaviour
.listen_addresses
.on_swarm_event(&FromSwarm::NewListenAddr(
libp2p_swarm::behaviour::NewListenAddr {
listener_id: libp2p_core::transport::ListenerId::next(),
addr: &listen_addr,
},
));
behaviour.events.clear();
let observed_matching: Multiaddr = "/ip4/203.0.113.1/tcp/8080".parse().unwrap();
behaviour.emit_new_external_addr_candidate_event(
libp2p_swarm::ConnectionId::new_unchecked(1),
&observed_matching,
);
assert_eq!(behaviour.events.len(), 1);
if let ToSwarm::NewExternalAddrCandidate(addr) = &behaviour.events[0] {
assert_eq!(addr, &observed_matching);
} else {
panic!("Expected NewExternalAddrCandidate event");
}
}
#[test]
fn test_address_translation_when_port_differs() {
use libp2p_identity::Keypair;
use libp2p_swarm::behaviour::FromSwarm;
let keypair = Keypair::generate_ed25519();
let config = Config::new("test/1.0.0".to_string(), keypair.public());
let mut behaviour = Behaviour::new(config);
let listen_addr: Multiaddr = "/ip4/0.0.0.0/tcp/8080".parse().unwrap();
behaviour
.listen_addresses
.on_swarm_event(&FromSwarm::NewListenAddr(
libp2p_swarm::behaviour::NewListenAddr {
listener_id: libp2p_core::transport::ListenerId::next(),
addr: &listen_addr,
},
));
behaviour.events.clear();
let observed_different: Multiaddr = "/ip4/203.0.113.1/tcp/54321".parse().unwrap();
behaviour.emit_new_external_addr_candidate_event(
libp2p_swarm::ConnectionId::new_unchecked(1),
&observed_different,
);
assert!(!behaviour.events.is_empty());
let mut found_translated = false;
for event in &behaviour.events {
if let ToSwarm::NewExternalAddrCandidate(addr) = event {
if addr.to_string().contains("203.0.113.1") && addr.to_string().contains("tcp/8080")
{
found_translated = true;
break;
}
}
}
assert!(
found_translated,
"Should have found a translated address with listening port"
);
}
#[test]
fn test_no_listening_addresses() {
use libp2p_identity::Keypair;
let keypair = Keypair::generate_ed25519();
let config = Config::new("test/1.0.0".to_string(), keypair.public());
let mut behaviour = Behaviour::new(config);
behaviour.events.clear();
let observed: Multiaddr = "/ip4/203.0.113.1/tcp/54321".parse().unwrap();
behaviour.emit_new_external_addr_candidate_event(
libp2p_swarm::ConnectionId::new_unchecked(1),
&observed,
);
assert_eq!(behaviour.events.len(), 1);
if let ToSwarm::NewExternalAddrCandidate(addr) = &behaviour.events[0] {
assert_eq!(addr, &observed);
} else {
panic!("Expected NewExternalAddrCandidate event");
}
}
}