commit
						b400e4305f
					
				
							
								
								
									
										10
									
								
								Cargo.lock
									
									
									
										generated
									
									
									
								
							
							
						
						
									
										10
									
								
								Cargo.lock
									
									
									
										generated
									
									
									
								
							| @ -3169,6 +3169,7 @@ dependencies = [ | ||||
|  "serde_derive", | ||||
|  "serde_json 1.0.74", | ||||
|  "sha2", | ||||
|  "socket_cs", | ||||
|  "sys-locale", | ||||
|  "systray", | ||||
|  "uuid", | ||||
| @ -3455,11 +3456,12 @@ dependencies = [ | ||||
| ] | ||||
| 
 | ||||
| [[package]] | ||||
| name = "socket2" | ||||
| version = "0.4.2" | ||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "5dc90fe6c7be1a323296982db1836d1ea9e47b6839496dde9a541bc496df3516" | ||||
| name = "socket_cs" | ||||
| version = "0.1.0" | ||||
| dependencies = [ | ||||
|  "async-trait", | ||||
|  "clap", | ||||
|  "hbb_common", | ||||
|  "libc", | ||||
|  "winapi 0.3.9", | ||||
| ] | ||||
|  | ||||
| @ -20,6 +20,7 @@ default = ["use_dasp"] | ||||
| whoami = "1.2" | ||||
| scrap = { path = "libs/scrap" } | ||||
| hbb_common = { path = "libs/hbb_common" } | ||||
| socket_cs = { path = "libs/socket_cs" } | ||||
| enigo = { path = "libs/enigo" } | ||||
| sys-locale = "0.1" | ||||
| serde_derive = "1.0" | ||||
| @ -82,7 +83,7 @@ psutil = { version = "3.2", features = [ "process" ], git = "https://github.com/ | ||||
| android_logger = "0.10" | ||||
| 
 | ||||
| [workspace] | ||||
| members = ["libs/scrap", "libs/hbb_common", "libs/enigo"] | ||||
| members = ["libs/scrap", "libs/hbb_common", "libs/enigo", "libs/socket_cs"] | ||||
| 
 | ||||
| [package.metadata.winres] | ||||
| LegalCopyright = "Copyright © 2020" | ||||
|  | ||||
| @ -2,7 +2,7 @@ fn main() { | ||||
|     std::fs::create_dir_all("src/protos").unwrap(); | ||||
|     protobuf_codegen_pure::Codegen::new() | ||||
|         .out_dir("src/protos") | ||||
|         .inputs(&["protos/rendezvous.proto", "protos/message.proto"]) | ||||
|         .inputs(&["protos/rendezvous.proto", "protos/base_proto.proto", "protos/message.proto", "protos/discovery.proto"]) | ||||
|         .include("protos") | ||||
|         .run() | ||||
|         .expect("Codegen failed."); | ||||
|  | ||||
							
								
								
									
										21
									
								
								libs/hbb_common/protos/base_proto.proto
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										21
									
								
								libs/hbb_common/protos/base_proto.proto
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,21 @@ | ||||
| syntax = "proto3"; | ||||
| package base; | ||||
| 
 | ||||
| message DisplayInfo { | ||||
|   sint32 x = 1; | ||||
|   sint32 y = 2; | ||||
|   int32 width = 3; | ||||
|   int32 height = 4; | ||||
|   string name = 5; | ||||
|   bool online = 6; | ||||
| } | ||||
| 
 | ||||
| message PeerInfo { | ||||
|   string username = 1; | ||||
|   string hostname = 2; | ||||
|   string platform = 3; | ||||
|   repeated DisplayInfo displays = 4; | ||||
|   int32 current_display = 5; | ||||
|   bool sas_enabled = 6; | ||||
|   string version = 7; | ||||
| } | ||||
							
								
								
									
										16
									
								
								libs/hbb_common/protos/discovery.proto
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										16
									
								
								libs/hbb_common/protos/discovery.proto
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,16 @@ | ||||
| syntax = "proto3"; | ||||
| package discovery; | ||||
| 
 | ||||
| import "base_proto.proto"; | ||||
| 
 | ||||
| message Discovery { | ||||
|   string id = 1; | ||||
|   base.PeerInfo peer = 2; | ||||
|   /// response port for current listening port(udp for now) | ||||
|   int32 port = 3; | ||||
| } | ||||
| 
 | ||||
| message DiscoveryBack { | ||||
|   string id = 1; | ||||
|   base.PeerInfo peer = 2; | ||||
| } | ||||
| @ -1,6 +1,8 @@ | ||||
| syntax = "proto3"; | ||||
| package hbb; | ||||
| 
 | ||||
| import "base_proto.proto"; | ||||
| 
 | ||||
| message VP9 { | ||||
|   bytes data = 1; | ||||
|   bool key = 2; | ||||
| @ -25,15 +27,6 @@ message VideoFrame { | ||||
|   } | ||||
| } | ||||
| 
 | ||||
| message DisplayInfo { | ||||
|   sint32 x = 1; | ||||
|   sint32 y = 2; | ||||
|   int32 width = 3; | ||||
|   int32 height = 4; | ||||
|   string name = 5; | ||||
|   bool online = 6; | ||||
| } | ||||
| 
 | ||||
| message PortForward { | ||||
|   string host = 1; | ||||
|   int32 port = 2; | ||||
| @ -58,20 +51,10 @@ message LoginRequest { | ||||
| 
 | ||||
| message ChatMessage { string text = 1; } | ||||
| 
 | ||||
| message PeerInfo { | ||||
|   string username = 1; | ||||
|   string hostname = 2; | ||||
|   string platform = 3; | ||||
|   repeated DisplayInfo displays = 4; | ||||
|   int32 current_display = 5; | ||||
|   bool sas_enabled = 6; | ||||
|   string version = 7; | ||||
| } | ||||
| 
 | ||||
| message LoginResponse { | ||||
|   oneof union { | ||||
|     string error = 1; | ||||
|     PeerInfo peer_info = 2; | ||||
|     base.PeerInfo peer_info = 2; | ||||
|   } | ||||
| } | ||||
| 
 | ||||
|  | ||||
| @ -1,4 +1,8 @@ | ||||
| pub mod compress; | ||||
| #[path = "./protos/base_proto.rs"] | ||||
| pub mod base_proto; | ||||
| #[path = "./protos/discovery.rs"] | ||||
| pub mod discovery_proto; | ||||
| #[path = "./protos/message.rs"] | ||||
| pub mod message_proto; | ||||
| #[path = "./protos/rendezvous.rs"] | ||||
|  | ||||
							
								
								
									
										14
									
								
								libs/socket_cs/Cargo.toml
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										14
									
								
								libs/socket_cs/Cargo.toml
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,14 @@ | ||||
| [package] | ||||
| name = "socket_cs" | ||||
| version = "0.1.0" | ||||
| edition = "2021" | ||||
| 
 | ||||
| # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html | ||||
| 
 | ||||
| [dependencies] | ||||
| hbb_common = { path = "../hbb_common" } | ||||
| async-trait = "0.1" | ||||
| 
 | ||||
| 
 | ||||
| [dev-dependencies] | ||||
| clap = "2.33" | ||||
							
								
								
									
										90
									
								
								libs/socket_cs/examples/discovery.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										90
									
								
								libs/socket_cs/examples/discovery.rs
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,90 @@ | ||||
| use hbb_common::{ | ||||
|     base_proto::PeerInfo, | ||||
|     discovery_proto::{Discovery as DiscoveryProto, DiscoveryBack as DiscoveryBackProto}, | ||||
|     env_logger::*, | ||||
|     log, protobuf, tokio, | ||||
| }; | ||||
| use socket_cs::{discovery::*, udp::*}; | ||||
| use std::env; | ||||
| 
 | ||||
| async fn lan_discover(port: u16, port_back: u16) { | ||||
|     let peer = PeerInfo { | ||||
|         username: "client username".to_owned(), | ||||
|         hostname: "client hostname".to_owned(), | ||||
|         ..Default::default() | ||||
|     }; | ||||
|     let client = DiscoveryClient::create(DiscoveryProto { | ||||
|         id: "client id".to_owned(), | ||||
|         peer: protobuf::MessageField::from_option(Some(peer)), | ||||
|         port: port_back as i32, | ||||
|         ..Default::default() | ||||
|     }) | ||||
|     .await | ||||
|     .unwrap(); | ||||
| 
 | ||||
|     client.lan_discover(port).await.unwrap(); | ||||
| } | ||||
| 
 | ||||
| async fn listen_discovery_back(port: u16) { | ||||
|     fn proc_peer(info: DiscoveryBackProto) { | ||||
|         log::info!( | ||||
|             "peer recived, id: {}, username: {}, hostname: {}", | ||||
|             info.id, | ||||
|             info.peer.as_ref().unwrap().username, | ||||
|             info.peer.as_ref().unwrap().hostname | ||||
|         ); | ||||
|     } | ||||
| 
 | ||||
|     let handlers = UdpHandlers::new().handle( | ||||
|         CMD_DISCOVERY_BACK.as_bytes().to_vec(), | ||||
|         Box::new(HandlerDiscoveryBack::new(proc_peer)), | ||||
|     ); | ||||
| 
 | ||||
|     let mut server = Server::create(port).unwrap(); | ||||
|     server.start(handlers).await.unwrap(); | ||||
| 
 | ||||
|     loop { | ||||
|         std::thread::sleep(std::time::Duration::from_millis(1000)); | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| async fn listen_discovery(port: u16) { | ||||
|     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 handlers = UdpHandlers::new().handle( | ||||
|         CMD_DISCOVERY.as_bytes().to_vec(), | ||||
|         Box::new(HandlerDiscovery::new(Some(|| true), info)), | ||||
|     ); | ||||
| 
 | ||||
|     let mut server = Server::create(port).unwrap(); | ||||
|     server.start(handlers).await.unwrap(); | ||||
|     loop { | ||||
|         std::thread::sleep(std::time::Duration::from_millis(1000)); | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| #[tokio::main] | ||||
| async fn main() { | ||||
|     init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "trace")); | ||||
| 
 | ||||
|     let args: Vec<String> = env::args().collect(); | ||||
| 
 | ||||
|     let port_back = 9801u16; | ||||
|     let server_port: u16 = 9802u16; | ||||
| 
 | ||||
|     if args.len() == 1 { | ||||
|         lan_discover(server_port, port_back).await; | ||||
|     } else if args.len() == 2 { | ||||
|         listen_discovery_back(port_back).await; | ||||
|     } else { | ||||
|         listen_discovery(server_port).await; | ||||
|     } | ||||
| } | ||||
							
								
								
									
										131
									
								
								libs/socket_cs/src/discovery.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										131
									
								
								libs/socket_cs/src/discovery.rs
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,131 @@ | ||||
| use super::udp::UdpRequest; | ||||
| use async_trait::async_trait; | ||||
| use hbb_common::{ | ||||
|     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"; | ||||
| 
 | ||||
| pub struct DiscoveryClient { | ||||
|     socket: UdpSocket, | ||||
|     send_data: Vec<u8>, | ||||
| } | ||||
| 
 | ||||
| fn make_send_data(cmd: &str, msg: &impl Message) -> ResultType<Vec<u8>> { | ||||
|     let info_bytes = msg.write_to_bytes()?; | ||||
|     let mut send_data = cmd.as_bytes().to_vec(); | ||||
|     send_data.push(crate::CMD_TOKEN); | ||||
|     send_data.extend(info_bytes); | ||||
|     Ok(send_data) | ||||
| } | ||||
| 
 | ||||
| impl DiscoveryClient { | ||||
|     pub async fn create(info: DiscoveryProto) -> ResultType<Self> { | ||||
|         let addr = "0.0.0.0:0"; | ||||
|         let socket = UdpSocket::bind(addr).await?; | ||||
|         log::trace!("succeeded to bind {} for discovery client", addr); | ||||
| 
 | ||||
|         socket.set_broadcast(true)?; | ||||
|         log::trace!("Broadcast mode set to ON"); | ||||
| 
 | ||||
|         let send_data = make_send_data(CMD_DISCOVERY, &info)?; | ||||
|         Ok(Self { socket, send_data }) | ||||
|     } | ||||
| 
 | ||||
|     pub async fn lan_discover(&self, peer_port: u16) -> ResultType<()> { | ||||
|         let addr = SocketAddr::from(([255, 255, 255, 255], peer_port)); | ||||
|         self.socket.send_to(&self.send_data, addr).await?; | ||||
|         Ok(()) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| pub struct HandlerDiscovery { | ||||
|     get_allow: Option<fn() -> bool>, | ||||
|     id: String, | ||||
|     send_data: Vec<u8>, | ||||
| } | ||||
| 
 | ||||
| impl HandlerDiscovery { | ||||
|     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 { | ||||
|             get_allow, | ||||
|             id: self_info.id, | ||||
|             send_data, | ||||
|         } | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| #[async_trait] | ||||
| impl crate::Handler<UdpRequest> for HandlerDiscovery { | ||||
|     async fn call(&self, request: UdpRequest) -> ResultType<()> { | ||||
|         let discovery = DiscoveryProto::parse_from_bytes(&request.data)?; | ||||
|         if discovery.id == self.id { | ||||
|             return Ok(()); | ||||
|         } | ||||
| 
 | ||||
|         let peer = discovery.peer.as_ref().take().unwrap(); | ||||
|         log::trace!( | ||||
|             "received discovery query from {} {}", | ||||
|             peer.username, | ||||
|             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?; | ||||
| 
 | ||||
|         let mut peer_addr = request.addr; | ||||
|         peer_addr.set_port(discovery.port as u16); | ||||
|         log::trace!("send self peer info to {}", peer_addr); | ||||
| 
 | ||||
|         let send_len = self.send_data.len(); | ||||
|         let mut cur_len = 0usize; | ||||
|         while cur_len < send_len { | ||||
|             let len = socket | ||||
|                 .send_to(&self.send_data[cur_len..], peer_addr) | ||||
|                 .await?; | ||||
|             cur_len += len; | ||||
|         } | ||||
|         log::trace!("send self peer info to {} done", peer_addr); | ||||
| 
 | ||||
|         Ok(()) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| pub struct HandlerDiscoveryBack { | ||||
|     proc: fn(info: DiscoveryBackProto), | ||||
| } | ||||
| 
 | ||||
| impl HandlerDiscoveryBack { | ||||
|     pub fn new(proc: fn(info: DiscoveryBackProto)) -> Self { | ||||
|         Self { proc } | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| #[async_trait] | ||||
| impl crate::Handler<UdpRequest> for HandlerDiscoveryBack { | ||||
|     async fn call(&self, request: UdpRequest) -> ResultType<()> { | ||||
|         log::trace!("recved discover back from {}", request.addr); | ||||
| 
 | ||||
|         let info = DiscoveryBackProto::parse_from_bytes(&request.data)?; | ||||
|         (self.proc)(info); | ||||
|         Ok(()) | ||||
|     } | ||||
| } | ||||
							
								
								
									
										12
									
								
								libs/socket_cs/src/lib.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										12
									
								
								libs/socket_cs/src/lib.rs
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,12 @@ | ||||
| use async_trait::async_trait; | ||||
| pub use hbb_common::ResultType; | ||||
| pub mod discovery; | ||||
| pub mod udp; | ||||
| 
 | ||||
| const CMD_TOKEN: u8 = '\n' as u8; | ||||
| 
 | ||||
| /// Use tower::Service may be better ?
 | ||||
| #[async_trait] | ||||
| pub trait Handler<Request>: Send + Sync { | ||||
|     async fn call(&self, request: Request) -> ResultType<()>; | ||||
| } | ||||
							
								
								
									
										187
									
								
								libs/socket_cs/src/udp.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										187
									
								
								libs/socket_cs/src/udp.rs
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,187 @@ | ||||
| use async_trait::async_trait; | ||||
| use hbb_common::{ | ||||
|     log, | ||||
|     tokio::{self, runtime::Runtime, sync::Notify, task::JoinHandle}, | ||||
|     udp::FramedSocket, | ||||
|     ResultType, | ||||
| }; | ||||
| use std::{collections::HashMap, future::Future, net::SocketAddr, sync::Arc}; | ||||
| 
 | ||||
| /// Simple udp server
 | ||||
| pub struct Server { | ||||
|     port: u16, | ||||
|     exit_notify: Arc<Notify>, | ||||
|     rt: Arc<Runtime>, | ||||
|     join_handler: Option<JoinHandle<()>>, | ||||
| } | ||||
| 
 | ||||
| pub struct UdpRequest { | ||||
|     pub data: Vec<u8>, | ||||
|     pub addr: SocketAddr, | ||||
| } | ||||
| 
 | ||||
| type UdpHandler = Box<dyn crate::Handler<UdpRequest>>; | ||||
| 
 | ||||
| pub struct UdpFnHandler<F>(F); | ||||
| 
 | ||||
| /// Handlers of udp server.
 | ||||
| /// After udp server received data. Command should be parsed.
 | ||||
| /// Handler will then be used to process data.
 | ||||
| pub struct UdpHandlers { | ||||
|     handlers: HashMap<Vec<u8>, UdpHandler>, | ||||
| } | ||||
| 
 | ||||
| impl Server { | ||||
|     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(&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); | ||||
| 
 | ||||
|         let rt = self.rt.clone(); | ||||
|         let join_handler = rt.clone().spawn(async move { | ||||
|             let handlers = Arc::new(handlers.handlers); | ||||
|             loop { | ||||
|                 tokio::select! { | ||||
|                     _ = exit_notify.notified() => { | ||||
|                         log::debug!("exit server graceful"); | ||||
|                         break; | ||||
|                     } | ||||
|                     n = server.next() => { | ||||
|                         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) => { | ||||
|                                         rt.spawn(async move { | ||||
|                                             let cmd = data[0..p].to_vec(); | ||||
|                                             match handlers.get(&cmd) { | ||||
|                                                 Some(h) => { | ||||
|                                                     let request = UdpRequest {data: data[p+1..].to_vec(), addr}; | ||||
|                                                     if let Err(_e) = h.call(request).await { | ||||
|                                                         // log::error!("handle {:?} failed, {}", cmd, e);
 | ||||
|                                                     } | ||||
|                                                 } | ||||
|                                                 None => { | ||||
|                                                     // log::warn!("no handler for {:?}", &cmd);
 | ||||
|                                                 } | ||||
|                                             } | ||||
|                                         }); | ||||
|                                     } | ||||
|                                     None => { | ||||
|                                         // log::error!("failed to parse command token");
 | ||||
|                                     } | ||||
|                                 } | ||||
| 
 | ||||
|                             } | ||||
|                             Some(Err(_e)) => { | ||||
|                                 // log::error!("recv error: {}", e)
 | ||||
|                             } | ||||
|                             None => { | ||||
|                                 log::error!("should never reach here"); | ||||
|                             } | ||||
|                         } | ||||
|                     } | ||||
|                 } | ||||
|             } | ||||
|         }); | ||||
| 
 | ||||
|         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 { | ||||
|     pub fn new() -> Self { | ||||
|         Self { | ||||
|             handlers: HashMap::new(), | ||||
|         } | ||||
|     } | ||||
|     /// Insert <cmd, handler> pair.
 | ||||
|     ///
 | ||||
|     /// # Example
 | ||||
|     ///
 | ||||
|     /// ```rust
 | ||||
|     /// extern crate socket_cs;
 | ||||
|     /// use socket_cs::{ResultType, udp::{UdpHandlers, UdpRequest}};
 | ||||
|     /// use async_trait::async_trait;
 | ||||
|     ///
 | ||||
|     /// struct SimpleHandler;
 | ||||
|     ///
 | ||||
|     /// #[async_trait]
 | ||||
|     /// impl socket_cs::Handler<UdpRequest> for SimpleHandler {
 | ||||
|     ///     async fn call(&self, _: UdpRequest) -> ResultType<()> {
 | ||||
|     ///         Ok(())
 | ||||
|     ///     }
 | ||||
|     /// }
 | ||||
|     /// async fn simple_ignore(_: UdpRequest) -> ResultType<()> {
 | ||||
|     ///     Ok(())
 | ||||
|     /// }
 | ||||
|     /// let handlers = UdpHandlers::new();
 | ||||
|     ///
 | ||||
|     /// handlers
 | ||||
|     /// .handle(b"cmd".to_vec(), Box::new(SimpleHandler))
 | ||||
|     /// .handle(b"cmd2".to_vec(), simple_ignore.into());
 | ||||
|     ///
 | ||||
|     /// ```
 | ||||
|     ///
 | ||||
|     /// **Notice** Same cmd where override the previous one.
 | ||||
|     ///
 | ||||
|     pub fn handle(mut self, cmd: Vec<u8>, h: UdpHandler) -> Self { | ||||
|         self.handlers.insert(cmd, h); | ||||
|         self | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| /// TODO: more generice Request.
 | ||||
| #[async_trait] | ||||
| impl<F, Fut> crate::Handler<UdpRequest> for UdpFnHandler<F> | ||||
| where | ||||
|     Fut: Future<Output = ResultType<()>> + Send, | ||||
|     F: Fn(UdpRequest) -> Fut + Send + Sync, | ||||
| { | ||||
|     async fn call(&self, request: UdpRequest) -> ResultType<()> { | ||||
|         self.0(request).await | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| impl<F, Fut> From<F> for UdpHandler | ||||
| where | ||||
|     Fut: Future<Output = ResultType<()>> + Send, | ||||
|     F: Fn(UdpRequest) -> Fut + Send + Sync + 'static, | ||||
| { | ||||
|     fn from(f: F) -> Self { | ||||
|         Box::new(UdpFnHandler(f)) | ||||
|     } | ||||
| } | ||||
| @ -10,6 +10,7 @@ use hbb_common::{ | ||||
|     bail, | ||||
|     config::{Config, PeerConfig, PeerInfoSerde, CONNECT_TIMEOUT, RELAY_PORT, RENDEZVOUS_TIMEOUT}, | ||||
|     log, | ||||
|     base_proto::*, | ||||
|     message_proto::*, | ||||
|     protobuf::Message as _, | ||||
|     rendezvous_proto::*, | ||||
|  | ||||
							
								
								
									
										17
									
								
								src/ipc.rs
									
									
									
									
									
								
							
							
						
						
									
										17
									
								
								src/ipc.rs
									
									
									
									
									
								
							| @ -91,6 +91,7 @@ pub enum Data { | ||||
|     RawMessage(Vec<u8>), | ||||
|     Socks(Option<config::Socks5Server>), | ||||
|     FS(FS), | ||||
|     SessionsUpdated, | ||||
|     Test, | ||||
| } | ||||
| 
 | ||||
| @ -406,6 +407,22 @@ pub fn get_id() -> String { | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| pub async fn get_id_async() -> String { | ||||
|     if let Ok(Some(v)) = get_config_async("id", 1_000).await { | ||||
|         // update salt also, so that next time reinstallation not causing first-time auto-login failure
 | ||||
|         if let Ok(Some(v2)) = get_config_async("salt", 1_000).await { | ||||
|             Config::set_salt(&v2); | ||||
|         } | ||||
|         if v != Config::get_id() { | ||||
|             Config::set_key_confirmed(false); | ||||
|             Config::set_id(&v); | ||||
|         } | ||||
|         v | ||||
|     } else { | ||||
|         Config::get_id() | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| pub fn get_password() -> String { | ||||
|     if let Ok(Some(v)) = get_config("password") { | ||||
|         Config::set_password(&v); | ||||
|  | ||||
| @ -7,6 +7,7 @@ use hbb_common::{ | ||||
|     config::{Config, CONNECT_TIMEOUT, RELAY_PORT}, | ||||
|     log, | ||||
|     message_proto::*, | ||||
|     base_proto::*, | ||||
|     protobuf::{Message as _, ProtobufEnum}, | ||||
|     rendezvous_proto::*, | ||||
|     sleep, | ||||
| @ -27,6 +28,7 @@ mod connection; | ||||
| pub mod input_service; | ||||
| mod service; | ||||
| mod video_service; | ||||
| pub mod udp; | ||||
| 
 | ||||
| use hbb_common::tcp::new_listener; | ||||
| 
 | ||||
| @ -263,6 +265,7 @@ pub fn check_zombie() { | ||||
| 
 | ||||
| #[tokio::main] | ||||
| pub async fn start_server(is_server: bool, _tray: bool) { | ||||
|     // TODO: Add a graceful shutdown handler, and attach all servers to that handler.
 | ||||
|     #[cfg(target_os = "linux")] | ||||
|     { | ||||
|         log::info!("DISPLAY={:?}", std::env::var("DISPLAY")); | ||||
| @ -275,6 +278,13 @@ pub async fn start_server(is_server: bool, _tray: bool) { | ||||
|                 std::process::exit(-1); | ||||
|             } | ||||
|         }); | ||||
|         let _server_guard = match udp::start_udp_server().await { | ||||
|             Ok(s) => Some(s), | ||||
|             Err(e) => { | ||||
|                 log::warn!("Failed to start udp server {}", e); | ||||
|                 None | ||||
|             } | ||||
|         }; | ||||
|         input_service::fix_key_down_timeout_loop(); | ||||
|         crate::RendezvousMediator::start_all().await; | ||||
|     } else { | ||||
|  | ||||
							
								
								
									
										125
									
								
								src/server/udp.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										125
									
								
								src/server/udp.rs
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,125 @@ | ||||
| /// udp server
 | ||||
| ///
 | ||||
| /// eg. discovery
 | ||||
| ///
 | ||||
| use hbb_common::{base_proto::PeerInfo, config::SERVER_UDP_PORT, ResultType}; | ||||
| use socket_cs::udp::{Server, UdpHandlers}; | ||||
| 
 | ||||
| /// Simple copy from ../connections.rs#send_logon_response
 | ||||
| /// Should be merged into one function.
 | ||||
| fn get_peer_info() -> PeerInfo { | ||||
|     let username = crate::platform::get_active_username(); | ||||
|     #[allow(unused_mut)] | ||||
|     let mut sas_enabled = false; | ||||
|     #[cfg(windows)] | ||||
|     if crate::platform::is_root() { | ||||
|         sas_enabled = true; | ||||
|     } | ||||
|     PeerInfo { | ||||
|         hostname: whoami::hostname(), | ||||
|         username, | ||||
|         platform: whoami::platform().to_string(), | ||||
|         version: crate::VERSION.to_owned(), | ||||
|         sas_enabled, | ||||
|         ..Default::default() | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| pub mod discovery { | ||||
|     use super::get_peer_info; | ||||
|     use crate::ipc; | ||||
|     use hbb_common::{ | ||||
|         config::{PeerConfig, PeerInfoSerde, SERVER_UDP_PORT}, | ||||
|         discovery_proto::{Discovery as DiscoveryProto, DiscoveryBack as DiscoveryBackProto}, | ||||
|         log, protobuf, tokio, ResultType, | ||||
|     }; | ||||
|     use socket_cs::{discovery::*, udp::UdpHandlers}; | ||||
| 
 | ||||
|     async fn get_discovery_back_info() -> DiscoveryBackProto { | ||||
|         let peer = get_peer_info(); | ||||
|         DiscoveryBackProto { | ||||
|             id: ipc::get_id_async().await, | ||||
|             peer: protobuf::MessageField::from_option(Some(peer)), | ||||
|             ..Default::default() | ||||
|         } | ||||
|     } | ||||
| 
 | ||||
|     /// process sicovery bakc(response)
 | ||||
|     /// 1. update current peers.
 | ||||
|     /// 2. notify index window to udpate recent sessions.
 | ||||
|     fn process_discovery_back(info: DiscoveryBackProto) { | ||||
|         let mut config = PeerConfig::load(info.id.as_str()); | ||||
| 
 | ||||
|         let peer = info.peer.as_ref().unwrap(); | ||||
|         let serde = PeerInfoSerde { | ||||
|             username: peer.username.clone(), | ||||
|             hostname: peer.hostname.clone(), | ||||
|             platform: peer.platform.clone(), | ||||
|         }; | ||||
|         config.info = serde; | ||||
|         config.store(info.id.as_str()); | ||||
| 
 | ||||
|         #[tokio::main(flavor = "current_thread")] | ||||
|         async fn notify_index_window() -> ResultType<()> { | ||||
|             let ms_timeout = 300; | ||||
|             let mut c = ipc::connect(ms_timeout, "_index").await?; | ||||
|             c.send(&ipc::Data::SessionsUpdated).await?; | ||||
|             Ok(()) | ||||
|         } | ||||
|         std::thread::spawn(move || { | ||||
|             if let Err(e) = notify_index_window() { | ||||
|                 log::error!("Failed to notify index window, {}", e); | ||||
|             } | ||||
|         }); | ||||
|     } | ||||
| 
 | ||||
|     /// launch lan discover when user click "discover" button.
 | ||||
|     pub fn launch_lan_discover() { | ||||
|         std::thread::spawn(move || { | ||||
|             if let Err(e) = lan_discover() { | ||||
|                 log::error!("Failed to lauch lan discover, {}", e); | ||||
|             } | ||||
|         }); | ||||
|     } | ||||
| 
 | ||||
|     #[tokio::main(flavor = "current_thread")] | ||||
|     pub async fn lan_discover() -> ResultType<()> { | ||||
|         let peer = get_peer_info(); | ||||
|         let client = DiscoveryClient::create(DiscoveryProto { | ||||
|             id: ipc::get_id_async().await, | ||||
|             peer: protobuf::MessageField::from_option(Some(peer)), | ||||
|             port: SERVER_UDP_PORT as i32, | ||||
|             ..Default::default() | ||||
|         }) | ||||
|         .await?; | ||||
| 
 | ||||
|         client.lan_discover(SERVER_UDP_PORT).await | ||||
|     } | ||||
| 
 | ||||
|     pub(super) async fn handle_discovery(handlers: UdpHandlers) -> UdpHandlers { | ||||
|         let info = get_discovery_back_info().await; | ||||
|         handlers | ||||
|             // handle discover request
 | ||||
|             .handle( | ||||
|                 CMD_DISCOVERY.as_bytes().to_vec(), | ||||
|                 Box::new(HandlerDiscovery::new( | ||||
|                     // Some(|| Config::get_option("enable-be-discovered") == "Y".to_owned()),
 | ||||
|                     Some(|| true), | ||||
|                     info, | ||||
|                 )), | ||||
|             ) | ||||
|             // handle discover back(response)
 | ||||
|             .handle( | ||||
|                 CMD_DISCOVERY_BACK.as_bytes().to_vec(), | ||||
|                 Box::new(HandlerDiscoveryBack::new(process_discovery_back)), | ||||
|             ) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| pub(super) async fn start_udp_server() -> ResultType<Server> { | ||||
|     let handlers = discovery::handle_discovery(UdpHandlers::new()).await; | ||||
| 
 | ||||
|     let mut server = Server::create(SERVER_UDP_PORT)?; | ||||
|     server.start(handlers).await?; | ||||
|     Ok(server) | ||||
| } | ||||
							
								
								
									
										56
									
								
								src/ui.rs
									
									
									
									
									
								
							
							
						
						
									
										56
									
								
								src/ui.rs
									
									
									
									
									
								
							| @ -86,6 +86,8 @@ pub fn start(args: &mut [String]) { | ||||
|         let childs: Childs = Default::default(); | ||||
|         let cloned = childs.clone(); | ||||
|         std::thread::spawn(move || check_zombie(cloned)); | ||||
|         let cloned = childs.clone(); | ||||
|         std::thread::spawn(move || start_ipc(cloned)); | ||||
|         crate::common::check_software_update(); | ||||
|         frame.event_handler(UI::new(childs)); | ||||
|         frame.sciter_handler(UIHostHandler {}); | ||||
| @ -626,6 +628,10 @@ impl UI { | ||||
|     fn is_xfce(&self) -> bool { | ||||
|         crate::platform::is_xfce() | ||||
|     } | ||||
| 
 | ||||
|     fn lan_discover(&self) { | ||||
|         crate::server::udp::discovery::launch_lan_discover(); | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| impl sciter::EventHandler for UI { | ||||
| @ -683,6 +689,7 @@ impl sciter::EventHandler for UI { | ||||
|         fn get_software_ext(); | ||||
|         fn open_url(String); | ||||
|         fn create_shortcut(String); | ||||
|         fn lan_discover(); | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| @ -714,6 +721,55 @@ pub fn check_zombie(childs: Childs) { | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| // TODO: Duplicated code.
 | ||||
| // Need more generic and better shutdown handler
 | ||||
| #[tokio::main(flavor = "current_thread")] | ||||
| async fn start_ipc(childs: Childs) { | ||||
|     match ipc::new_listener("_index").await { | ||||
|         Ok(mut incoming) => { | ||||
|             while let Some(result) = incoming.next().await { | ||||
|                 match result { | ||||
|                     Ok(stream) => { | ||||
|                         let mut stream = ipc::Connection::new(stream); | ||||
|                         let childs = childs.clone(); | ||||
|                         tokio::spawn(async move { | ||||
|                             loop { | ||||
|                                 tokio::select! { | ||||
|                                     res = stream.next() => { | ||||
|                                         match res { | ||||
|                                             Err(err) => { | ||||
|                                                 log::info!("cm ipc connection closed: {}", err); | ||||
|                                                 break; | ||||
|                                             } | ||||
|                                             Ok(Some(data)) => { | ||||
|                                                 match data { | ||||
|                                                     ipc::Data::SessionsUpdated => { | ||||
|                                                         childs.lock().unwrap().0 = true; | ||||
|                                                     } | ||||
|                                                     _ => { | ||||
|                                                     } | ||||
|                                                 } | ||||
|                                             } | ||||
|                                             _ => {} | ||||
|                                         } | ||||
|                                     } | ||||
|                                 } | ||||
|                             } | ||||
|                         }); | ||||
|                     } | ||||
|                     Err(err) => { | ||||
|                         log::error!("Couldn't get index client: {:?}", err); | ||||
|                     } | ||||
|                 } | ||||
|             } | ||||
|         } | ||||
|         Err(err) => { | ||||
|             log::error!("Failed to start index ipc server: {}", err); | ||||
|         } | ||||
|     } | ||||
|     std::process::exit(-1); | ||||
| } | ||||
| 
 | ||||
| // notice: avoiding create ipc connection repeatedly,
 | ||||
| // because windows named pipe has serious memory leak issue.
 | ||||
| #[tokio::main(flavor = "current_thread")] | ||||
|  | ||||
| @ -45,6 +45,10 @@ class ConnectStatus: Reactor.Component { | ||||
|         return translate("Ready"); | ||||
|     } | ||||
| 
 | ||||
|     event click $(.connect-status .link) () { | ||||
|     	handler.set_option("stop-service", ""); | ||||
|     } | ||||
| 
 | ||||
|     event click $(#start-service) () { | ||||
|         handler.set_option("stop-service", ""); | ||||
|     } | ||||
|  | ||||
| @ -7,6 +7,7 @@ use hbb_common::{ | ||||
|     allow_err, | ||||
|     config::{self, Config, PeerConfig}, | ||||
|     fs, log, | ||||
|     base_proto::*, | ||||
|     message_proto::*, | ||||
|     protobuf::Message as _, | ||||
|     rendezvous_proto::ConnType, | ||||
|  | ||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user