mirror of
https://github.com/rustdesk/rustdesk-server.git
synced 2025-07-03 08:15:25 +00:00
tcp punch
This commit is contained in:
parent
66bbe7baa8
commit
0034fd695b
735
Cargo.lock
generated
735
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@ -1 +1 @@
|
||||
Subproject commit 3d3d1bd1432f14341c7b90393206a787b95a1638
|
||||
Subproject commit 56827464aa46ec0ecb83944d7b783fb95a93f74c
|
@ -1,8 +1,27 @@
|
||||
use hbb_common::{
|
||||
allow_err, bytes::BytesMut, log, protobuf::parse_from_bytes, rendezvous_proto::*,
|
||||
tcp::new_listener, tokio, udp::FramedSocket, AddrMangle, ResultType,
|
||||
allow_err,
|
||||
bytes::Bytes,
|
||||
bytes::BytesMut,
|
||||
bytes_codec::BytesCodec,
|
||||
futures_util::{
|
||||
sink::SinkExt,
|
||||
stream::{SplitSink, StreamExt},
|
||||
},
|
||||
log,
|
||||
protobuf::{parse_from_bytes, Message as _},
|
||||
rendezvous_proto::*,
|
||||
tcp::new_listener,
|
||||
tokio::{self, net::TcpStream, sync::mpsc},
|
||||
tokio_util::codec::Framed,
|
||||
udp::FramedSocket,
|
||||
AddrMangle, ResultType,
|
||||
};
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
net::SocketAddr,
|
||||
sync::{Arc, Mutex},
|
||||
time::Instant,
|
||||
};
|
||||
use std::{collections::HashMap, net::SocketAddr, time::Instant};
|
||||
|
||||
pub struct Peer {
|
||||
socket_addr: SocketAddr,
|
||||
@ -14,6 +33,7 @@ const REG_TIMEOUT: i32 = 30_000;
|
||||
|
||||
pub struct RendezvousServer {
|
||||
peer_map: PeerMap,
|
||||
tcp_punch: Arc<Mutex<HashMap<SocketAddr, SplitSink<Framed<TcpStream, BytesCodec>, Bytes>>>>,
|
||||
}
|
||||
|
||||
impl RendezvousServer {
|
||||
@ -21,17 +41,38 @@ impl RendezvousServer {
|
||||
let mut socket = FramedSocket::new(addr).await?;
|
||||
let mut rs = Self {
|
||||
peer_map: PeerMap::new(),
|
||||
tcp_punch: Arc::new(Mutex::new(HashMap::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 (tx, mut rx) = mpsc::unbounded_channel::<(SocketAddr, String)>();
|
||||
let mut listener = new_listener(addr, true).await.unwrap();
|
||||
loop {
|
||||
tokio::select! {
|
||||
Some((addr, id)) = rx.recv() => {
|
||||
allow_err!(rs.handle_punch_hole_request(addr, &id, &mut socket).await);
|
||||
}
|
||||
Some(Ok((bytes, addr))) = socket.next() => {
|
||||
allow_err!(rs.handle_msg(&bytes, addr, &mut socket).await);
|
||||
}
|
||||
Ok((_, addr)) = listener.accept() => {
|
||||
Ok((stream, addr)) = listener.accept() => {
|
||||
log::debug!("Tcp connection from {:?}", addr);
|
||||
let (a, mut b) = Framed::new(stream, BytesCodec::new()).split();
|
||||
let tcp_punch = rs.tcp_punch.clone();
|
||||
tcp_punch.lock().unwrap().insert(addr, a);
|
||||
let tx = tx.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(Ok(bytes)) = b.next().await {
|
||||
if let Ok(msg_in) = parse_from_bytes::<RendezvousMessage>(&bytes) {
|
||||
match msg_in.union {
|
||||
Some(rendezvous_message::Union::punch_hole_request(ph)) => {
|
||||
tx.send((addr, ph.id)).ok();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
tcp_punch.lock().unwrap().remove(&addr);
|
||||
log::debug!("Tcp connection from {:?} closed", addr);
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -49,71 +90,20 @@ impl RendezvousServer {
|
||||
// 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, last_reg_time: Instant::now() });
|
||||
self.peer_map.insert(
|
||||
rp.id,
|
||||
Peer {
|
||||
socket_addr: addr,
|
||||
last_reg_time: Instant::now(),
|
||||
},
|
||||
);
|
||||
let mut msg_out = RendezvousMessage::new();
|
||||
msg_out.set_register_peer_response(RegisterPeerResponse::default());
|
||||
socket.send(&msg_out, addr).await?
|
||||
}
|
||||
}
|
||||
Some(rendezvous_message::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) {
|
||||
if peer.last_reg_time.elapsed().as_millis() as i32 >= REG_TIMEOUT {
|
||||
let mut msg_out = RendezvousMessage::new();
|
||||
msg_out.set_punch_hole_response(PunchHoleResponse {
|
||||
failure: punch_hole_response::Failure::OFFLINE.into(),
|
||||
..Default::default()
|
||||
});
|
||||
return socket.send(&msg_out, addr).await;
|
||||
}
|
||||
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: punch_hole_response::Failure::ID_NOT_EXIST.into(),
|
||||
..Default::default()
|
||||
});
|
||||
socket.send(&msg_out, addr).await?
|
||||
}
|
||||
self.handle_punch_hole_request(addr, &ph.id, socket).await?;
|
||||
}
|
||||
Some(rendezvous_message::Union::punch_hole_sent(phs)) => {
|
||||
// punch hole sent from B, tell A that B is ready to be connected
|
||||
@ -121,10 +111,11 @@ impl RendezvousServer {
|
||||
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),
|
||||
socket_addr: AddrMangle::encode(addr),
|
||||
..Default::default()
|
||||
});
|
||||
socket.send(&msg_out, addr_a).await?;
|
||||
self.send_to_tcp(&msg_out, addr_a).await?;
|
||||
}
|
||||
Some(rendezvous_message::Union::local_addr(la)) => {
|
||||
// forward local addrs of B to A
|
||||
@ -136,10 +127,91 @@ impl RendezvousServer {
|
||||
..Default::default()
|
||||
});
|
||||
socket.send(&msg_out, addr_a).await?;
|
||||
self.send_to_tcp(&msg_out, addr_a).await?;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_punch_hole_request(
|
||||
&mut self,
|
||||
addr: SocketAddr,
|
||||
id: &str,
|
||||
socket: &mut FramedSocket,
|
||||
) -> ResultType<()> {
|
||||
// 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(id) {
|
||||
if peer.last_reg_time.elapsed().as_millis() as i32 >= REG_TIMEOUT {
|
||||
let mut msg_out = RendezvousMessage::new();
|
||||
msg_out.set_punch_hole_response(PunchHoleResponse {
|
||||
failure: punch_hole_response::Failure::OFFLINE.into(),
|
||||
..Default::default()
|
||||
});
|
||||
return socket.send(&msg_out, addr).await;
|
||||
}
|
||||
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 {:?}",
|
||||
id,
|
||||
&peer.socket_addr,
|
||||
&addr
|
||||
);
|
||||
msg_out.set_fetch_local_addr(FetchLocalAddr {
|
||||
socket_addr,
|
||||
..Default::default()
|
||||
});
|
||||
} else {
|
||||
log::debug!(
|
||||
"Punch hole {:?} {:?} request from {:?}",
|
||||
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: punch_hole_response::Failure::ID_NOT_EXIST.into(),
|
||||
..Default::default()
|
||||
});
|
||||
socket.send(&msg_out, addr).await?
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_to_tcp(&mut self, msg: &RendezvousMessage, addr: SocketAddr) -> ResultType<()> {
|
||||
let tcp = self.tcp_punch.lock().unwrap().remove(&addr);
|
||||
if let Some(mut tcp) = tcp {
|
||||
if let Ok(bytes) = msg.write_to_bytes() {
|
||||
tokio::spawn(async move {
|
||||
allow_err!(tcp.send(Bytes::from(bytes)).await);
|
||||
log::debug!("Send punch hole to {} via tcp", addr);
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
Loading…
x
Reference in New Issue
Block a user