rustdesk/rustdesk-server
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server.git synced 2026-04-20 23:20:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor

Browse Source
This commit is contained in:
open-trade
2020-09-25 17:57:06 +08:00
parent e40b2419a4
commit 1417754b23
1 changed files with 28 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
src/rendezvous_server.rs
View File

@@ -169,16 +169,26 @@ impl RendezvousServer {
log::debug!("Tcp connection from {:?}", addr); log::debug!("Tcp connection from {:?}", addr);
let (a, mut b) = Framed::new(stream, BytesCodec::new()).split(); let (a, mut b) = Framed::new(stream, BytesCodec::new()).split();
let tcp_punch = rs.tcp_punch.clone(); let tcp_punch = rs.tcp_punch.clone();
tcp_punch.lock().unwrap().insert(addr, a);
let mut rs = rs.clone(); let mut rs = rs.clone();
tokio::spawn(async move { tokio::spawn(async move {
let mut sender = Some(a);
while let Ok(Some(Ok(bytes))) = timeout(30_000, 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) { if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
match msg_in.union { match msg_in.union {
Some(rendezvous_message::Union::punch_hole_request(ph)) => { Some(rendezvous_message::Union::punch_hole_request(ph)) => {
if let Some(sender) = sender.take() {
tcp_punch.lock().unwrap().insert(addr, sender);
} else {
break;
}
allow_err!(rs.handle_tcp_punch_hole_request(addr, ph.id).await); allow_err!(rs.handle_tcp_punch_hole_request(addr, ph.id).await);
} }
Some(rendezvous_message::Union::request_relay(mut rf)) => { Some(rendezvous_message::Union::request_relay(mut rf)) => {
if let Some(sender) = sender.take() {
tcp_punch.lock().unwrap().insert(addr, sender);
} else {
break;
}
if let Some(peer) = rs.pm.map.read().unwrap().get(&rf.id).map(|x| x.clone()) { if let Some(peer) = rs.pm.map.read().unwrap().get(&rf.id).map(|x| x.clone()) {
let mut msg_out = RendezvousMessage::new(); let mut msg_out = RendezvousMessage::new();
rf.socket_addr = AddrMangle::encode(addr); rf.socket_addr = AddrMangle::encode(addr);
@@ -191,12 +201,7 @@ impl RendezvousServer {
rfr.socket_addr = Default::default(); rfr.socket_addr = Default::default();
let mut msg_out = RendezvousMessage::new(); let mut msg_out = RendezvousMessage::new();
msg_out.set_request_relay_response(rfr); msg_out.set_request_relay_response(rfr);
let sender_b = rs.tcp_punch.lock().unwrap().remove(&addr_b); allow_err!(rs.send_to_tcp_sync(&msg_out, addr_b).await);
if let Some(mut sender_b) = sender_b {
if let Ok(bytes) = msg_out.write_to_bytes() {
allow_err!(sender_b.send(Bytes::from(bytes)).await);
}
}
break; break;
} }
Some(rendezvous_message::Union::punch_hole_sent(phs)) => { Some(rendezvous_message::Union::punch_hole_sent(phs)) => {
@@ -207,6 +212,19 @@ impl RendezvousServer {
allow_err!(rs.handle_local_addr(la, addr, None).await); allow_err!(rs.handle_local_addr(la, addr, None).await);
break; break;
} }
Some(rendezvous_message::Union::test_nat_request(_)) => {
let mut msg_out = RendezvousMessage::new();
msg_out.set_test_nat_response(TestNatResponse {
port: addr.port() as _,
..Default::default()
});
if let Some(tcp) = sender.as_mut() {
if let Ok(bytes) = msg_out.write_to_bytes() {
allow_err!(tcp.send(Bytes::from(bytes)).await);
}
}
break;
}
_ => { _ => {
break; break;
} }
@@ -215,7 +233,9 @@ impl RendezvousServer {
break; break;
} }
} }
rs.tcp_punch.lock().unwrap().remove(&addr); if sender.is_none() {
rs.tcp_punch.lock().unwrap().remove(&addr);
}
log::debug!("Tcp connection from {:?} closed", addr); log::debug!("Tcp connection from {:?} closed", addr);
}); });
} }
@@ -321,14 +341,6 @@ impl RendezvousServer {
socket.send(&msg_out, addr).await?; socket.send(&msg_out, addr).await?;
} }
} }
Some(rendezvous_message::Union::test_nat_request(_)) => {
let mut msg_out = RendezvousMessage::new();
msg_out.set_test_nat_response(TestNatResponse {
port: addr.port() as _,
..Default::default()
});
socket.send(&msg_out, addr).await?;
}
_ => {} _ => {}
} }
} }