rustdesk/rustdesk-server-demo
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server-demo.git synced 2025-07-01 07:15:29 +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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
README.md
View File

@ -1,2 +1,7 @@
# A working demo of RustDesk server implementation
This is a super simple working demo implementation with only one relay connection allowed, without NAT traverse, persistence, encryption and any other advanced features. But it can be your good starting point to write your own RustDesk server program.
## How to run
```
IP=<your public ip> cargo run
```

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();
}
_ => {}
}
}
}