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

//! Protocol handshakes and library synchronization between peers.

use std::{net::SocketAddr, sync::Arc};

use lanspread_proto::{Hello, HelloAck, PROTOCOL_VERSION};
use tokio::sync::{RwLock, mpsc::UnboundedSender};

use crate::{
    PeerEvent,
    context::{Ctx, PeerCtx},
    events,
    identity::default_features,
    library::{LocalLibraryState, build_library_snapshot},
    network::exchange_hello,
    peer_db::{PeerGameDB, PeerId, PeerUpsert},
};

#[derive(Clone)]
pub(crate) struct HandshakeCtx {
    peer_id: Arc<String>,
    local_peer_addr: Arc<RwLock<Option<SocketAddr>>>,
    local_library: Arc<RwLock<LocalLibraryState>>,
    peer_game_db: Arc<RwLock<PeerGameDB>>,
    tx_notify_ui: UnboundedSender<PeerEvent>,
}

impl HandshakeCtx {
    pub(crate) fn from_ctx(ctx: &Ctx, tx_notify_ui: &UnboundedSender<PeerEvent>) -> Self {
        Self {
            peer_id: ctx.peer_id.clone(),
            local_peer_addr: ctx.local_peer_addr.clone(),
            local_library: ctx.local_library.clone(),
            peer_game_db: ctx.peer_game_db.clone(),
            tx_notify_ui: tx_notify_ui.clone(),
        }
    }

    pub(crate) fn from_peer_ctx(ctx: &PeerCtx) -> Self {
        Self {
            peer_id: ctx.peer_id.clone(),
            local_peer_addr: ctx.local_peer_addr.clone(),
            local_library: ctx.local_library.clone(),
            peer_game_db: ctx.peer_game_db.clone(),
            tx_notify_ui: ctx.tx_notify_ui.clone(),
        }
    }
}

async fn required_listen_addr(
    local_peer_addr: &Arc<RwLock<Option<SocketAddr>>>,
) -> eyre::Result<SocketAddr> {
    (*local_peer_addr.read().await)
        .ok_or_else(|| eyre::eyre!("local peer listener address is not ready"))
}

pub(super) async fn build_hello_ack(ctx: &PeerCtx) -> eyre::Result<HelloAck> {
    let library_guard = ctx.local_library.read().await;
    let listen_addr = required_listen_addr(&ctx.local_peer_addr).await?;
    let library = build_library_snapshot(&library_guard);
    Ok(HelloAck {
        peer_id: ctx.peer_id.as_ref().clone(),
        proto_ver: PROTOCOL_VERSION,
        listen_addr,
        library,
        features: default_features(),
    })
}

async fn build_hello_from_state(ctx: &HandshakeCtx) -> eyre::Result<Hello> {
    let library_guard = ctx.local_library.read().await;
    let listen_addr = required_listen_addr(&ctx.local_peer_addr).await?;
    let library = build_library_snapshot(&library_guard);
    Ok(Hello {
        peer_id: ctx.peer_id.as_ref().clone(),
        proto_ver: PROTOCOL_VERSION,
        listen_addr,
        library,
        features: default_features(),
    })
}

pub(crate) async fn perform_handshake_with_peer(
    ctx: HandshakeCtx,
    peer_addr: SocketAddr,
    peer_id_hint: Option<PeerId>,
) -> eyre::Result<()> {
    let hello = build_hello_from_state(&ctx).await?;
    let ack = exchange_hello(peer_addr, hello).await?;

    if ack.proto_ver != PROTOCOL_VERSION {
        log::warn!(
            "Peer {peer_addr} uses incompatible protocol {} (expected {PROTOCOL_VERSION})",
            ack.proto_ver
        );
        return Ok(());
    }

    if ack.peer_id == *ctx.peer_id {
        log::trace!("Ignoring handshake with self for {peer_addr}");
        return Ok(());
    }

    if let Some(expected) = peer_id_hint.as_ref()
        && expected != &ack.peer_id
    {
        log::warn!(
            "Peer {peer_addr} id mismatch: mDNS advertised {expected}, hello ack returned {}",
            ack.peer_id
        );
        let _ = ctx.peer_game_db.write().await.remove_peer(expected);
    }

    let record_addr = ack.listen_addr;
    let upsert = record_remote_library(
        &ctx.peer_game_db,
        ack.peer_id.clone(),
        record_addr,
        ack.features.clone(),
        ack.library,
    )
    .await;

    after_peer_library_recorded(&ctx, upsert, record_addr).await;
    events::emit_peer_game_list(&ctx.peer_game_db, &ctx.tx_notify_ui).await;

    Ok(())
}

pub(super) async fn accept_inbound_hello(
    ctx: &PeerCtx,
    transport_addr: Option<SocketAddr>,
    hello: Hello,
) -> eyre::Result<HelloAck> {
    if hello.peer_id == *ctx.peer_id {
        log::trace!("Ignoring hello from self");
        return build_hello_ack(ctx).await;
    }

    if hello.proto_ver != PROTOCOL_VERSION {
        log::warn!(
            "Incompatible protocol from {transport_addr:?}: {}",
            hello.proto_ver
        );
        return build_hello_ack(ctx).await;
    }

    let addr = hello.listen_addr;
    let handshake_ctx = HandshakeCtx::from_peer_ctx(ctx);
    let upsert = record_remote_library(
        &ctx.peer_game_db,
        hello.peer_id.clone(),
        addr,
        hello.features.clone(),
        hello.library,
    )
    .await;

    after_peer_library_recorded(&handshake_ctx, upsert, addr).await;
    events::emit_peer_game_list(&ctx.peer_game_db, &ctx.tx_notify_ui).await;

    build_hello_ack(ctx).await
}

pub(super) fn spawn_library_resync(
    ctx: HandshakeCtx,
    peer_addr: SocketAddr,
    peer_id_hint: PeerId,
    reason: &'static str,
) {
    tokio::spawn(async move {
        if let Err(err) = perform_handshake_with_peer(ctx, peer_addr, Some(peer_id_hint)).await {
            log::warn!("Failed to {reason} library from {peer_addr}: {err}");
        }
    });
}

async fn record_remote_library(
    peer_game_db: &Arc<RwLock<PeerGameDB>>,
    peer_id: PeerId,
    peer_addr: SocketAddr,
    features: Vec<String>,
    snapshot: lanspread_proto::LibrarySnapshot,
) -> PeerUpsert {
    let mut db = peer_game_db.write().await;
    let upsert = db.upsert_peer(peer_id.clone(), peer_addr);
    db.apply_library_snapshot(&peer_id, snapshot);
    db.update_peer_features(&peer_id, features);
    upsert
}

async fn after_peer_library_recorded(
    ctx: &HandshakeCtx,
    upsert: PeerUpsert,
    peer_addr: SocketAddr,
) {
    if upsert.is_new {
        events::emit_peer_discovered(&ctx.peer_game_db, &ctx.tx_notify_ui, peer_addr).await;
    }
}

#[cfg(test)]
mod tests {
    use std::{
        collections::{HashMap, HashSet},
        net::SocketAddr,
        path::{Path, PathBuf},
        sync::Arc,
    };

    use lanspread_proto::{Availability, GameSummary, Hello, LibrarySnapshot, PROTOCOL_VERSION};
    use tokio::sync::{RwLock, mpsc};
    use tokio_util::{sync::CancellationToken, task::TaskTracker};

    use super::{HandshakeCtx, accept_inbound_hello, build_hello_from_state};
    use crate::{
        PeerEvent,
        UnpackFuture,
        Unpacker,
        context::Ctx,
        library::LocalLibraryState,
        peer_db::PeerGameDB,
    };

    struct NoopUnpacker;

    impl Unpacker for NoopUnpacker {
        fn unpack<'a>(&'a self, _archive: &'a Path, _dest: &'a Path) -> UnpackFuture<'a> {
            Box::pin(async { Ok(()) })
        }
    }

    fn addr(ip: [u8; 4], port: u16) -> SocketAddr {
        SocketAddr::from((ip, port))
    }

    fn test_handshake_ctx(local_peer_addr: Option<SocketAddr>) -> HandshakeCtx {
        let (tx_notify_ui, _rx_notify_ui) = mpsc::unbounded_channel();
        let peer_game_db = Arc::new(RwLock::new(PeerGameDB::new()));
        HandshakeCtx {
            peer_id: Arc::new("local-peer".to_string()),
            local_peer_addr: Arc::new(RwLock::new(local_peer_addr)),
            local_library: Arc::new(RwLock::new(LocalLibraryState::empty())),
            peer_game_db,
            tx_notify_ui,
        }
    }

    fn summary(id: &str) -> GameSummary {
        GameSummary {
            id: id.to_string(),
            name: id.to_string(),
            size: 42,
            downloaded: true,
            installed: true,
            eti_version: Some("20250101".to_string()),
            manifest_hash: 7,
            availability: Availability::Ready,
        }
    }

    #[tokio::test]
    async fn outbound_hello_requires_local_listener_addr() {
        let ctx = test_handshake_ctx(None);

        let err = build_hello_from_state(&ctx)
            .await
            .expect_err("hello without listener must fail");

        assert_eq!(err.to_string(), "local peer listener address is not ready");
    }

    #[tokio::test]
    async fn outbound_hello_carries_local_listener_addr() {
        let advertised = addr([10, 66, 0, 2], 40000);
        let ctx = test_handshake_ctx(Some(advertised));

        let hello = build_hello_from_state(&ctx)
            .await
            .expect("listener address is present");

        assert_eq!(hello.listen_addr, advertised);
    }

    #[tokio::test]
    async fn outbound_hello_carries_local_library_snapshot() {
        let ctx = test_handshake_ctx(Some(addr([10, 66, 0, 2], 40000)));
        ctx.local_library
            .write()
            .await
            .update_from_scan(HashMap::from([("game".to_string(), summary("game"))]), 7);

        let hello = build_hello_from_state(&ctx)
            .await
            .expect("listener address is present");

        assert_eq!(hello.library.library_rev, 7);
        assert_eq!(hello.library.games.len(), 1);
        assert_eq!(hello.library.games[0].id, "game");
    }

    #[tokio::test]
    async fn inbound_hello_applies_remote_library_snapshot() {
        let peer_game_db = Arc::new(RwLock::new(PeerGameDB::new()));
        let ctx = Ctx::new(
            peer_game_db.clone(),
            "local-peer".to_string(),
            PathBuf::new(),
            PathBuf::new(),
            Arc::new(NoopUnpacker),
            CancellationToken::new(),
            TaskTracker::new(),
            Arc::new(RwLock::new(HashSet::new())),
        );
        *ctx.local_peer_addr.write().await = Some(addr([127, 0, 0, 1], 4000));

        let (tx_notify_ui, mut rx_notify_ui) = mpsc::unbounded_channel();
        let peer_ctx = ctx.to_peer_ctx(tx_notify_ui);
        let remote_addr = addr([127, 0, 0, 1], 5000);
        let hello = Hello {
            peer_id: "remote-peer".to_string(),
            proto_ver: PROTOCOL_VERSION,
            listen_addr: remote_addr,
            library: LibrarySnapshot {
                library_rev: 3,
                games: vec![summary("remote-game")],
            },
            features: Vec::new(),
        };

        let ack = accept_inbound_hello(&peer_ctx, None, hello)
            .await
            .expect("current protocol hello should be accepted");

        assert_eq!(ack.peer_id, "local-peer");
        let snapshots = peer_game_db.read().await.peer_snapshots();
        assert_eq!(snapshots.len(), 1);
        assert_eq!(snapshots[0].addr, remote_addr);
        assert_eq!(snapshots[0].game_count, 1);
        assert_eq!(snapshots[0].games[0].id, "remote-game");

        assert!(matches!(
            rx_notify_ui
                .recv()
                .await
                .expect("peer discovery event should be emitted"),
            PeerEvent::PeerDiscovered(addr) if addr == remote_addr
        ));
        assert!(matches!(
            rx_notify_ui
                .recv()
                .await
                .expect("peer count event should be emitted"),
            PeerEvent::PeerCountUpdated(1)
        ));
        let PeerEvent::ListGames(games) = rx_notify_ui
            .recv()
            .await
            .expect("peer game list should be emitted")
        else {
            panic!("expected ListGames");
        };
        assert_eq!(games.len(), 1);
        assert_eq!(games[0].id, "remote-game");
        assert_eq!(games[0].peer_count, 1);
    }
}