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

2
libs/hbb_common

@ -1 +1 @@
Subproject commit a9536bcc127a9d8057a9b7ffbe486506a7b963de
Subproject commit 974815baaaa2efd886be73c3be8fb5931bcc29c4

140
src/rendezvous_server.rs
View File

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