From 3fdad5a40155e7621bd69270e15ec28abbbd464e Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Wed, 6 Oct 2021 02:29:47 +0800 Subject: [PATCH] reload config "servers" in runtime for sslocal - Trigger reload with signal SIGUSR1 - run_local is replaced by create_local then call run() on Server instance - DNS resolver reloader won't trigger multiple time when /etc/resolv.conf changed in 1 second --- Cargo.lock | 9 +- bin/sslocal.rs | 44 ++++- crates/shadowsocks-service/Cargo.toml | 1 + crates/shadowsocks-service/src/config.rs | 20 ++- crates/shadowsocks-service/src/lib.rs | 2 +- .../src/local/dns/server.rs | 2 +- .../src/local/loadbalancing/ping_balancer.rs | 159 ++++++++++++------ crates/shadowsocks-service/src/local/mod.rs | 33 +++- .../shadowsocks/src/dns_resolver/resolver.rs | 52 ++++-- 9 files changed, 236 insertions(+), 86 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 82a6b0ea..60a76418 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1246,9 +1246,9 @@ checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" [[package]] name = "quote" -version = "1.0.9" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3d0b9745dc2debf507c8422de05d7226cc1f0644216dfdfead988f9b1ab32a7" +checksum = "38bc8cc6a5f2e3655e0899c1b848643b2562f853f114bfec7be120678e3ace05" dependencies = [ "proc-macro2", ] @@ -1620,6 +1620,7 @@ dependencies = [ name = "shadowsocks-service" version = "1.12.0" dependencies = [ + "arc-swap 1.4.0", "async-trait", "byte_string", "byteorder", @@ -1989,9 +1990,9 @@ dependencies = [ [[package]] name = "tracing-attributes" -version = "0.1.17" +version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4f915eb6abf914599c200260efced9203504c4c37380af10cdf3b7d36970650" +checksum = "98863d0dd09fa59a1b79c6750ad80dbda6b75f4e71c437a6a1a8cb91a8bcbd77" dependencies = [ "proc-macro2", "quote", diff --git a/bin/sslocal.rs b/bin/sslocal.rs index eeae2ecd..0f7bee43 100644 --- a/bin/sslocal.rs +++ b/bin/sslocal.rs @@ -4,11 +4,11 @@ //! or you could specify a configuration file. The format of configuration file is defined //! in mod `config`. -use std::{net::IpAddr, process, time::Duration}; +use std::{net::IpAddr, path::PathBuf, process, time::Duration}; use clap::{clap_app, Arg, Error as ClapError, ErrorKind as ClapErrorKind}; use futures::future::{self, Either}; -use log::info; +use log::{error, info}; use tokio::{self, runtime::Builder}; #[cfg(feature = "local-redir")] @@ -18,7 +18,8 @@ use shadowsocks_service::shadowsocks::relay::socks5::Address; use shadowsocks_service::{ acl::AccessControl, config::{Config, ConfigType, LocalConfig, ProtocolType}, - run_local, + create_local, + local::loadbalancing::PingBalancer, shadowsocks::{ config::{Mode, ServerAddr, ServerConfig}, crypto::v1::{available_ciphers, CipherKind}, @@ -523,8 +524,16 @@ fn main() { }; runtime.block_on(async move { + let config_path = config.config_path.clone(); + + let instance = create_local(config).await.expect("create local"); + + if let Some(config_path) = config_path { + launch_reload_server_task(config_path, instance.server_balancer().clone()); + } + let abort_signal = monitor::create_signal_monitor(); - let server = run_local(config); + let server = instance.run(); tokio::pin!(abort_signal); tokio::pin!(server); @@ -545,3 +554,30 @@ fn main() { } }); } + +#[cfg(unix)] +fn launch_reload_server_task(config_path: PathBuf, balancer: PingBalancer) { + use tokio::signal::unix::{signal, SignalKind}; + + tokio::spawn(async move { + let mut sigusr1 = signal(SignalKind::user_defined1()).expect("signal"); + + while let Some(_) = sigusr1.recv().await { + let config = match Config::load_from_file(&config_path, ConfigType::Local) { + Ok(c) => c, + Err(err) => { + error!("auto-reload {} failed with error: {}", config_path.display(), err); + continue; + } + }; + + let servers = config.server; + info!("auto-reload {} with {} servers", config_path.display(), servers.len()); + + balancer.reset_servers(servers).await; + } + }); +} + +#[cfg(not(unix))] +fn launch_reload_server_task(_: PathBuf, _: PingBalancer) {} diff --git a/crates/shadowsocks-service/Cargo.toml b/crates/shadowsocks-service/Cargo.toml index 63756cd6..d3f7a181 100644 --- a/crates/shadowsocks-service/Cargo.toml +++ b/crates/shadowsocks-service/Cargo.toml @@ -80,6 +80,7 @@ qrcode = { version = "0.12", default-features = false } pin-project = "1.0" once_cell = "1.8" thiserror = "1.0" +arc-swap = "1.3" spin = { version = "0.9", features = ["std"] } lru_time_cache = "0.11" diff --git a/crates/shadowsocks-service/src/config.rs b/crates/shadowsocks-service/src/config.rs index 492a7aa3..08697c07 100644 --- a/crates/shadowsocks-service/src/config.rs +++ b/crates/shadowsocks-service/src/config.rs @@ -938,6 +938,10 @@ pub struct Config { /// Balancer config of local server pub balancer: BalancerConfig, + + /// Configuration file path, the actual path of the configuration. + /// This is normally for auto-reloading if implementation supports. + pub config_path: Option, } /// Configuration parsing error kind @@ -1044,6 +1048,8 @@ impl Config { security: SecurityConfig::default(), balancer: BalancerConfig::default(), + + config_path: None, } } @@ -1719,11 +1725,19 @@ impl Config { } /// Load Config from a File - pub fn load_from_file(filename: &str, config_type: ConfigType) -> Result { - let mut reader = OpenOptions::new().read(true).open(&Path::new(filename))?; + pub fn load_from_file>(filename: P, config_type: ConfigType) -> Result { + let filename = filename.as_ref(); + + let mut reader = OpenOptions::new().read(true).open(filename)?; let mut content = String::new(); reader.read_to_string(&mut content)?; - Config::load_from_str(&content[..], config_type) + + let mut config = Config::load_from_str(&content[..], config_type)?; + + // Record the path of the configuration for auto-reloading + config.config_path = Some(filename.to_owned()); + + Ok(config) } /// Check if there are any plugin are enabled with servers diff --git a/crates/shadowsocks-service/src/lib.rs b/crates/shadowsocks-service/src/lib.rs index 4ab5fc9e..875999db 100644 --- a/crates/shadowsocks-service/src/lib.rs +++ b/crates/shadowsocks-service/src/lib.rs @@ -51,7 +51,7 @@ use std::time::Duration; #[cfg(feature = "local")] -pub use self::local::run as run_local; +pub use self::local::create as create_local; #[cfg(feature = "manager")] pub use self::manager::run as run_manager; #[cfg(feature = "server")] diff --git a/crates/shadowsocks-service/src/local/dns/server.rs b/crates/shadowsocks-service/src/local/dns/server.rs index 663a29b7..d92f6dd7 100644 --- a/crates/shadowsocks-service/src/local/dns/server.rs +++ b/crates/shadowsocks-service/src/local/dns/server.rs @@ -328,7 +328,7 @@ fn should_forward_by_query(context: &ServiceContext, balancer: &PingBalancer, qu // // This happens normally because VPN or TUN device receives DNS queries from local servers' plugins // https://github.com/shadowsocks/shadowsocks-android/issues/2722 - for server in balancer.servers() { + for server in balancer.servers().as_ref() { let svr_cfg = server.server_config(); if let ServerAddr::DomainName(ref dn, ..) = svr_cfg.addr() { // Convert domain name to `Name` diff --git a/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs b/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs index 9cb3683b..37246066 100644 --- a/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs +++ b/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs @@ -12,6 +12,7 @@ use std::{ time::{Duration, Instant}, }; +use arc_swap::ArcSwap; use byte_string::ByteStr; use futures::future; use log::{debug, info, trace, warn}; @@ -86,15 +87,13 @@ impl PingBalancerBuilder { self.check_interval = intv; } - pub async fn build(self) -> PingBalancer { - assert!(!self.servers.is_empty(), "build PingBalancer without any servers"); - + fn find_best_idx(servers: &[Arc], mode: Mode) -> (usize, usize) { let mut best_tcp_idx = 0; let mut best_udp_idx = 0; - if self.mode.enable_tcp() { + if mode.enable_tcp() { let mut found_tcp_idx = false; - for (idx, server) in self.servers.iter().enumerate() { + for (idx, server) in servers.iter().enumerate() { if PingBalancerContext::check_server_tcp_enabled(server.server_config()) { best_tcp_idx = idx; found_tcp_idx = true; @@ -105,19 +104,19 @@ impl PingBalancerBuilder { if !found_tcp_idx { warn!( "no valid TCP server serving for TCP clients, consider disable TCP with \"mode\": \"udp_only\", currently chose {}", - ServerConfigFormatter::new(self.servers[best_tcp_idx].server_config()) + ServerConfigFormatter::new(servers[best_tcp_idx].server_config()) ); } else { trace!( "init chose TCP server {}", - ServerConfigFormatter::new(self.servers[best_tcp_idx].server_config()) + ServerConfigFormatter::new(servers[best_tcp_idx].server_config()) ); } } - if self.mode.enable_udp() { + if mode.enable_udp() { let mut found_udp_idx = false; - for (idx, server) in self.servers.iter().enumerate() { + for (idx, server) in servers.iter().enumerate() { if PingBalancerContext::check_server_udp_enabled(server.server_config()) { best_udp_idx = idx; found_udp_idx = true; @@ -128,18 +127,26 @@ impl PingBalancerBuilder { if !found_udp_idx { warn!( "no valid UDP server serving for UDP clients, consider disable UDP with \"mode\": \"tcp_only\", currently chose {}", - ServerConfigFormatter::new(self.servers[best_udp_idx].server_config()) + ServerConfigFormatter::new(servers[best_udp_idx].server_config()) ); } else { trace!( "init chose UDP server {}", - ServerConfigFormatter::new(self.servers[best_udp_idx].server_config()) + ServerConfigFormatter::new(servers[best_udp_idx].server_config()) ); } } + (best_tcp_idx, best_udp_idx) + } + + pub async fn build(self) -> PingBalancer { + assert!(!self.servers.is_empty(), "build PingBalancer without any servers"); + + let (best_tcp_idx, best_udp_idx) = PingBalancerBuilder::find_best_idx(&self.servers, self.mode); + let balancer_context = PingBalancerContext { - servers: self.servers, + servers: ArcSwap::new(Arc::new(self.servers)), best_tcp_idx: AtomicUsize::new(best_tcp_idx), best_udp_idx: AtomicUsize::new(best_udp_idx), context: self.context, @@ -152,7 +159,7 @@ impl PingBalancerBuilder { let shared_context = Arc::new(balancer_context); - let abortable = { + let checker_abortable = { let shared_context = shared_context.clone(); tokio::spawn(async move { shared_context.checker_task().await }) }; @@ -160,14 +167,14 @@ impl PingBalancerBuilder { PingBalancer { inner: Arc::new(PingBalancerInner { context: shared_context, - abortable, + checker_abortable, }), } } } struct PingBalancerContext { - servers: Vec>, + servers: ArcSwap>>, best_tcp_idx: AtomicUsize, best_udp_idx: AtomicUsize, context: Arc, @@ -178,21 +185,64 @@ struct PingBalancerContext { impl PingBalancerContext { fn best_tcp_server(&self) -> Arc { - self.servers[self.best_tcp_idx.load(Ordering::Relaxed)].clone() + let servers = self.servers.load(); + + // A guard if reset_servers is running + let mut best_tcp_idx = self.best_tcp_idx.load(Ordering::Relaxed); + if best_tcp_idx >= servers.len() { + best_tcp_idx = 0; + } + + servers[best_tcp_idx].clone() } fn best_udp_server(&self) -> Arc { - self.servers[self.best_udp_idx.load(Ordering::Relaxed)].clone() + let servers = self.servers.load(); + + // A guard if reset_servers is running + let mut best_udp_idx = self.best_udp_idx.load(Ordering::Relaxed); + if best_udp_idx >= servers.len() { + best_udp_idx = 0; + } + + servers[best_udp_idx].clone() } } impl PingBalancerContext { - async fn init_score(&self) { - assert!(!self.servers.is_empty(), "check PingBalancer without any servers"); + pub async fn reset_servers(&self, servers: Vec) { + let max_server_rtt = self.max_server_rtt; + let servers = servers + .into_iter() + .map(|s| Arc::new(ServerIdent::new(s, max_server_rtt))) + .collect(); + self.reset_servers_ident(servers).await + } - if self.probing_required() { - self.check_once(true).await; - } + pub async fn reset_servers_ident(&self, servers: Vec>) { + assert!(!servers.is_empty(), "servers shouldn't be empty"); + + // Restore best_tcp_idx and best_udp_idx + let (best_tcp_idx, best_udp_idx) = PingBalancerBuilder::find_best_idx(&servers, self.mode); + + // Create a new Arc servers + let new_servers = Arc::new(servers); + + // Replace into context and then run the checker task once + self.servers.store(new_servers); + self.best_tcp_idx.store(best_tcp_idx, Ordering::Release); + self.best_udp_idx.store(best_udp_idx, Ordering::Release); + + self.check_once(true).await; + } + + async fn init_score(&self) { + assert!( + !self.servers.load().is_empty(), + "check PingBalancer without any servers" + ); + + self.check_once(true).await; } fn check_server_tcp_enabled(svr_cfg: &ServerConfig) -> bool { @@ -203,11 +253,11 @@ impl PingBalancerContext { svr_cfg.mode().enable_udp() && svr_cfg.weight().udp_weight() > 0.0 } - fn probing_required(&self) -> bool { + fn probing_required(&self, servers: &[Arc]) -> bool { let mut tcp_count = 0; let mut udp_count = 0; - for server in self.servers.iter() { + for server in servers.iter() { let svr_cfg = server.server_config(); if self.mode.enable_tcp() && PingBalancerContext::check_server_tcp_enabled(svr_cfg) { tcp_count += 1; @@ -221,26 +271,25 @@ impl PingBalancerContext { } async fn checker_task(self: Arc) { - assert!(!self.servers.is_empty(), "check PingBalancer without any servers"); + assert!( + !self.servers.load().is_empty(), + "check PingBalancer without any servers" + ); - if !self.probing_required() { - self.checker_task_dummy().await - } else { - self.checker_task_real().await - } - } - - /// Dummy task that will do nothing if there only have one server in the balancer - async fn checker_task_dummy(self: Arc) { - future::pending().await + self.checker_task_real().await } /// Check each servers' score and update the best server's index async fn check_once(&self, first_run: bool) { - let mut vfut_tcp = Vec::with_capacity(self.servers.len()); - let mut vfut_udp = Vec::with_capacity(self.servers.len()); + let servers = self.servers.load(); + if !self.probing_required(&*servers) { + return; + } - for server in self.servers.iter() { + let mut vfut_tcp = Vec::with_capacity(servers.len()); + let mut vfut_udp = Vec::with_capacity(servers.len()); + + for server in servers.iter() { let svr_cfg = server.server_config(); if self.mode.enable_tcp() && PingBalancerContext::check_server_tcp_enabled(svr_cfg) { @@ -287,7 +336,7 @@ impl PingBalancerContext { let mut best_idx = 0; let mut best_score = u32::MAX; - for (idx, server) in self.servers.iter().enumerate() { + for (idx, server) in servers.iter().enumerate() { let score = server.tcp_score().score(); if score < best_score { best_idx = idx; @@ -299,19 +348,19 @@ impl PingBalancerContext { if first_run { info!( "chose best TCP server {}", - ServerConfigFormatter::new(self.servers[best_idx].server_config()) + ServerConfigFormatter::new(servers[best_idx].server_config()) ); } else { if best_idx != old_best_idx { info!( "switched best TCP server from {} to {}", - ServerConfigFormatter::new(self.servers[old_best_idx].server_config()), - ServerConfigFormatter::new(self.servers[best_idx].server_config()) + ServerConfigFormatter::new(servers[old_best_idx].server_config()), + ServerConfigFormatter::new(servers[best_idx].server_config()) ); } else { debug!( "kept best TCP server {}", - ServerConfigFormatter::new(self.servers[old_best_idx].server_config()) + ServerConfigFormatter::new(servers[old_best_idx].server_config()) ); } } @@ -322,7 +371,7 @@ impl PingBalancerContext { let mut best_idx = 0; let mut best_score = u32::MAX; - for (idx, server) in self.servers.iter().enumerate() { + for (idx, server) in servers.iter().enumerate() { let score = server.udp_score().score(); if score < best_score { best_idx = idx; @@ -334,19 +383,19 @@ impl PingBalancerContext { if first_run { info!( "chose best UDP server {}", - ServerConfigFormatter::new(self.servers[best_idx].server_config()) + ServerConfigFormatter::new(servers[best_idx].server_config()) ); } else { if best_idx != old_best_idx { info!( "switched best UDP server from {} to {}", - ServerConfigFormatter::new(self.servers[old_best_idx].server_config()), - ServerConfigFormatter::new(self.servers[best_idx].server_config()) + ServerConfigFormatter::new(servers[old_best_idx].server_config()), + ServerConfigFormatter::new(servers[best_idx].server_config()) ); } else { debug!( "kept best UDP server {}", - ServerConfigFormatter::new(self.servers[old_best_idx].server_config()) + ServerConfigFormatter::new(servers[old_best_idx].server_config()) ); } } @@ -366,12 +415,12 @@ impl PingBalancerContext { struct PingBalancerInner { context: Arc, - abortable: JoinHandle<()>, + checker_abortable: JoinHandle<()>, } impl Drop for PingBalancerInner { fn drop(&mut self) { - self.abortable.abort(); + self.checker_abortable.abort(); trace!("ping balancer stopped"); } } @@ -404,8 +453,14 @@ impl PingBalancer { } /// Get the server list - pub fn servers(&self) -> &[Arc] { - &self.inner.context.servers + pub fn servers(&self) -> Arc>> { + self.inner.context.servers.load().clone() + } + + /// Reset servers in load balancer. Designed for auto-reloading configuration file. + #[inline] + pub async fn reset_servers(&self, servers: Vec) { + self.inner.context.reset_servers(servers).await } } diff --git a/crates/shadowsocks-service/src/local/mod.rs b/crates/shadowsocks-service/src/local/mod.rs index 2af86239..7fc4fb4d 100644 --- a/crates/shadowsocks-service/src/local/mod.rs +++ b/crates/shadowsocks-service/src/local/mod.rs @@ -9,7 +9,7 @@ use std::{ }; use futures::{ - future, + future::{self, BoxFuture}, stream::{FuturesUnordered, StreamExt}, FutureExt, }; @@ -27,7 +27,10 @@ use crate::{ dns::build_dns_resolver, }; -use self::{context::ServiceContext, loadbalancing::PingBalancerBuilder}; +use self::{ + context::ServiceContext, + loadbalancing::{PingBalancer, PingBalancerBuilder}, +}; pub mod context; #[cfg(feature = "local-dns")] @@ -50,8 +53,26 @@ pub mod utils; /// This is borrowed from Go's `net` library's default setting pub(crate) const LOCAL_DEFAULT_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(15); +pub struct Server { + vfut: FuturesUnordered>>, + balancer: PingBalancer, +} + +impl Server { + /// Run local server + pub async fn run(self) -> io::Result<()> { + let (res, _) = self.vfut.into_future().await; + res.unwrap() + } + + /// Get the internal server balancer + pub fn server_balancer(&self) -> &PingBalancer { + &self.balancer + } +} + /// Starts a shadowsocks local server -pub async fn run(mut config: Config) -> io::Result<()> { +pub async fn create(mut config: Config) -> io::Result { assert!(config.config_type == ConfigType::Local && !config.local.is_empty()); assert!(!config.server.is_empty()); @@ -406,8 +427,10 @@ pub async fn run(mut config: Config) -> io::Result<()> { } // let (res, ..) = future::select_all(vfut).await; - let (res, _) = vfut.into_future().await; - res.unwrap() + // let (res, _) = vfut.into_future().await; + // res.unwrap() + + Ok(Server { vfut, balancer }) } #[cfg(feature = "local-flow-stat")] diff --git a/crates/shadowsocks/src/dns_resolver/resolver.rs b/crates/shadowsocks/src/dns_resolver/resolver.rs index 5ee60ae3..390d64a8 100644 --- a/crates/shadowsocks/src/dns_resolver/resolver.rs +++ b/crates/shadowsocks/src/dns_resolver/resolver.rs @@ -6,7 +6,7 @@ use std::{ fmt::{self, Debug}, io::{self, Error, ErrorKind}, net::SocketAddr, - time::Instant, + time::{Duration, Instant}, }; #[cfg(feature = "trust-dns")] @@ -16,9 +16,9 @@ use cfg_if::cfg_if; #[cfg(feature = "trust-dns")] use log::error; use log::{log_enabled, trace, Level}; -use tokio::net::lookup_host; #[cfg(feature = "trust-dns")] use tokio::task::JoinHandle; +use tokio::{net::lookup_host, time}; #[cfg(feature = "trust-dns")] use trust_dns_resolver::{config::ResolverConfig, TokioAsyncResolver}; @@ -143,42 +143,62 @@ async fn trust_dns_notify_update_dns(resolver: Arc) -> n use super::trust_dns_resolver::create_resolver; + static DNS_RESOLV_FILE_PATH: &str = "/etc/resolv.conf"; + let (tx, mut rx) = watch::channel::(Event::default()); let mut watcher: RecommendedWatcher = notify::recommended_watcher(move |ev_result: NotifyResult| match ev_result { Ok(ev) => { - trace!("received event {:?}", ev); + trace!("received {} event {:?}", DNS_RESOLV_FILE_PATH, ev); if let EventKind::Modify(..) = ev.kind { tx.send(ev).expect("watcher.send"); } } Err(err) => { - error!("watching /etc/resolv.conf error: {}", err); + error!("watching {} error: {}", DNS_RESOLV_FILE_PATH, err); } })?; // NOTE: It is an undefined behavior if this file get renamed or removed. - watcher.watch(Path::new("/etc/resolv.conf"), RecursiveMode::NonRecursive)?; + watcher.watch(Path::new(DNS_RESOLV_FILE_PATH), RecursiveMode::NonRecursive)?; + + // Delayed task + let mut update_task: Option> = None; while rx.changed().await.is_ok() { - trace!("received notify /etc/resolv.conf changed"); + trace!("received notify {} changed", DNS_RESOLV_FILE_PATH); - let new_resolver = match create_resolver(None, resolver.ipv6_first).await { - Ok(r) => r, - Err(err) => { - error!("failed to reload /etc/resolv.conf, error: {}", err); - continue; - } + // Kill the pending task + if let Some(t) = update_task.take() { + t.abort(); + } + + let task = { + let resolver = resolver.clone(); + tokio::spawn(async move { + // /etc/resolv.conf may be modified multiple time in 1 second + // Update once for all those Modify events + time::sleep(Duration::from_secs(1)).await; + + match create_resolver(None, resolver.ipv6_first).await { + Ok(r) => { + debug!("auto-reload {}", DNS_RESOLV_FILE_PATH); + + resolver.resolver.store(Arc::new(r)); + } + Err(err) => { + error!("failed to reload {}, error: {}", DNS_RESOLV_FILE_PATH, err); + } + } + }) }; - debug!("auto-reload /etc/resolv.conf"); - - resolver.resolver.store(Arc::new(new_resolver)); + update_task = Some(task); } - error!("auto-reload /etc/resolv.conf task exited unexpectly"); + error!("auto-reload {} task exited unexpectly", DNS_RESOLV_FILE_PATH); Ok(()) }