lanspread/crates/lanspread-peer/src/services/discovery.rs

//! mDNS peer discovery and discovery-time protocol negotiation.

use std::time::Duration;

use lanspread_mdns::{LANSPREAD_SERVICE_TYPE, MdnsBrowser, MdnsService, MdnsServicePoll};
use lanspread_proto::PROTOCOL_VERSION;
use tokio::sync::mpsc::UnboundedSender;

use crate::{
    PeerEvent,
    context::Ctx,
    events,
    peer_db::PeerId,
    services::handshake::{HandshakeCtx, perform_handshake_with_peer},
};

struct MdnsPeerInfo {
    addr: std::net::SocketAddr,
    peer_id: Option<PeerId>,
    proto_ver: Option<u32>,
    library_rev: u64,
    library_digest: u64,
}

/// Runs the peer discovery service using mDNS.
pub async fn run_peer_discovery(
    tx_notify_ui: UnboundedSender<PeerEvent>,
    ctx: Ctx,
) -> eyre::Result<()> {
    log::info!("Starting peer discovery task");

    if !wait_for_local_peer_addr(&ctx).await {
        return Ok(());
    }

    let service_type = LANSPREAD_SERVICE_TYPE.to_string();
    let (service_tx, mut service_rx) = tokio::sync::mpsc::unbounded_channel();
    let worker_shutdown = ctx.shutdown.clone();
    let service_type_clone = service_type.clone();

    let worker_handle = ctx
        .task_tracker
        .spawn_blocking(move || -> eyre::Result<()> {
            let browser = MdnsBrowser::new(&service_type_clone)?;
            while !worker_shutdown.is_cancelled() {
                match browser.next_service_timeout(None, Duration::from_millis(250))? {
                    MdnsServicePoll::Service(service) => {
                        if service_tx.send(service).is_err() {
                            log::debug!("Peer discovery consumer dropped; stopping worker");
                            break;
                        }
                    }
                    MdnsServicePoll::Timeout => {}
                    MdnsServicePoll::Closed => {
                        log::warn!("mDNS browser closed; stopping peer discovery worker");
                        break;
                    }
                }
            }
            Ok(())
        });

    loop {
        tokio::select! {
            () = ctx.shutdown.cancelled() => break,
            service = service_rx.recv() => {
                let Some(service) = service else {
                    break;
                };

                let info = parse_mdns_peer(&service);
                if is_self_advertisement(&info, &ctx).await {
                    log::trace!("Ignoring self advertisement at {}", info.addr);
                    continue;
                }

                handle_discovered_peer(info, &ctx, &tx_notify_ui).await;
            }
        }
    }

    match worker_handle.await {
        Ok(Ok(())) if ctx.shutdown.is_cancelled() => Ok(()),
        Ok(Ok(())) => {
            eyre::bail!("mDNS discovery worker exited unexpectedly");
        }
        Ok(Err(err)) if ctx.shutdown.is_cancelled() => {
            log::debug!("Peer discovery worker stopped during shutdown: {err}");
            Ok(())
        }
        Ok(Err(err)) => Err(err.wrap_err("peer discovery worker failed")),
        Err(err) if ctx.shutdown.is_cancelled() => {
            log::debug!("Peer discovery worker join ended during shutdown: {err}");
            Ok(())
        }
        Err(err) => Err(eyre::eyre!("peer discovery worker join error: {err}")),
    }
}

async fn wait_for_local_peer_addr(ctx: &Ctx) -> bool {
    loop {
        if ctx.local_peer_addr.read().await.is_some() {
            return true;
        }

        tokio::select! {
            () = ctx.shutdown.cancelled() => return false,
            () = tokio::time::sleep(Duration::from_millis(25)) => {}
        }
    }
}

fn parse_mdns_peer(service: &MdnsService) -> MdnsPeerInfo {
    MdnsPeerInfo {
        addr: service.addr,
        peer_id: service.properties.get("peer_id").cloned(),
        proto_ver: service
            .properties
            .get("proto_ver")
            .and_then(|value| value.parse::<u32>().ok()),
        library_rev: service
            .properties
            .get("library_rev")
            .and_then(|value| value.parse::<u64>().ok())
            .unwrap_or(0),
        library_digest: service
            .properties
            .get("library_digest")
            .and_then(|value| value.parse::<u64>().ok())
            .unwrap_or(0),
    }
}

async fn is_self_advertisement(info: &MdnsPeerInfo, ctx: &Ctx) -> bool {
    let guard = ctx.local_peer_addr.read().await;
    guard.as_ref().is_some_and(|addr| *addr == info.addr)
        || info
            .peer_id
            .as_ref()
            .is_some_and(|peer_id| peer_id == ctx.peer_id.as_ref())
}

async fn handle_discovered_peer(
    info: MdnsPeerInfo,
    ctx: &Ctx,
    tx_notify_ui: &UnboundedSender<PeerEvent>,
) {
    if info.proto_ver != Some(PROTOCOL_VERSION) {
        log::debug!(
            "Ignoring peer at {} with protocol {:?}; expected {PROTOCOL_VERSION}",
            info.addr,
            info.proto_ver
        );
        return;
    }

    let Some(peer_id) = info.peer_id.clone() else {
        log::debug!(
            "Ignoring current-protocol peer at {} without a peer_id TXT record",
            info.addr
        );
        return;
    };

    let upsert = {
        let mut db = ctx.peer_game_db.write().await;
        let upsert = db.upsert_peer(peer_id.clone(), info.addr);
        let features = db.peer_features(&peer_id);
        if info.library_rev > 0 || info.library_digest > 0 {
            db.update_peer_library(&peer_id, info.library_rev, info.library_digest, features);
        }
        upsert
    };

    if upsert.is_new {
        log::info!("Discovered peer at: {}", info.addr);
        events::emit_peer_discovered(&ctx.peer_game_db, tx_notify_ui, info.addr).await;
    }

    if upsert.is_new || upsert.addr_changed {
        spawn_protocol_negotiation(&info, ctx, tx_notify_ui, peer_id);
    }
}

fn spawn_protocol_negotiation(
    info: &MdnsPeerInfo,
    ctx: &Ctx,
    tx_notify_ui: &UnboundedSender<PeerEvent>,
    peer_id: PeerId,
) {
    let peer_addr = info.addr;
    let handshake_ctx = HandshakeCtx::from_ctx(ctx, tx_notify_ui);

    ctx.task_tracker.spawn(async move {
        if let Err(err) = perform_handshake_with_peer(handshake_ctx, peer_addr, Some(peer_id)).await
        {
            log::warn!("Failed to negotiate protocol with peer {peer_addr}: {err}");
        }
    });
}