From 65eef2b579610bf0e594befcbf2c3f492a0a11cc Mon Sep 17 00:00:00 2001
From: open-trade <zhouhuab@163.com>
Date: Tue, 11 Jan 2022 16:24:35 +0800
Subject: [PATCH] raw lan discovery

---
 libs/hbb_common/protos/rendezvous.proto | 11 ++++++
 libs/hbb_common/src/lib.rs              |  1 +
 src/rendezvous_mediator.rs              | 47 +++++++++++++++++++++++--
 3 files changed, 57 insertions(+), 2 deletions(-)

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/src/rendezvous_mediator.rs b/src/rendezvous_mediator.rs
index d46344517..a94cd092a 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,43 @@ async fn direct_server(server: ServerPtr) -> ResultType<()> {
         }
     }
 }
+
+pub fn create_multicast_socket() -> ResultType<FramedSocket> {
+    let port = (RENDEZVOUS_PORT + 3) as u16;
+    udp::bind_multicast(
+        &SocketAddrV4::new([0, 0, 0, 0].into(), port),
+        &SocketAddrV4::new([239, 255, 42, 98].into(), port),
+    )
+}
+
+async fn lan_discovery() -> ResultType<()> {
+    let mut socket = 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 peer = PeerDiscovery {
+                                    cmd: "pong".to_owned,
+                                    mac: hbb_common::mac_address::get_mac_address()?,
+                                    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).await?;
+                            }
+                        }
+                        _ => {}
+                    }
+                }
+            }
+        }
+    }
+    Ok(())
+}