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

refactor

Browse Source
This commit is contained in:
open-trade 2020-03-07 22:45:22 +08:00
parent 5be70826aa
commit d8b829818e
1 changed files with 33 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
src/rendezvous_server.rs
View File

@ -1,5 +1,5 @@
use super::message_proto::*; use super::message_proto::*;
use bytes::Bytes; use bytes::{Bytes, BytesMut};
use futures::SinkExt; use futures::SinkExt;
use protobuf::{parse_from_bytes, Message as _}; use protobuf::{parse_from_bytes, Message as _};
use std::{ use std::{
@ -8,11 +8,7 @@ use std::{
net::{Ipv4Addr, SocketAddr, SocketAddrV4}, net::{Ipv4Addr, SocketAddr, SocketAddrV4},
time::{Duration, SystemTime, UNIX_EPOCH}, time::{Duration, SystemTime, UNIX_EPOCH},
}; };
use tokio::{ use tokio::{net::UdpSocket, stream::StreamExt, time::delay_for};
net::UdpSocket,
stream::StreamExt,
time::{self, delay_for},
};
use tokio_util::{codec::BytesCodec, udp::UdpFramed}; use tokio_util::{codec::BytesCodec, udp::UdpFramed};
/// Certain router and firewalls scan the packet and if they /// Certain router and firewalls scan the packet and if they
@ -74,43 +70,43 @@ impl RendezvousServer {
peer_map: PeerMap::new(), peer_map: PeerMap::new(),
}; };
while let Some(Ok((bytes, addr))) = socket.next().await { while let Some(Ok((bytes, addr))) = socket.next().await {
if let SocketAddr::V4(addr_v4) = addr { rs.handle_msg(&bytes, addr, &mut socket).await?;
if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) {
match msg_in.union {
Some(Message_oneof_union::register_peer(rp)) => {
if rp.hbb_addr.len() > 0 {
rs.peer_map.insert(
rp.hbb_addr,
Peer {
socket_addr: addr_v4,
},
);
}
}
Some(Message_oneof_union::peek_peer(pp)) => {
rs.handle_peek_peer(&pp, addr, &mut socket).await?;
}
_ => {}
}
}
}
} }
Ok(()) Ok(())
} }
pub async fn handle_peek_peer( pub async fn handle_msg(
&self, &mut self,
pp: &PeekPeer, bytes: &BytesMut,
addr: SocketAddr, addr: SocketAddr,
socket: &mut FramedSocket, socket: &mut FramedSocket,
) -> ResultType { ) -> ResultType {
if let Some(peer) = self.peer_map.get(&pp.hbb_addr) { if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) {
let mut msg_out = Message::new(); if let SocketAddr::V4(addr_v4) = addr {
msg_out.set_peek_peer_response(PeekPeerResponse { match msg_in.union {
socket_addr: V4AddrMangle::encode(&peer.socket_addr), Some(Message_oneof_union::register_peer(rp)) => {
..Default::default() if rp.hbb_addr.len() > 0 {
}); self.peer_map.insert(
send_to(&msg_out, addr, socket).await?; rp.hbb_addr,
Peer {
socket_addr: addr_v4,
},
);
}
}
Some(Message_oneof_union::peek_peer(pp)) => {
if let Some(peer) = self.peer_map.get(&pp.hbb_addr) {
let mut msg_out = Message::new();
msg_out.set_peek_peer_response(PeekPeerResponse {
socket_addr: V4AddrMangle::encode(&peer.socket_addr),
..Default::default()
});
send_to(&msg_out, addr, socket).await?;
}
}
_ => {}
}
}
} }
Ok(()) Ok(())
} }
@ -159,7 +155,7 @@ mod tests {
}); });
send_to(&msg_out, to_addr, &mut socket).await; send_to(&msg_out, to_addr, &mut socket).await;
if let Ok(Some(Ok((bytes, _)))) = if let Ok(Some(Ok((bytes, _)))) =
time::timeout(Duration::from_millis(1), socket.next()).await tokio::time::timeout(Duration::from_millis(1), socket.next()).await
{ {
if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) { if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) {
assert_eq!( assert_eq!(