diff --git a/libs/hbb_common/protos/rendezvous.proto b/libs/hbb_common/protos/rendezvous.proto index 8508971fa..de81ed0f9 100644 --- a/libs/hbb_common/protos/rendezvous.proto +++ b/libs/hbb_common/protos/rendezvous.proto @@ -133,6 +133,16 @@ message LocalAddr { string version = 5; } +message PeerDiscovery { + string cmd = 1; + string mac = 2; + string id = 3; + string username = 4; + string hostname = 5; + string platform = 6; + string misc = 7; +} + message RendezvousMessage { oneof union { RegisterPeer register_peer = 6; @@ -151,5 +161,6 @@ message RendezvousMessage { RelayResponse relay_response = 19; TestNatRequest test_nat_request = 20; TestNatResponse test_nat_response = 21; + PeerDiscovery peer_discovery = 22; } } diff --git a/libs/hbb_common/src/lib.rs b/libs/hbb_common/src/lib.rs index f84221b70..10975e573 100644 --- a/libs/hbb_common/src/lib.rs +++ b/libs/hbb_common/src/lib.rs @@ -32,6 +32,7 @@ pub use sodiumoxide; pub use tokio_socks; pub use tokio_socks::IntoTargetAddr; pub use tokio_socks::TargetAddr; +pub use mac_address; #[cfg(feature = "quic")] pub type Stream = quic::Connection; diff --git a/libs/hbb_common/src/udp.rs b/libs/hbb_common/src/udp.rs index d1896bdbc..876180475 100644 --- a/libs/hbb_common/src/udp.rs +++ b/libs/hbb_common/src/udp.rs @@ -3,8 +3,8 @@ use anyhow::anyhow; use bytes::{Bytes, BytesMut}; use futures::{SinkExt, StreamExt}; use protobuf::Message; -use socket2::{Domain, Socket, Type}; -use std::net::SocketAddr; +use socket2::{Domain, Protocol, Socket, Type}; +use std::net::{SocketAddr, SocketAddrV4}; use tokio::net::{ToSocketAddrs, UdpSocket}; use tokio_socks::{udp::Socks5UdpFramed, IntoTargetAddr, TargetAddr, ToProxyAddrs}; use tokio_util::{codec::BytesCodec, udp::UdpFramed}; @@ -142,3 +142,18 @@ impl FramedSocket { } } } + +// const DEFAULT_MULTICAST: &str = "239.255.42.98"; + +pub fn bind_multicast(addr: &SocketAddrV4, multi_addr: &SocketAddrV4) -> ResultType { + assert!(multi_addr.ip().is_multicast(), "Must be multcast address"); + let socket = Socket::new(Domain::ipv4(), Type::dgram(), Some(Protocol::udp()))?; + socket.set_reuse_address(true)?; + socket.bind(&socket2::SockAddr::from(*addr))?; + socket.set_multicast_loop_v4(true)?; + socket.join_multicast_v4(multi_addr.ip(), addr.ip())?; + Ok(FramedSocket::Direct(UdpFramed::new( + UdpSocket::from_std(socket.into_udp_socket())?, + BytesCodec::new(), + ))) +} diff --git a/src/rendezvous_mediator.rs b/src/rendezvous_mediator.rs index d46344517..64638c15a 100644 --- a/src/rendezvous_mediator.rs +++ b/src/rendezvous_mediator.rs @@ -12,11 +12,11 @@ use hbb_common::{ self, select, time::{interval, Duration}, }, - udp::FramedSocket, + udp::{self, FramedSocket}, AddrMangle, IntoTargetAddr, ResultType, TargetAddr, }; use std::{ - net::SocketAddr, + net::{SocketAddr, SocketAddrV4}, sync::{ atomic::{AtomicBool, Ordering}, Arc, Mutex, @@ -59,6 +59,9 @@ impl RendezvousMediator { tokio::spawn(async move { allow_err!(direct_server(server_cloned).await); }); + tokio::spawn(async move { + allow_err!(lan_discovery().await); + }); loop { Config::reset_online(); if Config::get_option("stop-service").is_empty() { @@ -501,3 +504,48 @@ async fn direct_server(server: ServerPtr) -> ResultType<()> { } } } + +pub fn create_multicast_socket() -> ResultType<(FramedSocket, SocketAddr)> { + let port = (RENDEZVOUS_PORT + 3) as u16; + let maddr = SocketAddrV4::new([239, 255, 42, 98].into(), port); + Ok(( + udp::bind_multicast(&SocketAddrV4::new([0, 0, 0, 0].into(), port), &maddr)?, + SocketAddr::V4(maddr), + )) +} + +async fn lan_discovery() -> ResultType<()> { + let (mut socket, maddr) = create_multicast_socket()?; + loop { + select! { + Some(Ok((bytes, _))) = socket.next() => { + if let Ok(msg_in) = Message::parse_from_bytes(&bytes) { + match msg_in.union { + Some(rendezvous_message::Union::peer_discovery(p)) => { + if p.cmd == "ping" { + let mut msg_out = Message::new(); + let mac = if let Ok(Some(mac)) = mac_address::get_mac_address() { + mac.to_string() + } else { + "".to_owned() + }; + let peer = PeerDiscovery { + cmd: "pong".to_owned(), + mac, + id: Config::get_id(), + hostname: whoami::hostname(), + username: crate::platform::get_active_username(), + platform: whoami::platform().to_string(), + ..Default::default() + }; + msg_out.set_peer_discovery(peer); + socket.send(&msg_out, maddr).await?; + } + } + _ => {} + } + } + } + } + } +}