rustdesk/rustdesk-server-demo
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server-demo.git synced 2026-04-21 07:30:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Work now

Browse Source
This commit is contained in:
rustdesk
2021-05-27 20:44:13 +08:00
parent c18fa776ce
commit 89a8704909
2 changed files with 49 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
src/main.rs
View File

@@ -1,4 +1,5 @@
use hbb_common::{
bytes::BytesMut,
protobuf::Message as _,
rendezvous_proto::*,
tcp::{new_listener, FramedStream},
@@ -17,29 +18,15 @@ async fn main() {
loop {
tokio::select! {
Some(Ok((bytes, addr))) = socket.next() => {
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
match msg_in.union {
Some(rendezvous_message::Union::register_peer(rp)) => {
id_map.insert(rp.id, addr);
}
Some(rendezvous_message::Union::register_pk(_)) => {
let mut msg_out = RendezvousMessage::new();
msg_out.set_register_pk_response(RegisterPkResponse {
result: register_pk_response::Result::OK.into(),
..Default::default()
});
socket.send(&msg_out, addr).await.ok();
}
_ => {}
}
}
handle_udp(&mut socket, bytes, addr, &mut id_map).await;
}
Ok((stream, _)) = listener.accept() => {
Ok((stream, addr)) = listener.accept() => {
let mut stream = FramedStream::from(stream);
if let Some(Ok(bytes)) = stream.next_timeout(3000).await {
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
match msg_in.union {
Some(rendezvous_message::Union::punch_hole_request(ph)) => {
println!("punch_hole_request {:?}", addr);
if let Some(addr) = id_map.get(&ph.id) {
let mut msg_out = RendezvousMessage::new();
msg_out.set_request_relay(RequestRelay {
@@ -51,6 +38,7 @@ async fn main() {
}
}
Some(rendezvous_message::Union::relay_response(_)) => {
println!("relay_response {:?}", addr);
let mut msg_out = RendezvousMessage::new();
msg_out.set_relay_response(RelayResponse {
relay_server: relay_server.clone(),
@@ -64,7 +52,7 @@ async fn main() {
if let Ok((stream_b, _)) = rlistener.accept().await {
let mut stream_b = FramedStream::from(stream_b);
stream_b.next_timeout(3_000).await;
relay(stream_a, stream_b).await;
relay(stream_a, stream_b, &mut socket, &mut id_map).await;
}
}
}
@@ -78,13 +66,21 @@ async fn main() {
}
}
async fn relay(stream: FramedStream, peer: FramedStream) {
async fn relay(
stream: FramedStream,
peer: FramedStream,
socket: &mut FramedSocket,
id_map: &mut std::collections::HashMap<String, std::net::SocketAddr>,
) {
let mut peer = peer;
let mut stream = stream;
peer.set_raw();
stream.set_raw();
loop {
tokio::select! {
Some(Ok((bytes, addr))) = socket.next() => {
handle_udp(socket, bytes, addr, id_map).await;
}
res = peer.next() => {
if let Some(Ok(bytes)) = res {
stream.send_bytes(bytes.into()).await.ok();
@@ -102,3 +98,32 @@ async fn relay(stream: FramedStream, peer: FramedStream) {
}
}
}
async fn handle_udp(
socket: &mut FramedSocket,
bytes: BytesMut,
addr: std::net::SocketAddr,
id_map: &mut std::collections::HashMap<String, std::net::SocketAddr>,
) {
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
match msg_in.union {
Some(rendezvous_message::Union::register_peer(rp)) => {
println!("register_peer {:?}", addr);
id_map.insert(rp.id, addr);
let mut msg_out = RendezvousMessage::new();
msg_out.set_register_peer_response(RegisterPeerResponse::new());
socket.send(&msg_out, addr).await.ok();
}
Some(rendezvous_message::Union::register_pk(_)) => {
println!("register_pk {:?}", addr);
let mut msg_out = RendezvousMessage::new();
msg_out.set_register_pk_response(RegisterPkResponse {
result: register_pk_response::Result::OK.into(),
..Default::default()
});
socket.send(&msg_out, addr).await.ok();
}
_ => {}
}
}
}