use hbb_common::{ 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}; pub struct Peer { socket_addr: SocketAddr, } type PeerMap = HashMap; pub struct RendezvousServer { peer_map: PeerMap, } impl RendezvousServer { pub async fn start(addr: &str) -> ResultType<()> { let mut socket = FramedSocket::new(addr).await?; let mut rs = Self { peer_map: PeerMap::new(), }; // 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 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); } } } } pub async fn handle_msg( &mut self, bytes: &BytesMut, addr: SocketAddr, socket: &mut FramedSocket, ) -> ResultType<()> { if let Ok(msg_in) = parse_from_bytes::(&bytes) { match msg_in.union { Some(RendezvousMessage_oneof_union::register_peer(rp)) => { // B registered if rp.id.len() > 0 { log::debug!("New peer registered: {:?} {:?}", &rp.id, &addr); self.peer_map.insert(rp.id, Peer { socket_addr: addr }); let mut msg_out = RendezvousMessage::new(); msg_out.set_register_peer_response(RegisterPeerResponse::default()); socket.send(&msg_out, addr).await? } } Some(RendezvousMessage_oneof_union::punch_hole_request(ph)) => { // punch hole request from A, forward to B, // check if in same intranet first, // fetch local addrs if in same intranet. // because punch hole won't work if in the same intranet, // all routers will drop such self-connections. if let Some(peer) = self.peer_map.get(&ph.id) { let mut msg_out = RendezvousMessage::new(); let same_intranet = match peer.socket_addr { SocketAddr::V4(a) => match addr { SocketAddr::V4(b) => a.ip() == b.ip(), _ => false, }, SocketAddr::V6(a) => match addr { SocketAddr::V6(b) => a.ip() == b.ip(), _ => false, }, }; let socket_addr = AddrMangle::encode(&addr); if same_intranet { log::debug!( "Fetch local addr {:?} {:?} request from {:?}", &ph.id, &peer.socket_addr, &addr ); msg_out.set_fetch_local_addr(FetchLocalAddr { socket_addr, ..Default::default() }); } else { log::debug!( "Punch hole {:?} {:?} request from {:?}", &ph.id, &peer.socket_addr, &addr ); msg_out.set_punch_hole(PunchHole { socket_addr, ..Default::default() }); } socket.send(&msg_out, peer.socket_addr).await?; } else { let mut msg_out = RendezvousMessage::new(); msg_out.set_punch_hole_response(PunchHoleResponse { failure: PunchHoleResponse_Failure::ID_NOT_EXIST, ..Default::default() }); socket.send(&msg_out, addr).await? } } Some(RendezvousMessage_oneof_union::punch_hole_sent(phs)) => { // punch hole sent from B, tell A that B is ready to be connected let addr_a = AddrMangle::decode(&phs.socket_addr); log::debug!("Punch hole response to {:?} from {:?}", &addr_a, &addr); let mut msg_out = RendezvousMessage::new(); msg_out.set_punch_hole_response(PunchHoleResponse { socket_addr: AddrMangle::encode(&addr), ..Default::default() }); socket.send(&msg_out, addr_a).await?; } Some(RendezvousMessage_oneof_union::local_addr(la)) => { // forward local addrs of B to A let addr_a = AddrMangle::decode(&la.socket_addr); log::debug!("Local addrs response to {:?} from {:?}", &addr_a, &addr); let mut msg_out = RendezvousMessage::new(); msg_out.set_punch_hole_response(PunchHoleResponse { socket_addr: la.local_addr, ..Default::default() }); socket.send(&msg_out, addr_a).await?; } _ => {} } } Ok(()) } } #[cfg(test)] mod tests { use super::*; use hbb_common::{tokio}; #[allow(unused_must_use)] #[tokio::main] async fn test_rs_async() { let mut port_server: u16 = 0; let socket = FramedSocket::new("127.0.0.1:0").await.unwrap(); if let SocketAddr::V4(addr) = socket.get_ref().local_addr().unwrap() { port_server = addr.port(); } drop(socket); let addr_server = format!("127.0.0.1:{}", port_server); let f1 = RendezvousServer::start(&addr_server); let addr_server = addr_server.parse().unwrap(); let f2 = punch_hole(addr_server); tokio::try_join!(f1, f2); } async fn punch_hole(addr_server: SocketAddr) -> ResultType<()> { // B register it to server let mut socket_b = FramedSocket::new("127.0.0.1:0").await?; let local_addr_b = socket_b.get_ref().local_addr().unwrap(); let mut msg_out = RendezvousMessage::new(); msg_out.set_register_peer(RegisterPeer { id: "123".to_string(), ..Default::default() }); socket_b.send(&msg_out, addr_server).await?; // A send punch request to server let mut socket_a = FramedSocket::new("127.0.0.1:0").await?; let local_addr_a = socket_a.get_ref().local_addr().unwrap(); msg_out.set_punch_hole_request(PunchHoleRequest { id: "123".to_string(), ..Default::default() }); socket_a.send(&msg_out, addr_server).await?; println!( "A {:?} request punch hole to B {:?} via server {:?}", local_addr_a, local_addr_b, addr_server, ); // on B side, responsed to A's punch request forwarded from server if let Some(Ok((bytes, addr))) = socket_b.next_timeout(1000).await { assert_eq!(addr_server, addr); let msg_in = parse_from_bytes::(&bytes)?; let remote_addr_a = AddrMangle::decode(&msg_in.get_punch_hole().socket_addr); assert_eq!(local_addr_a, remote_addr_a); // B punch A socket_b .get_mut() .send_to(&b"SYN"[..], &remote_addr_a) .await?; msg_out.set_punch_hole_sent(PunchHoleSent { socket_addr: AddrMangle::encode(&remote_addr_a), ..Default::default() }); socket_b.send(&msg_out, addr_server).await?; } else { panic!("failed"); } // on A side socket_a.next().await; // skip "SYN" if let Some(Ok((bytes, addr))) = socket_a.next_timeout(1000).await { assert_eq!(addr_server, addr); let msg_in = parse_from_bytes::(&bytes)?; let remote_addr_b = AddrMangle::decode(&msg_in.get_punch_hole_response().socket_addr); assert_eq!(local_addr_b, remote_addr_b); } else { panic!("failed"); } Err("done".into()) } #[test] fn test_rs() { self::test_rs_async(); } }