refactored

1 files changed, 12 insertions, 12 deletions

diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs
index 9a47305..c1c9467 100644
--- a/src/rendezvous_server.rs
+++ b/src/rendezvous_server.rs
@@ -1,6 +1,6 @@
 use hbb_common::{
-    bytes::BytesMut, log, protobuf::parse_from_bytes, rendezvous_proto::*, udp::FramedSocket, tcp::new_listener,
-    AddrMangle, ResultType,
+    allow_err, bytes::BytesMut, log, protobuf::parse_from_bytes, rendezvous_proto::*,
+    tcp::new_listener, tokio, udp::FramedSocket, AddrMangle, ResultType,
 };
 use std::{collections::HashMap, net::SocketAddr};
 
@@ -20,19 +20,19 @@ impl RendezvousServer {
         let mut rs = Self {
             peer_map: PeerMap::new(),
         };
-        // used to test if udp/tcp share the same NAT port, yes in my test.
+        // tcp listener used to test if udp/tcp share the same NAT port, yes in my test.
         // also be used to help client to get local ip.
-        let addr = addr.to_string();
-        hbb_common::tokio::spawn(async {
-            let mut l = new_listener(addr, true).await.unwrap();
-            while let Ok((_, addr)) = l.accept().await {
-                log::debug!("Tcp connection from {:?}", addr);
+        let mut listener = new_listener(addr, true).await.unwrap();
+        loop {
+            tokio::select! {
+                Some(Ok((bytes, addr))) = socket.next() => {
+                    allow_err!(rs.handle_msg(&bytes, addr, &mut socket).await);
+                }
+                Ok((_, addr)) = listener.accept() => {
+                    log::debug!("Tcp connection from {:?}", addr);
+                }
             }
-        });
-        while let Some(Ok((bytes, addr))) = socket.next().await {
-            rs.handle_msg(&bytes, addr, &mut socket).await?;
         }
-        Ok(())
     }
 
     pub async fn handle_msg(