diff options
| -rw-r--r-- | Cargo.toml | 4 | ||||
| m--------- | libs/hbb_common | 0 | ||||
| -rw-r--r-- | src/hbbf/main.rs | 84 | ||||
| -rw-r--r-- | src/rendezvous_server.rs | 31 |
4 files changed, 16 insertions, 103 deletions
@@ -5,8 +5,8 @@ authors = ["open-trade <[email protected]>"] edition = "2018" [[bin]] -name = "hbbf" -path = "src/hbbf/main.rs" +name = "hbbr" +path = "src/hbbr/main.rs" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/libs/hbb_common b/libs/hbb_common -Subproject fdaa6c740c475f6023c93a148bc851655d0093a +Subproject 4b0de359f90425ce9c72903b7a9e5fcbe505a5c diff --git a/src/hbbf/main.rs b/src/hbbf/main.rs deleted file mode 100644 index a9df3df..0000000 --- a/src/hbbf/main.rs +++ /dev/null @@ -1,84 +0,0 @@ -use hbb_common::{ - env_logger::*, - log, - protobuf::Message as _, - rendezvous_proto::*, - sleep, - tcp::{new_listener, FramedStream}, - tokio, ResultType, -}; -use std::{ - collections::HashMap, - net::SocketAddr, - sync::{Arc, Mutex}, -}; - -lazy_static::lazy_static! { - static ref PEERS: Arc<Mutex<HashMap<String, FramedStream>>> = Arc::new(Mutex::new(HashMap::new())); -} - -#[tokio::main] -async fn main() -> ResultType<()> { - init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info")); - let addr = "0.0.0.0:21117"; - log::info!("Listening on {}", addr); - let mut listener = new_listener(addr, true).await?; - loop { - tokio::select! { - Ok((stream, addr)) = listener.accept() => { - tokio::spawn(async move { - make_pair(FramedStream::from(stream), addr).await.ok(); - }); - } - } - } -} - -async fn make_pair(stream: FramedStream, addr: SocketAddr) -> ResultType<()> { - let mut stream = stream; - if let Some(Ok(bytes)) = stream.next_timeout(30_000).await { - if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) { - if let Some(rendezvous_message::Union::request_relay(rf)) = msg_in.union { - if !rf.uuid.is_empty() { - let peer = PEERS.lock().unwrap().remove(&rf.uuid); - if let Some(peer) = peer { - log::info!("Forward request {} from {} got paired", rf.uuid, addr); - return relay(stream, peer).await; - } else { - log::info!("New relay request {} from {}", rf.uuid, addr); - PEERS.lock().unwrap().insert(rf.uuid.clone(), stream); - sleep(30.).await; - PEERS.lock().unwrap().remove(&rf.uuid); - } - } - } - } - } - Ok(()) -} - -async fn relay(stream: FramedStream, peer: FramedStream) -> ResultType<()> { - let mut peer = peer; - let mut stream = stream; - peer.set_raw(); - stream.set_raw(); - loop { - tokio::select! { - res = peer.next() => { - if let Some(Ok(bytes)) = res { - stream.send_bytes(bytes.into()).await?; - } else { - break; - } - }, - res = stream.next() => { - if let Some(Ok(bytes)) = res { - peer.send_bytes(bytes.into()).await?; - } else { - break; - } - }, - } - } - Ok(()) -} diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs index 77df3df..139b32b 100644 --- a/src/rendezvous_server.rs +++ b/src/rendezvous_server.rs @@ -10,6 +10,7 @@ use hbb_common::{ protobuf::Message as _, rendezvous_proto::*, tcp::new_listener, + timeout, tokio::{self, net::TcpStream, sync::mpsc}, tokio_util::codec::Framed, udp::FramedSocket, @@ -151,27 +152,28 @@ impl RendezvousServer { tcp_punch.lock().unwrap().insert(addr, a); let mut rs = rs.clone(); tokio::spawn(async move { - while let Some(Ok(bytes)) = b.next().await { + while let Ok(Some(Ok(bytes))) = timeout(30_000, b.next()).await { if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) { match msg_in.union { Some(rendezvous_message::Union::punch_hole_request(ph)) => { allow_err!(rs.handle_tcp_punch_hole_request(addr, ph.id).await); } Some(rendezvous_message::Union::request_relay(mut rf)) => { - if !rs.pm.is_in_memory(&rf.id) { - break; + if let Some(peer) = rs.pm.map.read().unwrap().get(&rf.id).map(|x| x.clone()) { + let mut msg_out = RendezvousMessage::new(); + rf.socket_addr = AddrMangle::encode(addr); + msg_out.set_request_relay(rf); + rs.tx.send((msg_out, peer.socket_addr)).ok(); } - let mut msg_out = RendezvousMessage::new(); - rf.socket_addr = AddrMangle::encode(addr); - msg_out.set_request_relay(rf); - rs.tx.send((msg_out, addr)).ok(); } Some(rendezvous_message::Union::request_relay_response(mut rfr)) => { let addr_b = AddrMangle::decode(&rfr.socket_addr); rfr.socket_addr = Default::default(); - let sender_b= rs.tcp_punch.lock().unwrap().remove(&addr_b); + let mut msg_out = RendezvousMessage::new(); + msg_out.set_request_relay_response(rfr); + let sender_b = rs.tcp_punch.lock().unwrap().remove(&addr_b); if let Some(mut sender_b) = sender_b { - if let Ok(bytes) = rfr.write_to_bytes() { + if let Ok(bytes) = msg_out.write_to_bytes() { allow_err!(sender_b.send(Bytes::from(bytes)).await); } } @@ -185,7 +187,9 @@ impl RendezvousServer { allow_err!(rs.handle_local_addr(&la, addr, None).await); break; } - _ => {} + _ => { + break; + } } } else { break; @@ -255,13 +259,6 @@ impl RendezvousServer { }); } } - Some(rendezvous_message::Union::request_relay(rf)) => { - if self.pm.is_in_memory(&rf.id) { - let mut msg_out = RendezvousMessage::new(); - msg_out.set_request_relay(rf); - socket.send(&msg_out, addr).await? - } - } Some(rendezvous_message::Union::punch_hole_sent(phs)) => { self.handle_hole_sent(phs, addr, Some(socket)).await?; } |