From e1e4c9a9652caf4ab0aa643ea80db2cc70854a2b Mon Sep 17 00:00:00 2001 From: open-trade Date: Thu, 12 Mar 2020 22:06:09 +0800 Subject: [PATCH] refactored --- src/rendezvous_server.rs | 24 ++++++++++++------------ 1 file 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(