From 89a87049098236d73de2af21bd0292b92a138491 Mon Sep 17 00:00:00 2001
From: rustdesk <info@rustdesk.com>
Date: Thu, 27 May 2021 20:44:13 +0800
Subject: [PATCH] Work now

---
 README.md   |  5 +++++
 src/main.rs | 63 +++++++++++++++++++++++++++++++++++++----------------
 2 files changed, 49 insertions(+), 19 deletions(-)

diff --git a/README.md b/README.md
index d771f46..5aacc9f 100644
--- a/README.md
+++ b/README.md
@@ -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
+```
diff --git a/src/main.rs b/src/main.rs
index e93f92c..694aa02 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -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();
+            }
+            _ => {}
+        }
+    }
+}