rustdesk/rustdesk-server
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server.git synced 2026-02-16 02:20:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

test local server

Browse Source
This commit is contained in:
open-trade
2020-03-09 16:16:21 +08:00
parent 52ea29bb94
commit 2812effd6d
2 changed files with 86 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
libs/hbb_common

Submodule libs/hbb_common updated: a9536bcc12...974815baaa

140
src/rendezvous_server.rs
View File

@@ -3,19 +3,14 @@ use futures::SinkExt;
use hbb_common::{ use hbb_common::{
message_proto::*, message_proto::*,
protobuf::{parse_from_bytes, Message as _}, protobuf::{parse_from_bytes, Message as _},
V4AddrMangle, AddrMangle,
};
use std::{
collections::HashMap,
error::Error,
net::{SocketAddr, SocketAddrV4},
time::Duration,
}; };
use std::{collections::HashMap, error::Error, net::SocketAddr, time::Duration};
use tokio::{net::UdpSocket, stream::StreamExt, time::delay_for}; use tokio::{net::UdpSocket, stream::StreamExt, time::delay_for};
use tokio_util::{codec::BytesCodec, udp::UdpFramed}; use tokio_util::{codec::BytesCodec, udp::UdpFramed};
pub struct Peer { pub struct Peer {
socket_addr: SocketAddrV4, socket_addr: SocketAddr,
} }
type PeerMap = HashMap<String, Peer>; type PeerMap = HashMap<String, Peer>;
@@ -48,41 +43,35 @@ impl RendezvousServer {
socket: &mut FramedSocket, socket: &mut FramedSocket,
) -> ResultType { ) -> ResultType {
if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) { if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) {
if let SocketAddr::V4(addr_v4) = addr { match msg_in.union {
match msg_in.union { Some(Message_oneof_union::register_peer(rp)) => {
Some(Message_oneof_union::register_peer(rp)) => { if rp.hbb_addr.len() > 0 {
if rp.hbb_addr.len() > 0 { self.peer_map
self.peer_map.insert( .insert(rp.hbb_addr, Peer { socket_addr: addr });
rp.hbb_addr,
Peer {
socket_addr: addr_v4,
},
);
}
} }
Some(Message_oneof_union::punch_hole_request(ph)) => { }
// punch hole request from A, forward to B Some(Message_oneof_union::punch_hole_request(ph)) => {
if let Some(peer) = self.peer_map.get(&ph.hbb_addr) { // punch hole request from A, forward to B
let mut msg_out = Message::new(); if let Some(peer) = self.peer_map.get(&ph.hbb_addr) {
msg_out.set_punch_hole(PunchHole {
socket_addr: V4AddrMangle::encode(&peer.socket_addr),
..Default::default()
});
send_to(&msg_out, addr, socket).await?;
}
}
Some(Message_oneof_union::punch_hole_sent(phs)) => {
// punch hole sent from B, tell A that B ready
let addr_a = V4AddrMangle::decode(&phs.socket_addr);
let mut msg_out = Message::new(); let mut msg_out = Message::new();
msg_out.set_punch_hole_response(PunchHoleResponse { msg_out.set_punch_hole(PunchHole {
socket_addr: V4AddrMangle::encode(&addr_v4), socket_addr: AddrMangle::encode(&addr),
..Default::default() ..Default::default()
}); });
send_to(&msg_out, SocketAddr::V4(addr_a), socket).await?; send_to(&msg_out, peer.socket_addr, socket).await?;
} }
_ => {}
} }
Some(Message_oneof_union::punch_hole_sent(phs)) => {
// punch hole sent from B, tell A that B ready
let addr_a = AddrMangle::decode(&phs.socket_addr);
let mut msg_out = Message::new();
msg_out.set_punch_hole_response(PunchHoleResponse {
socket_addr: AddrMangle::encode(&addr),
..Default::default()
});
send_to(&msg_out, addr_a, socket).await?;
}
_ => {}
} }
} }
Ok(()) Ok(())
@@ -109,36 +98,77 @@ mod tests {
#[allow(unused_must_use)] #[allow(unused_must_use)]
#[tokio::main] #[tokio::main]
async fn test_rs_async() { async fn test_rs_async() {
let server_addr = "0.0.0.0:21116"; let mut port_server: u16 = 0;
let f1 = RendezvousServer::start(server_addr); let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let to_addr = server_addr.parse().unwrap(); if let SocketAddr::V4(addr) = socket.local_addr().unwrap() {
port_server = addr.port();
}
drop(socket);
let addr_server = format!("127.0.0.1:{}", port_server);
let f1 = RendezvousServer::start(&addr_server);
let addr_server = addr_server.parse().unwrap();
let f2 = async { let f2 = async {
let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); // B register it to server
let local_addr = socket.local_addr().unwrap(); let socket_b = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let mut socket = UdpFramed::new(socket, BytesCodec::new()); let local_addr_b = socket_b.local_addr().unwrap();
let mut socket_b = UdpFramed::new(socket_b, BytesCodec::new());
let mut msg_out = Message::new(); let mut msg_out = Message::new();
msg_out.set_register_peer(RegisterPeer { msg_out.set_register_peer(RegisterPeer {
hbb_addr: "123".to_string(), hbb_addr: "123".to_string(),
..Default::default() ..Default::default()
}); });
send_to(&msg_out, to_addr, &mut socket).await; send_to(&msg_out, addr_server, &mut socket_b).await;
// A send punch request to server
let socket_a = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let local_addr_a = socket_a.local_addr().unwrap();
let mut socket_a = UdpFramed::new(socket_a, BytesCodec::new());
msg_out.set_punch_hole_request(PunchHoleRequest { msg_out.set_punch_hole_request(PunchHoleRequest {
hbb_addr: "123".to_string(), hbb_addr: "123".to_string(),
..Default::default() ..Default::default()
}); });
send_to(&msg_out, to_addr, &mut socket).await; send_to(&msg_out, addr_server, &mut socket_a).await;
if let Ok(Some(Ok((bytes, _)))) =
tokio::time::timeout(Duration::from_millis(1), socket.next()).await println!(
"A {:?} request punch hole to B {:?} via server {:?}",
local_addr_a, local_addr_b, addr_server,
);
// on B side, responsed to A's punch request forwarded from server
if let Ok(Some(Ok((bytes, addr)))) =
tokio::time::timeout(Duration::from_millis(1000), socket_b.next()).await
{ {
if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) { assert_eq!(addr_server, addr);
assert_eq!( let msg_in = parse_from_bytes::<Message>(&bytes).unwrap();
local_addr, let remote_addr_a = AddrMangle::decode(&msg_in.get_punch_hole().socket_addr[..]);
SocketAddr::V4(V4AddrMangle::decode( assert_eq!(local_addr_a, remote_addr_a);
&msg_in.get_punch_hole_response().socket_addr[..]
)) // B punch A
); socket_b
} .get_mut()
.send_to(&b"SYN"[..], &remote_addr_a)
.await;
msg_out.set_punch_hole_sent(PunchHoleSent {
socket_addr: AddrMangle::encode(&remote_addr_a),
..Default::default()
});
send_to(&msg_out, addr_server, &mut socket_b).await;
} }
// on A side
socket_a.next().await; // skip "SYN"
if let Ok(Some(Ok((bytes, addr)))) =
tokio::time::timeout(Duration::from_millis(1000), socket_a.next()).await
{
assert_eq!(addr_server, addr);
let msg_in = parse_from_bytes::<Message>(&bytes).unwrap();
let remote_addr_b =
AddrMangle::decode(&msg_in.get_punch_hole_response().socket_addr[..]);
println!("{:?}", msg_in);
assert_eq!(local_addr_b, remote_addr_b);
}
if true { if true {
Err(Box::new(simple_error::SimpleError::new("done"))) Err(Box::new(simple_error::SimpleError::new("done")))
} else { } else {