fix on request_relay

Cargo.toml

@@ -5,8 +5,8 @@ authors = ["open-trade <info@opentradesolutions.com>"]
 edition = "2018"
 
 [[bin]]
-name = "hbbf"
+name = "hbbr"
-path = "src/hbbf/main.rs"
+path = "src/hbbr/main.rs"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

libs/hbb_common

@@ -1 +1 @@
-Subproject commit fdaa6c740c475f6023c93a148bc851655d0093ac
+Subproject commit 4b0de359f90425ce9c72903b7a9e5fcbe505a5cc

src/hbbf/main.rs

@@ -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(())
-}

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) {
+                                        if let Some(peer) = rs.pm.map.read().unwrap().get(&rf.id).map(|x| x.clone()) {
-                                            break;
-                                        }
                                             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();
+                                            rs.tx.send((msg_out, peer.socket_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 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?;
                 }