mirror/rustdesk-server
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server.git synced 2025-01-18 23:35:18 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix on request_relay

Browse Source
This commit is contained in:
open-trade 2020-09-18 00:47:15 +08:00
parent fcd949c556
commit 0a1fbadb4f
4 changed files with 17 additions and 104 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
Cargo.toml
View File

@ -5,8 +5,8 @@ authors = ["open-trade <info@opentradesolutions.com>"]
edition = "2018"
[[bin]]
name = "hbbf"
path = "src/hbbf/main.rs"
name = "hbbr"
path = "src/hbbr/main.rs"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

2
libs/hbb_common

@ -1 +1 @@
Subproject commit fdaa6c740c475f6023c93a148bc851655d0093ac
Subproject commit 4b0de359f90425ce9c72903b7a9e5fcbe505a5cc

84
src/hbbf/main.rs
View File

@ -1,84 +0,0 @@
use hbb_common::{
env_logger::*,
log,
protobuf::Message as _,
rendezvous_proto::*,
sleep,
tcp::{new_listener, FramedStream},
tokio, ResultType,
};
use std::{
collections::HashMap,
net::SocketAddr,
sync::{Arc, Mutex},
};
lazy_static::lazy_static! {
static ref PEERS: Arc<Mutex<HashMap<String, FramedStream>>> = Arc::new(Mutex::new(HashMap::new()));
}
#[tokio::main]
async fn main() -> ResultType<()> {
init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info"));
let addr = "0.0.0.0:21117";
log::info!("Listening on {}", addr);
let mut listener = new_listener(addr, true).await?;
loop {
tokio::select! {
Ok((stream, addr)) = listener.accept() => {
tokio::spawn(async move {
make_pair(FramedStream::from(stream), addr).await.ok();
});
}
}
}
}
async fn make_pair(stream: FramedStream, addr: SocketAddr) -> ResultType<()> {
let mut stream = stream;
if let Some(Ok(bytes)) = stream.next_timeout(30_000).await {
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
if let Some(rendezvous_message::Union::request_relay(rf)) = msg_in.union {
if !rf.uuid.is_empty() {
let peer = PEERS.lock().unwrap().remove(&rf.uuid);
if let Some(peer) = peer {
log::info!("Forward request {} from {} got paired", rf.uuid, addr);
return relay(stream, peer).await;
} else {
log::info!("New relay request {} from {}", rf.uuid, addr);
PEERS.lock().unwrap().insert(rf.uuid.clone(), stream);
sleep(30.).await;
PEERS.lock().unwrap().remove(&rf.uuid);
}
}
}
}
}
Ok(())
}
async fn relay(stream: FramedStream, peer: FramedStream) -> ResultType<()> {
let mut peer = peer;
let mut stream = stream;
peer.set_raw();
stream.set_raw();
loop {
tokio::select! {
res = peer.next() => {
if let Some(Ok(bytes)) = res {
stream.send_bytes(bytes.into()).await?;
} else {
break;
}
},
res = stream.next() => {
if let Some(Ok(bytes)) = res {
peer.send_bytes(bytes.into()).await?;
} else {
break;
}
},
}
}
Ok(())
}

31
src/rendezvous_server.rs
View File

@ -10,6 +10,7 @@ use hbb_common::{
protobuf::Message as _,
rendezvous_proto::*,
tcp::new_listener,
timeout,
tokio::{self, net::TcpStream, sync::mpsc},
tokio_util::codec::Framed,
udp::FramedSocket,
@ -151,27 +152,28 @@ impl RendezvousServer {
tcp_punch.lock().unwrap().insert(addr, a);
let mut rs = rs.clone();
tokio::spawn(async move {
while let Some(Ok(bytes)) = 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) {
match msg_in.union {
Some(rendezvous_message::Union::punch_hole_request(ph)) => {
allow_err!(rs.handle_tcp_punch_hole_request(addr, ph.id).await);
}
Some(rendezvous_message::Union::request_relay(mut rf)) => {
if !rs.pm.is_in_memory(&rf.id) {
break;
if let Some(peer) = rs.pm.map.read().unwrap().get(&rf.id).map(|x| x.clone()) {
let mut msg_out = RendezvousMessage::new();
rf.socket_addr = AddrMangle::encode(addr);
msg_out.set_request_relay(rf);
rs.tx.send((msg_out, peer.socket_addr)).ok();
}
let mut msg_out = RendezvousMessage::new();
rf.socket_addr = AddrMangle::encode(addr);
msg_out.set_request_relay(rf);
rs.tx.send((msg_out, addr)).ok();
}
Some(rendezvous_message::Union::request_relay_response(mut rfr)) => {
let addr_b = AddrMangle::decode(&rfr.socket_addr);
rfr.socket_addr = Default::default();
let sender_b= rs.tcp_punch.lock().unwrap().remove(&addr_b);
let mut msg_out = RendezvousMessage::new();
msg_out.set_request_relay_response(rfr);
let sender_b = rs.tcp_punch.lock().unwrap().remove(&addr_b);
if let Some(mut sender_b) = sender_b {
if let Ok(bytes) = rfr.write_to_bytes() {
if let Ok(bytes) = msg_out.write_to_bytes() {
allow_err!(sender_b.send(Bytes::from(bytes)).await);
}
}
@ -185,7 +187,9 @@ impl RendezvousServer {
allow_err!(rs.handle_local_addr(&la, addr, None).await);
break;
}
_ => {}
_ => {
break;
}
}
} else {
break;
@ -255,13 +259,6 @@ impl RendezvousServer {
});
}
}
Some(rendezvous_message::Union::request_relay(rf)) => {
if self.pm.is_in_memory(&rf.id) {
let mut msg_out = RendezvousMessage::new();
msg_out.set_request_relay(rf);
socket.send(&msg_out, addr).await?
}
}
Some(rendezvous_message::Union::punch_hole_sent(phs)) => {
self.handle_hole_sent(phs, addr, Some(socket)).await?;
}