use crate::dns::{build_query, build_query_response, build_service_discovery_response};
use crate::query::MdnsPacket;
use async_io::{Async, Timer};
use futures::prelude::*;
use if_watch::{IfEvent, IfWatcher};
use lazy_static::lazy_static;
use libp2p_core::{
address_translation, connection::ConnectionId, multiaddr::Protocol, Multiaddr, PeerId,
};
use libp2p_swarm::{
protocols_handler::DummyProtocolsHandler, NetworkBehaviour, NetworkBehaviourAction,
PollParameters, ProtocolsHandler,
};
use smallvec::SmallVec;
use socket2::{Domain, Socket, Type};
use std::{
cmp,
collections::VecDeque,
fmt, io, iter,
net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket},
pin::Pin,
task::Context,
task::Poll,
time::{Duration, Instant},
};
lazy_static! {
static ref IPV4_MDNS_MULTICAST_ADDRESS: SocketAddr =
SocketAddr::from((Ipv4Addr::new(224, 0, 0, 251), 5353));
}
#[derive(Clone, Debug)]
pub struct MdnsConfig {
pub ttl: Duration,
pub query_interval: Duration,
}
impl Default for MdnsConfig {
fn default() -> Self {
Self {
ttl: Duration::from_secs(6 * 60),
query_interval: Duration::from_secs(5 * 60),
}
}
}
#[derive(Debug)]
pub struct Mdns {
recv_socket: Async<UdpSocket>,
send_socket: Async<UdpSocket>,
if_watch: IfWatcher,
recv_buffer: [u8; 4096],
send_buffer: VecDeque<Vec<u8>>,
discovered_nodes: SmallVec<[(PeerId, Multiaddr, Instant); 8]>,
closest_expiration: Option<Timer>,
events: VecDeque<MdnsEvent>,
query_interval: Duration,
ttl: Duration,
timeout: Timer,
}
impl Mdns {
pub async fn new(config: MdnsConfig) -> io::Result<Self> {
let recv_socket = {
let socket = Socket::new(
Domain::IPV4,
Type::DGRAM,
Some(socket2::Protocol::UDP),
)?;
socket.set_reuse_address(true)?;
#[cfg(unix)]
socket.set_reuse_port(true)?;
socket.bind(&SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 5353).into())?;
let socket = UdpSocket::from(socket);
socket.set_multicast_loop_v4(true)?;
socket.set_multicast_ttl_v4(255)?;
Async::new(socket)?
};
let send_socket = {
let socket = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0))?;
Async::new(socket)?
};
let if_watch = if_watch::IfWatcher::new().await?;
Ok(Self {
recv_socket,
send_socket,
if_watch,
recv_buffer: [0; 4096],
send_buffer: Default::default(),
discovered_nodes: SmallVec::new(),
closest_expiration: None,
events: Default::default(),
query_interval: config.query_interval,
ttl: config.ttl,
timeout: Timer::interval(config.query_interval),
})
}
pub fn has_node(&self, peer_id: &PeerId) -> bool {
self.discovered_nodes().any(|p| p == peer_id)
}
pub fn discovered_nodes(&self) -> impl ExactSizeIterator<Item = &PeerId> {
self.discovered_nodes.iter().map(|(p, _, _)| p)
}
fn inject_mdns_packet(&mut self, packet: MdnsPacket, params: &impl PollParameters) {
self.timeout.set_interval(self.query_interval);
match packet {
MdnsPacket::Query(query) => {
for packet in build_query_response(
query.query_id(),
*params.local_peer_id(),
params.listened_addresses(),
self.ttl,
) {
self.send_buffer.push_back(packet);
}
}
MdnsPacket::Response(response) => {
let obs_ip = Protocol::from(response.remote_addr().ip());
let obs_port = Protocol::Udp(response.remote_addr().port());
let observed: Multiaddr = iter::once(obs_ip).chain(iter::once(obs_port)).collect();
let mut discovered: SmallVec<[_; 4]> = SmallVec::new();
for peer in response.discovered_peers() {
if peer.id() == params.local_peer_id() {
continue;
}
let new_expiration = Instant::now() + peer.ttl();
let mut addrs: Vec<Multiaddr> = Vec::new();
for addr in peer.addresses() {
if let Some(new_addr) = address_translation(&addr, &observed) {
addrs.push(new_addr.clone())
}
addrs.push(addr.clone())
}
for addr in addrs {
if let Some((_, _, cur_expires)) = self
.discovered_nodes
.iter_mut()
.find(|(p, a, _)| p == peer.id() && *a == addr)
{
*cur_expires = cmp::max(*cur_expires, new_expiration);
} else {
self.discovered_nodes
.push((*peer.id(), addr.clone(), new_expiration));
}
discovered.push((*peer.id(), addr));
}
}
self.closest_expiration = self
.discovered_nodes
.iter()
.fold(None, |exp, &(_, _, elem_exp)| {
Some(exp.map(|exp| cmp::min(exp, elem_exp)).unwrap_or(elem_exp))
})
.map(Timer::at);
self.events
.push_back(MdnsEvent::Discovered(DiscoveredAddrsIter {
inner: discovered.into_iter(),
}));
}
MdnsPacket::ServiceDiscovery(disc) => {
let resp = build_service_discovery_response(disc.query_id(), self.ttl);
self.send_buffer.push_back(resp);
}
}
}
}
impl NetworkBehaviour for Mdns {
type ProtocolsHandler = DummyProtocolsHandler;
type OutEvent = MdnsEvent;
fn new_handler(&mut self) -> Self::ProtocolsHandler {
DummyProtocolsHandler::default()
}
fn addresses_of_peer(&mut self, peer_id: &PeerId) -> Vec<Multiaddr> {
let now = Instant::now();
self.discovered_nodes
.iter()
.filter(move |(p, _, expires)| p == peer_id && *expires > now)
.map(|(_, addr, _)| addr.clone())
.collect()
}
fn inject_connected(&mut self, _: &PeerId) {}
fn inject_disconnected(&mut self, _: &PeerId) {}
fn inject_event(
&mut self,
_: PeerId,
_: ConnectionId,
ev: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
) {
void::unreachable(ev)
}
fn poll(
&mut self,
cx: &mut Context<'_>,
params: &mut impl PollParameters,
) -> Poll<
NetworkBehaviourAction<
<Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
Self::OutEvent,
>,
> {
while let Poll::Ready(event) = Pin::new(&mut self.if_watch).poll(cx) {
let multicast = From::from([224, 0, 0, 251]);
let socket = self.recv_socket.get_ref();
match event {
Ok(IfEvent::Up(inet)) => {
if inet.addr().is_loopback() {
continue;
}
if let IpAddr::V4(addr) = inet.addr() {
log::trace!("joining multicast on iface {}", addr);
if let Err(err) = socket.join_multicast_v4(&multicast, &addr) {
log::error!("join multicast failed: {}", err);
} else {
self.send_buffer.push_back(build_query());
}
}
}
Ok(IfEvent::Down(inet)) => {
if inet.addr().is_loopback() {
continue;
}
if let IpAddr::V4(addr) = inet.addr() {
log::trace!("leaving multicast on iface {}", addr);
if let Err(err) = socket.leave_multicast_v4(&multicast, &addr) {
log::error!("leave multicast failed: {}", err);
}
}
}
Err(err) => log::error!("if watch returned an error: {}", err),
}
}
while self.recv_socket.poll_readable(cx).is_ready() {
match self
.recv_socket
.recv_from(&mut self.recv_buffer)
.now_or_never()
{
Some(Ok((len, from))) => {
if let Some(packet) = MdnsPacket::new_from_bytes(&self.recv_buffer[..len], from)
{
self.inject_mdns_packet(packet, params);
}
}
Some(Err(err)) => log::error!("Failed reading datagram: {}", err),
_ => {}
}
}
if Pin::new(&mut self.timeout).poll_next(cx).is_ready() {
self.send_buffer.push_back(build_query());
}
if !self.send_buffer.is_empty() {
while self.send_socket.poll_writable(cx).is_ready() {
if let Some(packet) = self.send_buffer.pop_front() {
match self
.send_socket
.send_to(&packet, *IPV4_MDNS_MULTICAST_ADDRESS)
.now_or_never()
{
Some(Ok(_)) => {}
Some(Err(err)) => log::error!("{}", err),
None => self.send_buffer.push_front(packet),
}
} else {
break;
}
}
}
if let Some(event) = self.events.pop_front() {
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event));
}
if let Some(ref mut closest_expiration) = self.closest_expiration {
if let Poll::Ready(now) = Pin::new(closest_expiration).poll(cx) {
let mut expired = SmallVec::<[(PeerId, Multiaddr); 4]>::new();
while let Some(pos) = self
.discovered_nodes
.iter()
.position(|(_, _, exp)| *exp < now)
{
let (peer_id, addr, _) = self.discovered_nodes.remove(pos);
expired.push((peer_id, addr));
}
if !expired.is_empty() {
let event = MdnsEvent::Expired(ExpiredAddrsIter {
inner: expired.into_iter(),
});
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event));
}
}
}
Poll::Pending
}
}
#[derive(Debug)]
pub enum MdnsEvent {
Discovered(DiscoveredAddrsIter),
Expired(ExpiredAddrsIter),
}
pub struct DiscoveredAddrsIter {
inner: smallvec::IntoIter<[(PeerId, Multiaddr); 4]>,
}
impl Iterator for DiscoveredAddrsIter {
type Item = (PeerId, Multiaddr);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl ExactSizeIterator for DiscoveredAddrsIter {}
impl fmt::Debug for DiscoveredAddrsIter {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("DiscoveredAddrsIter").finish()
}
}
pub struct ExpiredAddrsIter {
inner: smallvec::IntoIter<[(PeerId, Multiaddr); 4]>,
}
impl Iterator for ExpiredAddrsIter {
type Item = (PeerId, Multiaddr);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl ExactSizeIterator for ExpiredAddrsIter {}
impl fmt::Debug for ExpiredAddrsIter {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("ExpiredAddrsIter").finish()
}
}