Browse Source

refactored

open-trade 5 years ago
parent
commit
e1e4c9a965
1 changed files with 12 additions and 12 deletions
  1. 12 12
      src/rendezvous_server.rs

+ 12 - 12
src/rendezvous_server.rs

@@ -1,6 +1,6 @@
1
 use hbb_common::{
1
 use hbb_common::{
2
-    bytes::BytesMut, log, protobuf::parse_from_bytes, rendezvous_proto::*, udp::FramedSocket, tcp::new_listener,
3
-    AddrMangle, ResultType,
2
+    allow_err, bytes::BytesMut, log, protobuf::parse_from_bytes, rendezvous_proto::*,
3
+    tcp::new_listener, tokio, udp::FramedSocket, AddrMangle, ResultType,
4
 };
4
 };
5
 use std::{collections::HashMap, net::SocketAddr};
5
 use std::{collections::HashMap, net::SocketAddr};
6
 
6
 
@@ -20,19 +20,19 @@ impl RendezvousServer {
20
         let mut rs = Self {
20
         let mut rs = Self {
21
             peer_map: PeerMap::new(),
21
             peer_map: PeerMap::new(),
22
         };
22
         };
23
-        // used to test if udp/tcp share the same NAT port, yes in my test.
23
+        // tcp listener used to test if udp/tcp share the same NAT port, yes in my test.
24
         // also be used to help client to get local ip.
24
         // also be used to help client to get local ip.
25
-        let addr = addr.to_string();
26
-        hbb_common::tokio::spawn(async {
27
-            let mut l = new_listener(addr, true).await.unwrap();
28
-            while let Ok((_, addr)) = l.accept().await {
29
-                log::debug!("Tcp connection from {:?}", addr);
25
+        let mut listener = new_listener(addr, true).await.unwrap();
26
+        loop {
27
+            tokio::select! {
28
+                Some(Ok((bytes, addr))) = socket.next() => {
29
+                    allow_err!(rs.handle_msg(&bytes, addr, &mut socket).await);
30
+                }
31
+                Ok((_, addr)) = listener.accept() => {
32
+                    log::debug!("Tcp connection from {:?}", addr);
33
+                }
30
             }
34
             }
31
-        });
32
-        while let Some(Ok((bytes, addr))) = socket.next().await {
33
-            rs.handle_msg(&bytes, addr, &mut socket).await?;
34
         }
35
         }
35
-        Ok(())
36
     }
36
     }
37
 
37
 
38
     pub async fn handle_msg(
38
     pub async fn handle_msg(