use clap::App; use hbb_common::{ allow_err, anyhow::{Context, Result}, get_version_number, log, tokio, ResultType }; use ini::Ini; use sodiumoxide::crypto::sign; use std::{ io::prelude::*, io::Read, net::{IpAddr, Ipv4Addr, SocketAddr}, time::{Instant, SystemTime}, }; pub fn parse_bind_address(value: &str) -> Result> { let value = value.trim(); if value.is_empty() { Ok(None) } else { value .parse() .with_context(|| format!("Invalid bind address: {value}")) .map(Some) } } pub async fn listen_tcp( bind_addr: Option, port: u16, ) -> ResultType { if let Some(bind_addr) = bind_addr { hbb_common::tcp::new_listener(SocketAddr::new(bind_addr, port), true).await } else { hbb_common::tcp::listen_any(port).await } } pub fn console_addr(bind_addr: Option, port: u16) -> Option { let bind_addr = bind_addr?; if bind_addr.is_unspecified() || bind_addr == IpAddr::V4(Ipv4Addr::LOCALHOST) { return None; } Some(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port)) } // The runtime console (check_cmd) is reached via 127.0.0.1, so when the bind // address does not already accept connections to 127.0.0.1 (it is neither the // any-address nor 127.0.0.1 itself), the console gets a dedicated listener // there; it is never bound to the external bind address. pub async fn listen_console( bind_addr: Option, port: u16, ) -> ResultType> { match console_addr(bind_addr, port) { Some(addr) => { let listener = hbb_common::tcp::new_listener(addr, true).await?; log::info!("Listening on tcp {} for the console", addr); Ok(Some(listener)) } None => Ok(None), } } pub async fn accept_or_pending( listener: Option<&hbb_common::tokio::net::TcpListener>, ) -> std::io::Result<(hbb_common::tokio::net::TcpStream, SocketAddr)> { match listener { Some(listener) => listener.accept().await, None => std::future::pending().await, } } #[allow(dead_code)] pub(crate) fn get_expired_time() -> Instant { let now = Instant::now(); now.checked_sub(std::time::Duration::from_secs(3600)) .unwrap_or(now) } #[allow(dead_code)] pub(crate) fn test_if_valid_server(host: &str, name: &str) -> ResultType { use std::net::ToSocketAddrs; let res = if host.contains(':') { host.to_socket_addrs()?.next().context("") } else { format!("{}:{}", host, 0) .to_socket_addrs()? .next() .context("") }; if res.is_err() { log::error!("Invalid {} {}: {:?}", name, host, res); } res } #[allow(dead_code)] pub(crate) fn get_servers(s: &str, tag: &str) -> Vec { let servers: Vec = s .split(',') .filter(|x| !x.is_empty() && test_if_valid_server(x, tag).is_ok()) .map(|x| x.to_owned()) .collect(); log::info!("{}={:?}", tag, servers); servers } #[allow(dead_code)] #[inline] fn arg_name(name: &str) -> String { name.to_uppercase().replace('_', "-") } #[allow(dead_code)] #[inline] pub fn set_arg(name: &str, value: &str) { std::env::set_var(arg_name(name), value); } #[allow(dead_code)] pub fn init_args(args: &str, name: &str, about: &str) { let matches = App::new(name) .version(crate::version::VERSION) .author("Purslane Ltd. ") .about(about) .args_from_usage(args) .get_matches(); if let Ok(v) = Ini::load_from_file(".env") { if let Some(section) = v.section(None::) { section .iter() .for_each(|(k, v)| set_arg(k, v)); } } if let Some(config) = matches.value_of("config") { if let Ok(v) = Ini::load_from_file(config) { if let Some(section) = v.section(None::) { section .iter() .for_each(|(k, v)| set_arg(k, v)); } } } for (k, v) in matches.args { if let Some(v) = v.vals.first() { set_arg(k, &v.to_string_lossy()); } } } #[allow(dead_code)] pub fn get_arg_opt(name: &str) -> Option { let dashed = arg_name(name); let underscored = dashed.replace('-', "_"); let lower_dashed = dashed.to_lowercase(); let lower_underscored = underscored.to_lowercase(); for alias in [&dashed, &underscored, &lower_dashed, &lower_underscored] { if let Ok(value) = std::env::var(alias) { return Some(value); } } let mut aliases = std::env::vars_os() .filter_map(|(key, value)| { let key = key.into_string().ok()?; if arg_name(&key) == dashed { Some((key, value.into_string().ok()?)) } else { None } }) .collect::>(); aliases.sort_by(|a, b| a.0.cmp(&b.0)); aliases.into_iter().next().map(|(_, value)| value) } #[allow(dead_code)] #[inline] pub fn get_arg(name: &str) -> String { get_arg_or(name, "".to_owned()) } #[allow(dead_code)] #[inline] pub fn get_arg_or(name: &str, default: String) -> String { get_arg_opt(name).unwrap_or(default) } #[allow(dead_code)] #[inline] pub fn now() -> u64 { SystemTime::now() .duration_since(SystemTime::UNIX_EPOCH) .map(|x| x.as_secs()) .unwrap_or_default() } pub fn gen_sk(wait: u64) -> (String, Option) { let sk_file = "id_ed25519"; if wait > 0 && !std::path::Path::new(sk_file).exists() { std::thread::sleep(std::time::Duration::from_millis(wait)); } if let Ok(mut file) = std::fs::File::open(sk_file) { let mut contents = String::new(); if file.read_to_string(&mut contents).is_ok() { let contents = contents.trim(); let sk = base64::decode(contents).unwrap_or_default(); if sk.len() == sign::SECRETKEYBYTES { let mut tmp = [0u8; sign::SECRETKEYBYTES]; tmp[..].copy_from_slice(&sk); let pk = base64::encode(&tmp[sign::SECRETKEYBYTES / 2..]); log::info!("Private key comes from {}", sk_file); return (pk, Some(sign::SecretKey(tmp))); } else { // don't use log here, since it is async println!("Fatal error: malformed private key in {sk_file}."); std::process::exit(1); } } } else { let gen_func = || { let (tmp, sk) = sign::gen_keypair(); (base64::encode(tmp), sk) }; let (mut pk, mut sk) = gen_func(); for _ in 0..300 { if !pk.contains('/') && !pk.contains(':') { break; } (pk, sk) = gen_func(); } let pub_file = format!("{sk_file}.pub"); if let Ok(mut f) = std::fs::File::create(&pub_file) { f.write_all(pk.as_bytes()).ok(); if let Ok(mut f) = std::fs::File::create(sk_file) { let s = base64::encode(&sk); if f.write_all(s.as_bytes()).is_ok() { log::info!("Private/public key written to {}/{}", sk_file, pub_file); log::debug!("Public key: {}", pk); return (pk, Some(sk)); } } } } ("".to_owned(), None) } #[cfg(unix)] pub async fn listen_signal() -> Result<()> { use hbb_common::tokio; use hbb_common::tokio::signal::unix::{signal, SignalKind}; tokio::spawn(async { let mut s = signal(SignalKind::terminate())?; let terminate = s.recv(); let mut s = signal(SignalKind::interrupt())?; let interrupt = s.recv(); let mut s = signal(SignalKind::quit())?; let quit = s.recv(); tokio::select! { _ = terminate => { log::info!("signal terminate"); } _ = interrupt => { log::info!("signal interrupt"); } _ = quit => { log::info!("signal quit"); } } Ok(()) }) .await? } #[cfg(not(unix))] pub async fn listen_signal() -> Result<()> { let () = std::future::pending().await; unreachable!(); } pub fn check_software_update() { const ONE_DAY_IN_SECONDS: u64 = 60 * 60 * 24; std::thread::spawn(move || loop { std::thread::spawn(move || allow_err!(check_software_update_())); std::thread::sleep(std::time::Duration::from_secs(ONE_DAY_IN_SECONDS)); }); } #[tokio::main(flavor = "current_thread")] async fn check_software_update_() -> hbb_common::ResultType<()> { let (request, url) = hbb_common::version_check_request(hbb_common::VER_TYPE_RUSTDESK_SERVER.to_string()); let latest_release_response = reqwest::Client::builder().build()? .post(url) .json(&request) .send() .await?; let bytes = latest_release_response.bytes().await?; let resp: hbb_common::VersionCheckResponse = serde_json::from_slice(&bytes)?; let response_url = resp.url; let latest_release_version = response_url.rsplit('/').next().unwrap_or_default(); if get_version_number(&latest_release_version) > get_version_number(crate::version::VERSION) { log::info!("new version is available: {}", latest_release_version); } Ok(()) } #[cfg(test)] mod tests { use super::*; use std::net::{Ipv4Addr, Ipv6Addr}; #[test] fn argument_names_ignore_case_and_separator() { let aliases = [ "RUSTDESK-CONFIG-ALIAS-TEST", "RUSTDESK_CONFIG_ALIAS_TEST", "rustdesk-config-alias-test", "rustdesk_config_alias_test", "RustDesk_Config-Alias_Test", ]; for alias in aliases { std::env::remove_var(alias); } for alias in aliases { std::env::set_var(alias, alias); assert_eq!(get_arg("RUSTDESK_CONFIG_ALIAS_TEST"), alias); std::env::remove_var(alias); } set_arg("rustdesk_config_alias_test", "normalized"); assert_eq!( std::env::var("RUSTDESK-CONFIG-ALIAS-TEST").unwrap(), "normalized" ); std::env::set_var("RUSTDESK_CONFIG_ALIAS_TEST", "inherited"); set_arg("rustdesk-config-alias-test", "higher-priority"); assert_eq!(get_arg("rustdesk_config_alias_test"), "higher-priority"); std::env::remove_var("RUSTDESK-CONFIG-ALIAS-TEST"); std::env::remove_var("RUSTDESK_CONFIG_ALIAS_TEST"); } #[test] fn parses_bind_address() { assert_eq!(parse_bind_address("").unwrap(), None); assert_eq!( parse_bind_address("127.0.0.1").unwrap(), Some(IpAddr::V4(Ipv4Addr::LOCALHOST)) ); assert_eq!( parse_bind_address("::1").unwrap(), Some(IpAddr::V6(Ipv6Addr::LOCALHOST)) ); assert!(parse_bind_address("not-an-ip").is_err()); } #[hbb_common::tokio::test] async fn tcp_listener_uses_bind_address() { let bind_addr = IpAddr::V4(Ipv4Addr::LOCALHOST); let listener = listen_tcp(Some(bind_addr), 0).await.unwrap(); assert_eq!(listener.local_addr().unwrap().ip(), bind_addr); } #[test] fn console_addr_only_when_bind_does_not_cover_ipv4_localhost() { for bind_addr in [ None, Some(IpAddr::V4(Ipv4Addr::UNSPECIFIED)), Some(IpAddr::V6(Ipv6Addr::UNSPECIFIED)), Some(IpAddr::V4(Ipv4Addr::LOCALHOST)), ] { assert_eq!(console_addr(bind_addr, 21117), None); } for bind_addr in [ Some(IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1))), Some("2001:db8::1".parse().unwrap()), Some(IpAddr::V6(Ipv6Addr::LOCALHOST)), ] { assert_eq!( console_addr(bind_addr, 21117), Some(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 21117)) ); } } #[hbb_common::tokio::test] async fn console_listener_binds_ipv4_localhost() { let listener = listen_console(Some(IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1))), 0) .await .unwrap() .unwrap(); assert_eq!( listener.local_addr().unwrap().ip(), IpAddr::V4(Ipv4Addr::LOCALHOST) ); assert!(listen_console(None, 0).await.unwrap().is_none()); assert!(listen_console(Some(IpAddr::V4(Ipv4Addr::LOCALHOST)), 0) .await .unwrap() .is_none()); } }