weyne/rustdesk
1
0
Fork 0
mirror of https://github.com/weyne85/rustdesk.git synced 2025-10-29 17:00:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

lan discovery will be done soon

Browse Source 
Signed-off-by: fufesou <shuanglongchen@yeah.net>
This commit is contained in:
fufesou
2021-12-19 19:58:08 +08:00
parent bcbe9ccbe5
commit 5682b088de
11 changed files with 254 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
libs/hbb_common/protos/discovery.proto
View File

@@ -8,3 +8,8 @@ message Discovery {
/// response port for current listening port(udp for now)
int32 port = 2;
}
message DiscoveryBack {
string id = 1;
base.PeerInfo peer = 2;
}

2
libs/hbb_common/src/config.rs
View File

@@ -55,6 +55,8 @@ pub const RENDEZVOUS_SERVERS: &'static [&'static str] = &[
pub const RENDEZVOUS_PORT: i32 = 21116;
pub const RELAY_PORT: i32 = 21117;
pub const SERVER_UDP_PORT: u16 = 21001; // udp
#[derive(Debug, Default, Serialize, Deserialize, Clone)]
pub struct Config {
#[serde(default)]

33
libs/socket_cs/examples/discovery.rs
View File

@@ -1,13 +1,11 @@
use hbb_common::{
base_proto::PeerInfo,
discovery_proto::Discovery as DiscoveryProto,
discovery_proto::{Discovery as DiscoveryProto, DiscoveryBack as DiscoveryBackProto},
env_logger::*,
log, protobuf,
tokio::{self, sync::Notify},
log, protobuf, tokio,
};
use socket_cs::{discovery::*, udp::*};
use std::env;
use std::sync::Arc;
async fn lan_discover(port: u16, port_back: u16) {
let peer = PeerInfo {
@@ -27,21 +25,21 @@ async fn lan_discover(port: u16, port_back: u16) {
}
async fn listen_discovery_back(port: u16) {
fn proc_peer(peer: PeerInfo) {
fn proc_peer(info: DiscoveryBackProto) {
log::info!(
"peer recived, username: {}, hostname: {}",
peer.username,
peer.hostname
"peer recived, id: {}, username: {}, hostname: {}",
info.id,
info.peer.as_ref().unwrap().username,
info.peer.as_ref().unwrap().hostname
);
}
let exit_notify = Notify::new();
let handlers = UdpHandlers::new().handle(
CMD_DISCOVERY_BACK.as_bytes().to_vec(),
Box::new(HandlerDiscoveryBack::new(proc_peer)),
);
let server = Server::new(port, Arc::new(exit_notify));
let mut server = Server::create(port).unwrap();
server.start(handlers).await.unwrap();
loop {
@@ -50,19 +48,22 @@ async fn listen_discovery_back(port: u16) {
}
async fn listen_discovery(port: u16) {
let peer = PeerInfo {
username: "server username".to_owned(),
hostname: "server hostname".to_owned(),
let info = DiscoveryBackProto {
id: "server id".to_owned(),
peer: protobuf::MessageField::from_option(Some(PeerInfo {
username: "server username".to_owned(),
hostname: "server hostname".to_owned(),
..Default::default()
})),
..Default::default()
};
let exit_notify = Notify::new();
let handlers = UdpHandlers::new().handle(
CMD_DISCOVERY.as_bytes().to_vec(),
Box::new(HandlerDiscovery::new(peer)),
Box::new(HandlerDiscovery::new(Some(|| true), info)),
);
let server = Server::new(port, Arc::new(exit_notify));
let mut server = Server::create(port).unwrap();
server.start(handlers).await.unwrap();
loop {
std::thread::sleep(std::time::Duration::from_millis(1000));

35
libs/socket_cs/src/discovery.rs
View File

@@ -1,15 +1,17 @@
use super::udp::UdpRequest;
use async_trait::async_trait;
use hbb_common::{
base_proto::PeerInfo, discovery_proto::Discovery as DiscoveryProto, log, protobuf::Message,
tokio::net::UdpSocket, ResultType,
discovery_proto::{Discovery as DiscoveryProto, DiscoveryBack as DiscoveryBackProto},
log,
protobuf::Message,
tokio::net::UdpSocket,
ResultType,
};
use std::net::SocketAddr;
pub const CMD_DISCOVERY: &str = "discovery";
pub const CMD_DISCOVERY_BACK: &str = "discovery_back";
// TODO: make sure if UdpFramed is needed, or UdpSocket just works fine.
pub struct DiscoveryClient {
socket: UdpSocket,
send_data: Vec<u8>,
@@ -44,13 +46,17 @@ impl DiscoveryClient {
}
pub struct HandlerDiscovery {
get_allow: Option<fn() -> bool>,
send_data: Vec<u8>,
}
impl HandlerDiscovery {
pub fn new(self_info: PeerInfo) -> Self {
pub fn new(get_allow: Option<fn() -> bool>, self_info: DiscoveryBackProto) -> Self {
let send_data = make_send_data(CMD_DISCOVERY_BACK, &self_info).unwrap();
Self { send_data }
Self {
get_allow,
send_data,
}
}
}
@@ -67,6 +73,17 @@ impl crate::Handler<UdpRequest> for HandlerDiscovery {
peer.hostname
);
let allowed = self.get_allow.map_or(false, |f| f());
if !allowed {
log::info!(
"received discovery query from {} {} {}, but discovery is not allowed",
request.addr,
peer.hostname,
peer.username
);
return Ok(());
}
let addr = "0.0.0.0:0";
let socket = UdpSocket::bind(addr).await?;
@@ -89,11 +106,11 @@ impl crate::Handler<UdpRequest> for HandlerDiscovery {
}
pub struct HandlerDiscoveryBack {
proc: fn(peer_info: PeerInfo),
proc: fn(info: DiscoveryBackProto),
}
impl HandlerDiscoveryBack {
pub fn new(proc: fn(peer_info: PeerInfo)) -> Self {
pub fn new(proc: fn(info: DiscoveryBackProto)) -> Self {
Self { proc }
}
}
@@ -103,8 +120,8 @@ impl crate::Handler<UdpRequest> for HandlerDiscoveryBack {
async fn call(&self, request: UdpRequest) -> ResultType<()> {
log::trace!("recved discover back from {}", request.addr);
let peer = PeerInfo::parse_from_bytes(&request.data)?;
(self.proc)(peer);
let info = DiscoveryBackProto::parse_from_bytes(&request.data)?;
(self.proc)(info);
Ok(())
}
}

42
libs/socket_cs/src/udp.rs
View File

@@ -1,7 +1,7 @@
use async_trait::async_trait;
use hbb_common::{
log,
tokio::{self, sync::Notify},
tokio::{self, runtime::Runtime, sync::Notify, task::JoinHandle},
udp::FramedSocket,
ResultType,
};
@@ -14,6 +14,8 @@ use std::sync::Arc;
pub struct Server {
port: u16,
exit_notify: Arc<Notify>,
rt: Arc<Runtime>,
join_handler: Option<JoinHandle<()>>,
}
pub struct UdpRequest {
@@ -33,19 +35,27 @@ pub struct UdpHandlers {
}
impl Server {
pub fn new(port: u16, exit_notify: Arc<Notify>) -> Self {
Self { port, exit_notify }
pub fn create(port: u16) -> ResultType<Self> {
let rt = Arc::new(Runtime::new()?);
let exit_notify = Arc::new(Notify::new());
Ok(Self {
port,
exit_notify,
rt,
join_handler: None,
})
}
/// Start server with the handlers.
pub async fn start(&self, handlers: UdpHandlers) -> ResultType<()> {
pub async fn start(&mut self, handlers: UdpHandlers) -> ResultType<()> {
let exit_notify = self.exit_notify.clone();
let addr = SocketAddr::from(([0, 0, 0, 0], self.port));
let mut server = FramedSocket::new(addr).await?;
log::trace!("succeeded to bind {} for discovery server", addr);
tokio::spawn(async move {
let rt = self.rt.clone();
let join_handler = rt.clone().spawn(async move {
let handlers = Arc::new(handlers.handlers);
loop {
tokio::select! {
@@ -56,11 +66,12 @@ impl Server {
n = server.next() => {
log::info!("received message");
let handlers = handlers.clone();
let rt = rt.clone();
match n {
Some(Ok((data, addr))) => {
match data.iter().position(|x| x == &crate::CMD_TOKEN) {
Some(p) => {
tokio::spawn(async move {
rt.spawn(async move {
let cmd = data[0..p].to_vec();
match handlers.get(&cmd) {
Some(h) => {
@@ -92,8 +103,27 @@ impl Server {
}
}
});
self.join_handler = Some(join_handler);
Ok(())
}
pub async fn shutdonw(&mut self) {
self.exit_notify.notify_one();
if let Some(h) = self.join_handler.take() {
if let Err(e) = h.await {
log::error!("failed to join udp server, {}", e);
}
}
}
}
impl Drop for Server {
fn drop(&mut self) {
self.rt.clone().block_on(async {
self.shutdonw().await;
})
}
}
impl UdpHandlers {