lan discovery will be done soon

Signed-off-by: fufesou <shuanglongchen@yeah.net>

libs/socket_cs/src/discovery.rs

@@ -1,15 +1,17 @@
 use super::udp::UdpRequest;
 use async_trait::async_trait;
 use hbb_common::{
-    base_proto::PeerInfo, discovery_proto::Discovery as DiscoveryProto, log, protobuf::Message,
-    tokio::net::UdpSocket, ResultType,
+    discovery_proto::{Discovery as DiscoveryProto, DiscoveryBack as DiscoveryBackProto},
+    log,
+    protobuf::Message,
+    tokio::net::UdpSocket,
+    ResultType,
 };
 use std::net::SocketAddr;
 
 pub const CMD_DISCOVERY: &str = "discovery";
 pub const CMD_DISCOVERY_BACK: &str = "discovery_back";
 
-// TODO: make sure if UdpFramed is needed, or UdpSocket just works fine.
 pub struct DiscoveryClient {
     socket: UdpSocket,
     send_data: Vec<u8>,
@@ -44,13 +46,17 @@ impl DiscoveryClient {
 }
 
 pub struct HandlerDiscovery {
+    get_allow: Option<fn() -> bool>,
     send_data: Vec<u8>,
 }
 
 impl HandlerDiscovery {
-    pub fn new(self_info: PeerInfo) -> Self {
+    pub fn new(get_allow: Option<fn() -> bool>, self_info: DiscoveryBackProto) -> Self {
         let send_data = make_send_data(CMD_DISCOVERY_BACK, &self_info).unwrap();
-        Self { send_data }
+        Self {
+            get_allow,
+            send_data,
+        }
     }
 }
 
@@ -67,6 +73,17 @@ impl crate::Handler<UdpRequest> for HandlerDiscovery {
             peer.hostname
         );
 
+        let allowed = self.get_allow.map_or(false, |f| f());
+        if !allowed {
+            log::info!(
+                "received discovery query from {} {} {}, but discovery is not allowed",
+                request.addr,
+                peer.hostname,
+                peer.username
+            );
+            return Ok(());
+        }
+
         let addr = "0.0.0.0:0";
         let socket = UdpSocket::bind(addr).await?;
 
@@ -89,11 +106,11 @@ impl crate::Handler<UdpRequest> for HandlerDiscovery {
 }
 
 pub struct HandlerDiscoveryBack {
-    proc: fn(peer_info: PeerInfo),
+    proc: fn(info: DiscoveryBackProto),
 }
 
 impl HandlerDiscoveryBack {
-    pub fn new(proc: fn(peer_info: PeerInfo)) -> Self {
+    pub fn new(proc: fn(info: DiscoveryBackProto)) -> Self {
         Self { proc }
     }
 }
@@ -103,8 +120,8 @@ impl crate::Handler<UdpRequest> for HandlerDiscoveryBack {
     async fn call(&self, request: UdpRequest) -> ResultType<()> {
         log::trace!("recved discover back from {}", request.addr);
 
-        let peer = PeerInfo::parse_from_bytes(&request.data)?;
-        (self.proc)(peer);
+        let info = DiscoveryBackProto::parse_from_bytes(&request.data)?;
+        (self.proc)(info);
         Ok(())
     }
 }

libs/socket_cs/src/udp.rs

@@ -1,7 +1,7 @@
 use async_trait::async_trait;
 use hbb_common::{
     log,
-    tokio::{self, sync::Notify},
+    tokio::{self, runtime::Runtime, sync::Notify, task::JoinHandle},
     udp::FramedSocket,
     ResultType,
 };
@@ -14,6 +14,8 @@ use std::sync::Arc;
 pub struct Server {
     port: u16,
     exit_notify: Arc<Notify>,
+    rt: Arc<Runtime>,
+    join_handler: Option<JoinHandle<()>>,
 }
 
 pub struct UdpRequest {
@@ -33,19 +35,27 @@ pub struct UdpHandlers {
 }
 
 impl Server {
-    pub fn new(port: u16, exit_notify: Arc<Notify>) -> Self {
-        Self { port, exit_notify }
+    pub fn create(port: u16) -> ResultType<Self> {
+        let rt = Arc::new(Runtime::new()?);
+        let exit_notify = Arc::new(Notify::new());
+        Ok(Self {
+            port,
+            exit_notify,
+            rt,
+            join_handler: None,
+        })
     }
 
     /// Start server with the handlers.
-    pub async fn start(&self, handlers: UdpHandlers) -> ResultType<()> {
+    pub async fn start(&mut self, handlers: UdpHandlers) -> ResultType<()> {
         let exit_notify = self.exit_notify.clone();
 
         let addr = SocketAddr::from(([0, 0, 0, 0], self.port));
         let mut server = FramedSocket::new(addr).await?;
         log::trace!("succeeded to bind {} for discovery server", addr);
 
-        tokio::spawn(async move {
+        let rt = self.rt.clone();
+        let join_handler = rt.clone().spawn(async move {
             let handlers = Arc::new(handlers.handlers);
             loop {
                 tokio::select! {
@@ -56,11 +66,12 @@ impl Server {
                     n = server.next() => {
                         log::info!("received message");
                         let handlers = handlers.clone();
+                        let rt = rt.clone();
                         match n {
                             Some(Ok((data, addr))) => {
                                 match data.iter().position(|x| x == &crate::CMD_TOKEN) {
                                     Some(p) => {
-                                        tokio::spawn(async move {
+                                        rt.spawn(async move {
                                             let cmd = data[0..p].to_vec();
                                             match handlers.get(&cmd) {
                                                 Some(h) => {
@@ -92,8 +103,27 @@ impl Server {
                 }
             }
         });
+
+        self.join_handler = Some(join_handler);
         Ok(())
     }
+
+    pub async fn shutdonw(&mut self) {
+        self.exit_notify.notify_one();
+        if let Some(h) = self.join_handler.take() {
+            if let Err(e) = h.await {
+                log::error!("failed to join udp server, {}", e);
+            }
+        }
+    }
+}
+
+impl Drop for Server {
+    fn drop(&mut self) {
+        self.rt.clone().block_on(async {
+            self.shutdonw().await;
+        })
+    }
 }
 
 impl UdpHandlers {