diff --git a/crates/vmate-cli/src/commands/all.rs b/crates/vmate-cli/src/commands/all.rs index 4ef06c0..d98ae94 100644 --- a/crates/vmate-cli/src/commands/all.rs +++ b/crates/vmate-cli/src/commands/all.rs @@ -3,115 +3,32 @@ use crate::cli::AllArgs; use crate::commands::connect::{persist_connect_defaults, resolve_connect}; use crate::settings::Settings; -use crate::ui::progress::{ProgressReporter, VerboseReporter}; use anyhow::Result; use clap_verbosity_flag::Verbosity; use std::sync::Arc; -use tokio_util::sync::CancellationToken; use vmate_core::connect::{Candidate, ConnectOptions, ConnectQueue, ConnectService}; -use vmate_core::db::ConfigRepo; -use vmate_core::db::pool::init_pool; -use vmate_core::geo::IpInfoGeoLocator; -use vmate_core::ovpn::process::{RealOpenVpnRunner, RealVpnTester, VpnTester}; -use vmate_core::scan::{ScanOptions, ScanProgress, ScanService}; +use vmate_core::ovpn::process::RealOpenVpnRunner; use vmate_core::settings::UserSettings; -use vmate_core::system::{ - CleanupGuard, ProcessKiller, RealProcessKiller, require_root_for, shutdown_signal, -}; +use vmate_core::system::{ProcessKiller, RealProcessKiller, require_root_for}; pub async fn run(settings: &Settings, args: &AllArgs, verbose: &Verbosity) -> Result<()> { // --save-defaults is a pure settings operation: persist and exit without // scanning or connecting (no root, no OpenVPN, no DB needed). if settings.save_defaults { crate::commands::scan::persist_scan_defaults(&args.scan)?; - let mut us = UserSettings::load(); - persist_connect_defaults(&mut us, &args.connect)?; + persist_connect_defaults(&args.connect)?; return Ok(()); } require_root_for("run OpenVPN tests and connections", settings.no_elevate)?; let us = UserSettings::load(); - let (workers, limit, timeout) = crate::commands::scan::resolve_scan_defaults(&args.scan, &us); let connect = resolve_connect(&us, &args.connect); - let dir = match &args.scan.dir { - Some(dir) => dir.clone(), - None => crate::commands::scan::materialize_builtins(&args.scan.provider, &args.scan.proto)?, - }; - - let pool = init_pool(&settings.db_path).await?; - let repo = Arc::new(ConfigRepo::new(pool)); - let killer: Arc = Arc::new(RealProcessKiller { - killall_enabled: settings.killall_enabled, - }); - let tester: Arc = Arc::new(RealVpnTester { - bin: settings.openvpn_bin.clone(), - killer: killer.clone(), - }); - let geo = Arc::new(IpInfoGeoLocator::new( - repo.clone(), - settings.ipinfo_token.clone(), - )); - - let scan_service = ScanService { - tester, - geo, - repo: repo.clone(), - }; - - let scan_options = ScanOptions { - dir, - limit, - timeout, - workers, - modify: args.scan.modify, - backup: args.scan.backup, - no_save: args.scan.no_save, - filter: settings.filter.clone(), - }; - - let cancel = CancellationToken::new(); - let cancel_task = cancel.clone(); - let signal_task = tokio::spawn(async move { - let _ = shutdown_signal().await; - cancel_task.cancel(); - }); - let _guard = CleanupGuard::new(killer.clone(), settings.killall_enabled); - - let progress: Arc = if crate::app::is_verbose(verbose) { - Arc::new(VerboseReporter) - } else { - Arc::new(ProgressReporter::new(settings.filter.to_display())) - }; - - let report = scan_service.scan(&scan_options, progress, cancel).await?; - signal_task.abort(); - - println!(); - println!("--- Scan Result ---"); - println!("Scanned: {}", report.scanned); - println!("Tested: {}", report.tested); - println!("Success: {}", report.success); - println!("Matched: {}", report.matched); - println!("Filter: {}", report.filter); - for m in &report.matched_configs { - println!("{} -- {}", m.country, m.path.display()); - } - - // Export this scan's fresh matches (the scan already stored successes, - // so `vmate-cli recent` is updated as normal). - if let Some(export_dir) = &args.scan.export { - let dest = vmate_core::paths::expand_path(export_dir); - let result = - vmate_core::export::export_configs_from_matches(&report.matched_configs, &dest).await?; - println!( - "Exported {} of {} configs to {}", - result.exported, - result.total, - result.dest.display() - ); - } + // The scan preamble (wiring, options, report, export) is shared with + // `scan`; `all` keeps only the connect half. + let (report, repo) = + crate::commands::scan::scan_pipeline(settings, &args.scan, verbose).await?; if args.no_connect { return Ok(()); @@ -138,8 +55,13 @@ pub async fn run(settings: &Settings, args: &AllArgs, verbose: &Verbosity) -> Re } let queue = ConnectQueue::new(candidates); + let registry = Arc::new(vmate_core::system::ProcessRegistry::new()); let runner = Arc::new(RealOpenVpnRunner { bin: settings.openvpn_bin.clone(), + registry: registry.clone(), + }); + let killer: Arc = Arc::new(RealProcessKiller { + killall_enabled: settings.killall_enabled, }); let options = ConnectOptions { connect_timeout: connect.connect_timeout, @@ -157,6 +79,7 @@ pub async fn run(settings: &Settings, args: &AllArgs, verbose: &Verbosity) -> Re let service = ConnectService { runner, killer, + registry, repo, options, }; diff --git a/crates/vmate-cli/src/commands/connect.rs b/crates/vmate-cli/src/commands/connect.rs index 92b9041..3c5875d 100644 --- a/crates/vmate-cli/src/commands/connect.rs +++ b/crates/vmate-cli/src/commands/connect.rs @@ -36,19 +36,14 @@ pub(crate) fn resolve_connect(us: &UserSettings, args: &ConnectArgs) -> Resolved /// Persist the explicitly-passed connect defaults, then confirm where they /// were written. -pub(crate) fn persist_connect_defaults(us: &mut UserSettings, args: &ConnectArgs) -> Result<()> { - if let Some(v) = args.connect_timeout { - us.connect_timeout_secs = Some(v.as_secs()); - } - if let Some(v) = args.cooldown { - us.cooldown_secs = Some(v.as_secs()); - } - if let Some(v) = args.retry_count { - us.retry_count = Some(v); - } - if let Some(v) = args.stability_grace { - us.stability_grace_secs = Some(v.as_secs()); - } +pub(crate) fn persist_connect_defaults(args: &ConnectArgs) -> Result<()> { + let mut us = UserSettings::load(); + us.persist_connect(&vmate_core::settings::ConnectDefaults { + connect_timeout: args.connect_timeout, + cooldown: args.cooldown, + retry_count: args.retry_count, + stability_grace: args.stability_grace, + }); us.save()?; println!( "Saved connect defaults to {}", @@ -61,8 +56,7 @@ pub async fn run(settings: &Settings, args: &ConnectArgs, verbose: &Verbosity) - // --save-defaults is a pure settings operation: persist and exit without // connecting (no root, no OpenVPN, no DB needed). if settings.save_defaults { - let mut us = UserSettings::load(); - persist_connect_defaults(&mut us, args)?; + persist_connect_defaults(args)?; return Ok(()); } @@ -95,8 +89,10 @@ pub async fn run(settings: &Settings, args: &ConnectArgs, verbose: &Verbosity) - let killer: Arc = Arc::new(RealProcessKiller { killall_enabled: settings.killall_enabled, }); + let registry = Arc::new(vmate_core::system::ProcessRegistry::new()); let runner = Arc::new(RealOpenVpnRunner { bin: settings.openvpn_bin.clone(), + registry: registry.clone(), }); let options = ConnectOptions { @@ -115,6 +111,7 @@ pub async fn run(settings: &Settings, args: &ConnectArgs, verbose: &Verbosity) - let service = ConnectService { runner, killer, + registry, repo, options, }; diff --git a/crates/vmate-cli/src/commands/doctor.rs b/crates/vmate-cli/src/commands/doctor.rs index 57b9620..14d352a 100644 --- a/crates/vmate-cli/src/commands/doctor.rs +++ b/crates/vmate-cli/src/commands/doctor.rs @@ -25,8 +25,7 @@ pub async fn run(settings: &Settings) -> Result<()> { format!("SQLite DB ({})", settings.db_path.display()), "ok".to_string(), ]); - let repo = ConfigRepo::new(pool); - match repo.journal_mode().await { + match vmate_core::db::pool::journal_mode(&pool).await { Ok(mode) if mode.eq_ignore_ascii_case("wal") => { table.add_row(["WAL mode".to_string(), "ok".to_string()]); } @@ -38,6 +37,7 @@ pub async fn run(settings: &Settings) -> Result<()> { } } + let repo = ConfigRepo::new(pool); let success = repo .count_configs(ConfigStatus::Success) .await diff --git a/crates/vmate-cli/src/commands/recent.rs b/crates/vmate-cli/src/commands/recent.rs index f87676c..0e0c1b4 100644 --- a/crates/vmate-cli/src/commands/recent.rs +++ b/crates/vmate-cli/src/commands/recent.rs @@ -68,12 +68,8 @@ fn print_plain(entries: &[StoredConfig]) -> Result<()> { .last_success_at .map(|t| t.format("%Y-%m-%d %H:%M").to_string()) .unwrap_or_else(|| "-".to_string()); - let path = if vmate_core::builtin::is_builtin_path(Path::new(&entry.path)) { - Path::new(&entry.path) - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or(&entry.path) - .to_string() + let path = if let Some(name) = vmate_core::builtin::display_name(Path::new(&entry.path)) { + name } else if term::stdout_is_tty() { hyperlink::osc8_file_hyperlink(&entry.path).unwrap_or_else(|| entry.path.clone()) } else { diff --git a/crates/vmate-cli/src/commands/scan.rs b/crates/vmate-cli/src/commands/scan.rs index f5df3b5..e5e42cb 100644 --- a/crates/vmate-cli/src/commands/scan.rs +++ b/crates/vmate-cli/src/commands/scan.rs @@ -28,6 +28,19 @@ pub async fn run(settings: &Settings, args: &ScanArgs, verbose: &Verbosity) -> R require_root_for("run OpenVPN tests", settings.no_elevate)?; + scan_pipeline(settings, args, verbose).await?; + Ok(()) +} + +/// The full scan orchestration: DB/repo/killer/tester/geo wiring, option +/// resolution, the scan itself, the report and the export. Returns the report +/// and repo so `all` can go on to connect. Shared by `scan` and `all` so the +/// preamble (and its report format) lives in one place. +pub(crate) async fn scan_pipeline( + settings: &Settings, + args: &ScanArgs, + verbose: &Verbosity, +) -> Result<(ScanReport, Arc)> { let us = UserSettings::load(); let (workers, limit, timeout) = resolve_scan_defaults(args, &us); @@ -35,12 +48,14 @@ pub async fn run(settings: &Settings, args: &ScanArgs, verbose: &Verbosity) -> R let pool = init_pool(&settings.db_path).await?; let repo = Arc::new(ConfigRepo::new(pool)); + let registry = Arc::new(vmate_core::system::ProcessRegistry::new()); let killer: Arc = Arc::new(RealProcessKiller { killall_enabled: settings.killall_enabled, }); let tester: Arc = Arc::new(RealVpnTester { bin: settings.openvpn_bin.clone(), killer: killer.clone(), + registry: registry.clone(), }); let geo = Arc::new(IpInfoGeoLocator::new( repo.clone(), @@ -72,7 +87,7 @@ pub async fn run(settings: &Settings, args: &ScanArgs, verbose: &Verbosity) -> R let _ = shutdown_signal().await; cancel_task.cancel(); }); - let _guard = CleanupGuard::new(killer.clone(), settings.killall_enabled); + let _guard = CleanupGuard::new(killer.clone(), registry.clone(), settings.killall_enabled); let progress: Arc = if crate::app::is_verbose(verbose) { Arc::new(VerboseReporter) @@ -99,7 +114,7 @@ pub async fn run(settings: &Settings, args: &ScanArgs, verbose: &Verbosity) -> R ); } - Ok(()) + Ok((report, repo)) } /// Resolve the scan `workers`/`limit`/`timeout` as @@ -115,24 +130,15 @@ pub(crate) fn resolve_scan_defaults( ) } -/// Apply the explicitly-passed scan default flags onto `us`. -fn apply_scan_defaults(us: &mut UserSettings, args: &ScanArgs) { - if let Some(v) = args.max { - us.max_workers = Some(v as u64); - } - if let Some(v) = args.limit { - us.limit = Some(v as u64); - } - if let Some(t) = args.timeout { - us.scan_timeout_secs = Some(t.as_secs()); - } -} - /// Persist the explicitly-passed scan default flags to the user config and /// print a confirmation. Only the flags actually passed are written. pub(crate) fn persist_scan_defaults(args: &ScanArgs) -> Result<()> { let mut us = UserSettings::load(); - apply_scan_defaults(&mut us, args); + us.persist_scan(&vmate_core::settings::ScanDefaults { + max_workers: args.max.map(|v| v as u64), + limit: args.limit.map(|v| v as u64), + timeout: args.timeout, + }); us.save()?; println!("Saved scan defaults to {}", UserSettings::path()?.display()); Ok(()) @@ -301,15 +307,17 @@ mod tests { } #[test] - fn apply_scan_defaults_only_writes_explicitly_passed_keys() { + fn persist_scan_only_writes_explicitly_passed_keys() { + use vmate_core::settings::ScanDefaults; let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("settings.json"); let mut us = UserSettings::default(); - apply_scan_defaults( - &mut us, - &scan_args(Some(500), None, Some(Duration::from_secs(20))), - ); + us.persist_scan(&ScanDefaults { + max_workers: Some(500), + limit: None, + timeout: Some(Duration::from_secs(20)), + }); us.save_to(&path).unwrap(); let loaded = UserSettings::load_from(&path); @@ -320,7 +328,11 @@ mod tests { // A later save with a different explicit key must not resurrect keys // that were never passed. let mut us2 = loaded; - apply_scan_defaults(&mut us2, &scan_args(None, Some(25), None)); + us2.persist_scan(&ScanDefaults { + max_workers: None, + limit: Some(25), + timeout: None, + }); us2.save_to(&path).unwrap(); let loaded2 = UserSettings::load_from(&path); assert_eq!(loaded2.max_workers, Some(500)); diff --git a/crates/vmate-cli/src/ui/progress.rs b/crates/vmate-cli/src/ui/progress.rs index e28d94d..e280d30 100644 --- a/crates/vmate-cli/src/ui/progress.rs +++ b/crates/vmate-cli/src/ui/progress.rs @@ -30,20 +30,7 @@ impl ProgressReporter { } } - fn update(&self, tested: Option, ok: Option, matched: Option) { - let mut state = match self.state.lock() { - Ok(s) => s, - Err(e) => e.into_inner(), - }; - if let Some(t) = tested { - state.0 = t; - } - if let Some(o) = ok { - state.1 = o; - } - if let Some(m) = matched { - state.2 = m; - } + fn render(&self, state: &(usize, usize, usize)) { self.bar.set_position(state.0 as u64); self.bar.set_message(format!( "tested {} | ok {} | matched {}", @@ -57,16 +44,31 @@ impl ScanProgress for ProgressReporter { self.bar.set_length(total as u64); } - fn tested(&self, n: usize) { - self.update(Some(n), None, None); + fn tested(&self) { + let mut state = match self.state.lock() { + Ok(s) => s, + Err(e) => e.into_inner(), + }; + state.0 += 1; + self.render(&state); } - fn ok(&self, n: usize) { - self.update(None, Some(n), None); + fn ok(&self) { + let mut state = match self.state.lock() { + Ok(s) => s, + Err(e) => e.into_inner(), + }; + state.1 += 1; + self.render(&state); } - fn matched(&self, n: usize) { - self.update(None, None, Some(n)); + fn matched(&self) { + let mut state = match self.state.lock() { + Ok(s) => s, + Err(e) => e.into_inner(), + }; + state.2 += 1; + self.render(&state); } fn success(&self, _path: &Path, _country: &CountryCode) {} @@ -86,11 +88,11 @@ impl ScanProgress for VerboseReporter { println!("Testing {total} configs"); } - fn tested(&self, _n: usize) {} + fn tested(&self) {} - fn ok(&self, _n: usize) {} + fn ok(&self) {} - fn matched(&self, _n: usize) {} + fn matched(&self) {} fn success(&self, path: &Path, country: &CountryCode) { println!("[SUCCESS] {} --- {}", path.display(), country); diff --git a/crates/vmate-cli/src/ui/recent.rs b/crates/vmate-cli/src/ui/recent.rs index d941003..d0a4bf4 100644 --- a/crates/vmate-cli/src/ui/recent.rs +++ b/crates/vmate-cli/src/ui/recent.rs @@ -24,15 +24,7 @@ use vmate_core::db::models::StoredConfig; /// Path shown for an entry: for a built-in config, its remote display /// (`-`, from the file stem); otherwise the filesystem path. fn display_path(path: &str) -> String { - if vmate_core::builtin::is_builtin_path(Path::new(path)) { - Path::new(path) - .file_stem() - .and_then(|s| s.to_str()) - .unwrap_or(path) - .to_string() - } else { - path.to_string() - } + vmate_core::builtin::display_name(Path::new(path)).unwrap_or_else(|| path.to_string()) } type Term = Terminal>; diff --git a/crates/vmate-core/src/builtin/mod.rs b/crates/vmate-core/src/builtin/mod.rs index 200897b..553b8d8 100644 --- a/crates/vmate-core/src/builtin/mod.rs +++ b/crates/vmate-core/src/builtin/mod.rs @@ -90,6 +90,36 @@ pub struct BuiltinConfig { pub proto: Proto, } +impl BuiltinConfig { + /// Reconstruct the identity from a materialized path + /// (`///-.ovpn`). The path layout + /// is owned here — this is the single place that parses it back out. + pub fn from_path(path: &Path) -> Option { + if !is_builtin_path(path) { + return None; + } + let stem = path.file_stem()?.to_str()?; + let (host, port) = stem.rsplit_once('-')?; + if host.is_empty() || port.is_empty() || !port.chars().all(|c| c.is_ascii_digit()) { + return None; + } + let proto_dir = path.parent()?; + let proto = Proto::from_name(proto_dir.file_name()?.to_str()?)?; + let provider_dir = proto_dir.parent()?; + let provider = Provider::from_name(provider_dir.file_name()?.to_str()?)?; + Some(BuiltinConfig { + provider, + remote: format!("remote {host} {port}"), + proto, + }) + } + + /// The `host` and `port` of this config's remote. + pub fn host_port(&self) -> Option<(&str, &str)> { + remote_host_port(&self.remote) + } +} + /// Header template with `{proto}` and `{remote}` placeholders, substituted by /// [`build_config`]. const HEADER_TEMPLATE: &str = "\ @@ -254,27 +284,26 @@ pub fn is_builtin_path(path: &Path) -> bool { } } +/// The user-facing name for a built-in config path: `-` (its file +/// stem). `None` when the path is not a built-in. The display convention lives +/// here so recent/connect don't re-derive the layout. +pub fn display_name(path: &Path) -> Option { + if !is_builtin_path(path) { + return None; + } + path.file_stem().map(|s| s.to_string_lossy().into_owned()) +} + /// For a built-in path, the export file name /// `{provider}_{host}-{port}_{COUNTRY}.ovpn`, e.g. /// `vpn-gate_public-vpn-38.opengw.net-1195_JP.ovpn`. Returns `None` if the path /// is not a built-in or the remote can't be parsed. pub fn export_name(path: &Path, country: &str) -> Option { - if !is_builtin_path(path) { - return None; - } - // Layout: ///-.ovpn - let provider_name = path.parent()?.parent()?.file_name()?.to_str()?; - let provider = Provider::from_name(provider_name)?; - let stem = path.file_stem()?.to_str()?; - let dash = stem.rfind('-')?; - let host = &stem[..dash]; - let port = &stem[dash + 1..]; - if host.is_empty() || port.is_empty() || !port.chars().all(|c| c.is_ascii_digit()) { - return None; - } + let cfg = BuiltinConfig::from_path(path)?; + let (host, port) = cfg.host_port()?; Some(format!( "{}_{}-{}_{}.ovpn", - provider.name(), + cfg.provider.name(), host, port, country @@ -715,4 +744,28 @@ mod tests { // Unparseable remote (no numeric port) yields None. assert_eq!(export_name(&dir.join("just-a-host.ovpn"), "JP"), None); } + + #[test] + fn from_path_reconstructs_identity() { + let dir = paths::builtin_dir().unwrap().join("vpn-gate").join("tcp"); + let path = dir.join("public-vpn-38.opengw.net-1195.ovpn"); + let cfg = BuiltinConfig::from_path(&path).unwrap(); + assert_eq!(cfg.provider, Provider::VpnGate); + assert_eq!(cfg.proto, Proto::Tcp); + assert_eq!(cfg.host_port(), Some(("public-vpn-38.opengw.net", "1195"))); + assert_eq!(cfg.remote, "remote public-vpn-38.opengw.net 1195"); + // Not under the builtins dir -> None. + assert!(BuiltinConfig::from_path(Path::new("/tmp/x-1.ovpn")).is_none()); + } + + #[test] + fn display_name_is_host_port_stem() { + let dir = paths::builtin_dir().unwrap().join("vpn-gate").join("udp"); + let path = dir.join("public-vpn-38.opengw.net-1195.ovpn"); + assert_eq!( + display_name(&path).unwrap(), + "public-vpn-38.opengw.net-1195" + ); + assert_eq!(display_name(Path::new("/tmp/foo.ovpn")), None); + } } diff --git a/crates/vmate-core/src/connect/service.rs b/crates/vmate-core/src/connect/service.rs index f9b9227..b47ed40 100644 --- a/crates/vmate-core/src/connect/service.rs +++ b/crates/vmate-core/src/connect/service.rs @@ -18,7 +18,7 @@ use crate::connect::session::Phase2Exit; use crate::connect::state::ConnectionStatus; use crate::db::ConfigRepo; use crate::ovpn::process::{ConnectOutcome, OpenVpnRunner, connect_args, monitor_connect}; -use crate::system::killer::{CleanupGuard, ProcessKiller}; +use crate::system::killer::{CleanupGuard, ProcessKiller, ProcessRegistry}; use anyhow::Result; use std::path::Path; use std::sync::Arc; @@ -40,28 +40,74 @@ pub struct ConnectOptions { pub killall_enabled: bool, } +/// What a failed attempt should do with the same config. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RetryDecision { + /// Retry the config once more. + Retry, + /// Drop the config from history (budget exhausted). + Drop, +} + +/// Pure retry/drop budget: how many failed attempts a config is allowed before +/// it is dropped from history. Synchronous and process-free, so the retry +/// policy is unit-testable without spawning OpenVPN. +#[derive(Debug, Clone, Copy)] +pub struct RetryBudget { + pub max: u32, + pub failures: u32, +} + +impl RetryBudget { + /// A budget that drops a config after `max` failed attempts. + pub fn new(max: u32) -> Self { + Self { + max: max.max(1), + failures: 0, + } + } + + /// Reset the counter after a session that proved stable. + pub fn reset(&mut self) { + self.failures = 0; + } + + /// Register one failed attempt and decide whether to retry or drop. + pub fn on_failure(&mut self) -> RetryDecision { + self.failures += 1; + if self.failures >= self.max { + RetryDecision::Drop + } else { + RetryDecision::Retry + } + } +} + /// Orchestrates the connect loop. pub struct ConnectService { pub runner: Arc, pub killer: Arc, + pub registry: Arc, pub repo: Arc, pub options: ConnectOptions, } impl ConnectService { pub async fn run(&self, mut queue: ConnectQueue, host: &mut dyn ConnectHost) -> Result<()> { - let _guard = CleanupGuard::new(self.killer.clone(), self.options.killall_enabled); + let _guard = CleanupGuard::new( + self.killer.clone(), + self.registry.clone(), + self.options.killall_enabled, + ); let cancel = CancellationToken::new(); // Go parity: try once, reconnect once, then drop from history. The // budget is configurable via ConnectOptions::retry_count. - let max_failures = self.options.retry_count; - while let Some(first) = queue.next_candidate() { if cancel.is_cancelled() { break; } let candidate = first; - let mut failures: u32 = 0; + let mut budget = RetryBudget::new(self.options.retry_count); let mut last_reason: Option = None; let mut reconnecting = false; @@ -73,17 +119,21 @@ impl ConnectService { let message = if reconnecting { format!("Reconnecting to {}", candidate.country) - } else if failures == 0 { + } else if budget.failures == 0 { format!("Connecting to {}", candidate.country) } else { match &last_reason { Some(reason) => format!( - "Retrying {} (failure {failures}/{max_failures}): {reason}", - candidate.country + "Retrying {} (failure {}/{max}): {reason}", + candidate.country, + budget.failures, + max = budget.max ), None => format!( - "Retrying {} (failure {failures}/{max_failures})", - candidate.country + "Retrying {} (failure {}/{max})", + candidate.country, + budget.failures, + max = budget.max ), } }; @@ -109,7 +159,6 @@ impl ConnectService { return Err(e); } }; - let pid = handle.child.id().unwrap_or(0); let outcome = monitor_connect( &mut handle.lines, self.options.connect_timeout, @@ -120,14 +169,14 @@ impl ConnectService { match outcome { ConnectOutcome::Cancelled => { - self.kill_handle(pid, &mut handle).await; + handle.kill_graceful(self.killer.as_ref()).await; host.finish().await?; return Ok(()); } ConnectOutcome::Next => { // User pressed `n` during the handshake: KEEP in DB, // defer in-session, move to the next candidate. - self.kill_handle(pid, &mut handle).await; + handle.kill_graceful(self.killer.as_ref()).await; let _ = self.repo.mark_skipped(candidate.id).await; queue.skip(candidate); break; @@ -139,7 +188,7 @@ impl ConnectService { let exit = self .monitor_connected(&candidate, &mut handle, host, &cancel) .await?; - self.kill_handle(pid, &mut handle).await; // targeted kill only + handle.kill_graceful(self.killer.as_ref()).await; // targeted kill only match exit { Phase2Exit::Quit => { host.finish().await?; @@ -164,16 +213,16 @@ impl ConnectService { // connect-then-crash flakiness and still counts — // resetting there would retry it forever. if connected_at.elapsed() >= self.options.connect_stability_grace { - failures = 0; + budget.reset(); } let _ = self .repo .note_connect_failure(Path::new(&candidate.path)) .await; - failures += 1; last_reason = Some(reason.clone()); - if failures >= max_failures { - self.drop_candidate(&candidate, failures, host).await?; + if budget.on_failure() == RetryDecision::Drop { + self.drop_candidate(&candidate, budget.failures, host) + .await?; break; } host.notify(&format!("{reason}; retrying {}", candidate.country)) @@ -202,15 +251,15 @@ impl ConnectService { unreachable!() } }; - self.kill_handle(pid, &mut handle).await; // targeted kill only + handle.kill_graceful(self.killer.as_ref()).await; // targeted kill only let _ = self .repo .note_connect_failure(Path::new(&candidate.path)) .await; - failures += 1; last_reason = Some(reason.clone()); - if failures >= max_failures { - self.drop_candidate(&candidate, failures, host).await?; + if budget.on_failure() == RetryDecision::Drop { + self.drop_candidate(&candidate, budget.failures, host) + .await?; break; } host.notify(&format!("{reason}; retrying {}", candidate.country)) @@ -237,6 +286,32 @@ mod tests { use async_trait::async_trait; use std::sync::Mutex; + #[test] + fn budget_drops_after_max_failures() { + let mut b = RetryBudget::new(2); + assert_eq!(b.on_failure(), RetryDecision::Retry); + assert_eq!(b.on_failure(), RetryDecision::Drop); + } + + #[test] + fn budget_retry_count_is_configurable() { + let mut b = RetryBudget::new(3); + assert_eq!(b.on_failure(), RetryDecision::Retry); + assert_eq!(b.on_failure(), RetryDecision::Retry); + assert_eq!(b.on_failure(), RetryDecision::Drop); + } + + #[test] + fn budget_reset_lets_a_stable_session_retry_again() { + // A long-lived session's crash resets the budget, so two network + // blips never delete a working config. + let mut b = RetryBudget::new(2); + assert_eq!(b.on_failure(), RetryDecision::Retry); + b.reset(); + assert_eq!(b.on_failure(), RetryDecision::Retry); + assert_eq!(b.on_failure(), RetryDecision::Drop); + } + struct FakeHost { commands: Mutex>, messages: Mutex>, @@ -296,7 +371,9 @@ mod tests { /// Spawns `sh -c 'echo "Initialization Sequence Completed"'`: the handshake /// succeeds, then the process exits immediately, so the connected session /// crashes right after connecting (a connect-then-crash). - struct ConnectThenExitRunner; + struct ConnectThenExitRunner { + registry: Arc, + } impl OpenVpnRunner for ConnectThenExitRunner { fn spawn(&self, _args: &[String]) -> Result { @@ -306,6 +383,7 @@ mod tests { "-c".into(), "echo 'Initialization Sequence Completed'".into(), ], + &self.registry, ) } fn bin(&self) -> &str { @@ -315,11 +393,17 @@ mod tests { /// Spawns a real `sh -c 'echo AUTH_FAILED'`, so the handshake monitor /// immediately sees an error line and the connection fails. - struct FailHandshakeRunner; + struct FailHandshakeRunner { + registry: Arc, + } impl OpenVpnRunner for FailHandshakeRunner { fn spawn(&self, _args: &[String]) -> Result { - crate::ovpn::process::spawn_openvpn("sh", &["-c".into(), "echo AUTH_FAILED".into()]) + crate::ovpn::process::spawn_openvpn( + "sh", + &["-c".into(), "echo AUTH_FAILED".into()], + &self.registry, + ) } fn bin(&self) -> &str { "sh" @@ -328,11 +412,17 @@ mod tests { /// Spawns `sh -c 'sleep 30'`: no completion or error lines, so the /// handshake hangs until a key is pressed or the timeout fires. - struct SlowRunner; + struct SlowRunner { + registry: Arc, + } impl OpenVpnRunner for SlowRunner { fn spawn(&self, _args: &[String]) -> Result { - crate::ovpn::process::spawn_openvpn("sh", &["-c".into(), "sleep 30".into()]) + crate::ovpn::process::spawn_openvpn( + "sh", + &["-c".into(), "sleep 30".into()], + &self.registry, + ) } fn bin(&self) -> &str { "sh" @@ -343,7 +433,9 @@ mod tests { /// the handshake succeeds, the session stays up past the stability grace, /// then the process exits — a long-lived session's crash, not a /// connect-then-crash. - struct ConnectThenStableThenExitRunner; + struct ConnectThenStableThenExitRunner { + registry: Arc, + } impl OpenVpnRunner for ConnectThenStableThenExitRunner { fn spawn(&self, _args: &[String]) -> Result { @@ -353,6 +445,7 @@ mod tests { "-c".into(), "echo 'Initialization Sequence Completed'; sleep 0.2".into(), ], + &self.registry, ) } fn bin(&self) -> &str { @@ -367,6 +460,7 @@ mod tests { ConnectService { runner: Arc::new(NeverSpawn), killer, + registry: Arc::new(ProcessRegistry::new()), repo, options: ConnectOptions { connect_timeout: Duration::from_secs(1), @@ -440,9 +534,13 @@ mod tests { let killer: Arc = Arc::new(crate::system::killer::RealProcessKiller { killall_enabled: false, }); + let registry = Arc::new(ProcessRegistry::new()); let service = ConnectService { - runner: Arc::new(FailHandshakeRunner), + runner: Arc::new(FailHandshakeRunner { + registry: registry.clone(), + }), killer, + registry, repo: repo.clone(), options: ConnectOptions { connect_timeout: Duration::from_secs(1), @@ -472,6 +570,7 @@ mod tests { /// (`sleep 30`), so a configurable retry budget can be observed mid-flight. struct FailTwiceThenHang { spawns: std::sync::atomic::AtomicUsize, + registry: Arc, } impl OpenVpnRunner for FailTwiceThenHang { @@ -479,9 +578,17 @@ mod tests { use std::sync::atomic::Ordering; let n = self.spawns.fetch_add(1, Ordering::SeqCst); if n < 2 { - crate::ovpn::process::spawn_openvpn("sh", &["-c".into(), "echo AUTH_FAILED".into()]) + crate::ovpn::process::spawn_openvpn( + "sh", + &["-c".into(), "echo AUTH_FAILED".into()], + &self.registry, + ) } else { - crate::ovpn::process::spawn_openvpn("sh", &["-c".into(), "sleep 30".into()]) + crate::ovpn::process::spawn_openvpn( + "sh", + &["-c".into(), "sleep 30".into()], + &self.registry, + ) } } fn bin(&self) -> &str { @@ -549,7 +656,9 @@ mod tests { .await .unwrap(); - let make_service = |runner: Arc, repo: Arc| { + let make_service = |runner: Arc, + registry: Arc, + repo: Arc| { let killer: Arc = Arc::new(crate::system::killer::RealProcessKiller { killall_enabled: false, @@ -557,9 +666,14 @@ mod tests { ConnectService { runner, killer, + registry, repo, options: ConnectOptions { - connect_timeout: Duration::from_secs(1), + // Long handshake timeout so the host's `q` (returned + // once two failures are seen) always wins the race over a + // timeout — otherwise a scheduler-delayed poll can turn + // the intended cancel into a 3rd failure and a drop. + connect_timeout: Duration::from_secs(30), connect_stability_grace: Duration::from_secs(1), retry_count: 3, killall_enabled: false, @@ -574,10 +688,13 @@ mod tests { failures_seen: 0, quit_after: 2, }; + let registry = Arc::new(ProcessRegistry::new()); let service = make_service( Arc::new(FailTwiceThenHang { spawns: std::sync::atomic::AtomicUsize::new(0), + registry: registry.clone(), }), + registry, repo.clone(), ); service @@ -600,7 +717,14 @@ mod tests { ); // Phase 2: from a clean slate, three real failures drop the config. - let service = make_service(Arc::new(FailHandshakeRunner), repo.clone()); + let registry = Arc::new(ProcessRegistry::new()); + let service = make_service( + Arc::new(FailHandshakeRunner { + registry: registry.clone(), + }), + registry, + repo.clone(), + ); let mut host = FakeHost::new(vec![]); service .run( @@ -658,9 +782,13 @@ mod tests { let killer: Arc = Arc::new(crate::system::killer::RealProcessKiller { killall_enabled: false, }); + let registry = Arc::new(ProcessRegistry::new()); let service = ConnectService { - runner: Arc::new(ConnectThenExitRunner), + runner: Arc::new(ConnectThenExitRunner { + registry: registry.clone(), + }), killer, + registry, repo: repo.clone(), options: ConnectOptions { connect_timeout: Duration::from_secs(1), @@ -762,9 +890,13 @@ mod tests { let killer: Arc = Arc::new(crate::system::killer::RealProcessKiller { killall_enabled: false, }); + let registry = Arc::new(ProcessRegistry::new()); let service = ConnectService { - runner: Arc::new(ConnectThenStableThenExitRunner), + runner: Arc::new(ConnectThenStableThenExitRunner { + registry: registry.clone(), + }), killer, + registry, repo: repo.clone(), options: ConnectOptions { connect_timeout: Duration::from_secs(1), @@ -821,9 +953,13 @@ mod tests { let killer: Arc = Arc::new(crate::system::killer::RealProcessKiller { killall_enabled: false, }); + let registry = Arc::new(ProcessRegistry::new()); let service = ConnectService { - runner: Arc::new(SlowRunner), + runner: Arc::new(SlowRunner { + registry: registry.clone(), + }), killer, + registry, repo: repo.clone(), options: ConnectOptions { connect_timeout: Duration::from_secs(30), // quit must beat this diff --git a/crates/vmate-core/src/connect/session.rs b/crates/vmate-core/src/connect/session.rs index 4182c4e..d187f43 100644 --- a/crates/vmate-core/src/connect/session.rs +++ b/crates/vmate-core/src/connect/session.rs @@ -76,19 +76,6 @@ impl ConnectService { } } - /// Targeted kill of one OpenVPN process tree: SIGTERM the group, wait up - /// to `KILL_GRACE`, then SIGKILL if needed. Deliberately NO global - /// `killall` here — the `CleanupGuard` (per-process registry + optional - /// killall) is the global safety net. - pub(crate) async fn kill_handle(&self, pid: u32, handle: &mut OpenVpnHandle) { - crate::system::killer::kill_process_tree_graceful( - self.killer.as_ref(), - pid, - &mut handle.child, - ) - .await; - } - /// Go parity: after `MAX_FAILURES`, remove the config from history entirely. pub(crate) async fn drop_candidate( &self, diff --git a/crates/vmate-core/src/db/pool.rs b/crates/vmate-core/src/db/pool.rs index 7d08f93..1e69238 100644 --- a/crates/vmate-core/src/db/pool.rs +++ b/crates/vmate-core/src/db/pool.rs @@ -1,6 +1,7 @@ //! SQLite connection pool with WAL mode and automatic migrations. use anyhow::{Context, Result}; +use sqlx::Row; use sqlx::SqlitePool; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; use std::path::Path; @@ -45,3 +46,12 @@ pub async fn init_pool(db_path: &Path) -> Result { Ok(pool) } + +/// The current SQLite journal mode (used by `vmate-cli doctor`). +pub async fn journal_mode(pool: &DbPool) -> Result { + let row = sqlx::query("PRAGMA journal_mode") + .fetch_one(pool) + .await + .context("journal_mode check failed")?; + Ok(row.get::("journal_mode")) +} diff --git a/crates/vmate-core/src/db/repo/geo_cache.rs b/crates/vmate-core/src/db/repo/geo_cache.rs index 68e280b..c3a27a1 100644 --- a/crates/vmate-core/src/db/repo/geo_cache.rs +++ b/crates/vmate-core/src/db/repo/geo_cache.rs @@ -1,4 +1,4 @@ -//! IP -> country cache and journal-mode diagnostics. +//! IP -> country cache. use crate::country::CountryCode; use crate::db::repo::ConfigRepo; @@ -42,13 +42,4 @@ impl ConfigRepo { .context("cache_country_for_ip failed")?; Ok(()) } - - /// Current journal mode (used by `vmate-cli doctor`). - pub async fn journal_mode(&self) -> Result { - let row = sqlx::query("PRAGMA journal_mode") - .fetch_one(&self.pool) - .await - .context("journal_mode check failed")?; - Ok(row.get::("journal_mode")) - } } diff --git a/crates/vmate-core/src/db/repo/mod.rs b/crates/vmate-core/src/db/repo/mod.rs index 86f2255..a6a25a1 100644 --- a/crates/vmate-core/src/db/repo/mod.rs +++ b/crates/vmate-core/src/db/repo/mod.rs @@ -241,7 +241,10 @@ mod tests { #[tokio::test] async fn journal_mode_is_wal() { - let (repo, _dir) = test_repo().await; - assert_eq!(repo.journal_mode().await.unwrap(), "wal"); + let dir = tempfile::tempdir().expect("tempdir"); + let pool = init_pool(&dir.path().join("test.db")) + .await + .expect("init pool"); + assert_eq!(crate::db::pool::journal_mode(&pool).await.unwrap(), "wal"); } } diff --git a/crates/vmate-core/src/ovpn/process.rs b/crates/vmate-core/src/ovpn/process.rs index a90688a..276599a 100644 --- a/crates/vmate-core/src/ovpn/process.rs +++ b/crates/vmate-core/src/ovpn/process.rs @@ -2,7 +2,7 @@ use crate::connect::{ConnectHost, UserCommand}; use crate::ovpn::monitor::{VpnLineClass, classify_line}; -use crate::system::killer::ProcessKiller; +use crate::system::killer::{ProcessKiller, ProcessRegistry}; use anyhow::{Context, Result, anyhow}; use async_trait::async_trait; use std::path::{Path, PathBuf}; @@ -57,11 +57,29 @@ pub struct OpenVpnHandle { pub lines: mpsc::Receiver, } -/// Spawn OpenVPN in a new process group. +impl OpenVpnHandle { + /// Stop the process tree gracefully: SIGTERM the group, wait up to + /// `KILL_GRACE`, then SIGKILL stragglers. The handle owns its teardown, so + /// callers never reach for the raw pid or a kill policy. + pub async fn kill_graceful(&mut self, killer: &dyn ProcessKiller) { + crate::system::killer::kill_process_tree_graceful( + killer, + self.child.id().unwrap_or(0), + &mut self.child, + ) + .await; + } +} + +/// Spawn OpenVPN in a new process group, registering the pid in `registry`. /// /// Each spawned process gets its own process group so it can be killed with /// the whole tree via a single negative-pid SIGKILL. -pub fn spawn_openvpn(bin: &str, args: &[String]) -> Result { +pub fn spawn_openvpn( + bin: &str, + args: &[String], + registry: &ProcessRegistry, +) -> Result { let mut cmd = tokio::process::Command::new(bin); cmd.args(args); cmd.stdout(Stdio::piped()); @@ -77,9 +95,9 @@ pub fn spawn_openvpn(bin: &str, args: &[String]) -> Result { .spawn() .with_context(|| format!("failed to spawn OpenVPN binary: {bin}"))?; - // Track the pid so the cleanup guard can kill exactly the processes vmate - // spawned (a pid of 0 — unknown — is ignored by the registry). - crate::system::killer::register_process(child.id().unwrap_or(0)); + // Track the pid so the cleanup guard can kill exactly the processes this + // session spawned (a pid of 0 — unknown — is ignored by the registry). + registry.register(child.id().unwrap_or(0)); let stdout = child .stdout @@ -123,11 +141,12 @@ pub trait OpenVpnRunner: Send + Sync { /// Production runner that executes the configured `openvpn` binary. pub struct RealOpenVpnRunner { pub bin: String, + pub registry: Arc, } impl OpenVpnRunner for RealOpenVpnRunner { fn spawn(&self, args: &[String]) -> Result { - spawn_openvpn(&self.bin, args) + spawn_openvpn(&self.bin, args, &self.registry) } fn bin(&self) -> &str { @@ -237,16 +256,15 @@ pub async fn test_openvpn_config( timeout: Duration, cancel: CancellationToken, killer: &dyn ProcessKiller, + registry: &ProcessRegistry, ) -> Result { let args = test_args(config); - let mut handle = spawn_openvpn(bin, &args)?; - let pid = handle.child.id().unwrap_or(0); + let mut handle = spawn_openvpn(bin, &args, registry)?; let outcome = monitor_test(&mut handle.lines, timeout, cancel).await; - // Always clean up the process group; never leak workers. Graceful: SIGTERM - // the group, then SIGKILL only if it is still alive after KILL_GRACE. - crate::system::killer::kill_process_tree_graceful(killer, pid, &mut handle.child).await; + // Always clean up the process tree; never leak workers. + handle.kill_graceful(killer).await; Ok(matches!(outcome, MonitorOutcome::Success)) } @@ -266,6 +284,7 @@ pub trait VpnTester: Send + Sync { pub struct RealVpnTester { pub bin: String, pub killer: Arc, + pub registry: Arc, } #[async_trait] @@ -276,7 +295,15 @@ impl VpnTester for RealVpnTester { timeout: Duration, cancel: CancellationToken, ) -> Result { - test_openvpn_config(&self.bin, config, timeout, cancel, self.killer.as_ref()).await + test_openvpn_config( + &self.bin, + config, + timeout, + cancel, + self.killer.as_ref(), + &self.registry, + ) + .await } } @@ -299,8 +326,8 @@ mod tests { /// wait out the SIGKILL escalation. #[tokio::test] async fn kill_process_tree_graceful_sigterms_and_exits_quickly() { - crate::system::killer::clear_registry(); - let mut handle = spawn_openvpn("sh", &["-c".into(), "sleep 5".into()]).unwrap(); + let registry = crate::system::killer::ProcessRegistry::new(); + let mut handle = spawn_openvpn("sh", &["-c".into(), "sleep 5".into()], ®istry).unwrap(); let pid = handle.child.id().unwrap_or(0); assert!(pid != 0); diff --git a/crates/vmate-core/src/scan/report.rs b/crates/vmate-core/src/scan/report.rs index 3b40d5a..0c3057a 100644 --- a/crates/vmate-core/src/scan/report.rs +++ b/crates/vmate-core/src/scan/report.rs @@ -44,9 +44,12 @@ pub struct ScanOptions { /// and must be internally synchronized. pub trait ScanProgress: Send + Sync { fn total(&self, total: usize); - fn tested(&self, n: usize); - fn ok(&self, n: usize); - fn matched(&self, n: usize); + /// One config finished testing. + fn tested(&self); + /// One config succeeded. + fn ok(&self); + /// One config matched the filter. + fn matched(&self); fn success(&self, path: &Path, country: &CountryCode); fn failed(&self, path: &Path); /// Called when the scan finishes (lets a progress bar clear itself). diff --git a/crates/vmate-core/src/scan/service.rs b/crates/vmate-core/src/scan/service.rs index ab39192..c565670 100644 --- a/crates/vmate-core/src/scan/service.rs +++ b/crates/vmate-core/src/scan/service.rs @@ -108,7 +108,7 @@ impl ScanService { }; counts.tested.fetch_add(1, Ordering::SeqCst); - progress.tested(counts.tested.load(Ordering::SeqCst)); + progress.tested(); if !ok { progress.failed(&path); @@ -119,7 +119,7 @@ impl ScanService { } counts.ok.fetch_add(1, Ordering::SeqCst); - progress.ok(counts.ok.load(Ordering::SeqCst)); + progress.ok(); // Geo lookup is independent of the test and must never block the // whole scan on a slow HTTP call; failures degrade to UNKNOWN. @@ -159,7 +159,7 @@ impl ScanService { country: lookup.country, }); counts.matched.fetch_add(1, Ordering::SeqCst); - progress.matched(counts.matched.load(Ordering::SeqCst)); + progress.matched(); if results.len() >= limit { cancel.cancel(); } diff --git a/crates/vmate-core/src/settings.rs b/crates/vmate-core/src/settings.rs index 5cc9257..ad08bc2 100644 --- a/crates/vmate-core/src/settings.rs +++ b/crates/vmate-core/src/settings.rs @@ -126,6 +126,56 @@ impl UserSettings { .or(self.stability_grace_secs.map(Duration::from_secs)) .unwrap_or(Duration::from_secs(5)) } + + /// Persist the explicitly-passed scan defaults onto this struct. Only the + /// flags that were actually given are written; the rest keep their value. + pub fn persist_scan(&mut self, values: &ScanDefaults) { + if let Some(v) = values.max_workers { + self.max_workers = Some(v); + } + if let Some(v) = values.limit { + self.limit = Some(v); + } + if let Some(t) = values.timeout { + self.scan_timeout_secs = Some(t.as_secs()); + } + } + + /// Persist the explicitly-passed connect defaults onto this struct. Only + /// the flags that were actually given are written; the rest keep theirs. + pub fn persist_connect(&mut self, values: &ConnectDefaults) { + if let Some(v) = values.connect_timeout { + self.connect_timeout_secs = Some(v.as_secs()); + } + if let Some(v) = values.cooldown { + self.cooldown_secs = Some(v.as_secs()); + } + if let Some(v) = values.retry_count { + self.retry_count = Some(v); + } + if let Some(v) = values.stability_grace { + self.stability_grace_secs = Some(v.as_secs()); + } + } +} + +/// Explicitly-passed scan default values (from the CLI flags that were actually +/// given) — the input to [`UserSettings::persist_scan`]. +#[derive(Debug, Clone, Copy, Default)] +pub struct ScanDefaults { + pub max_workers: Option, + pub limit: Option, + pub timeout: Option, +} + +/// Explicitly-passed connect default values (from the CLI flags that were +/// actually given) — the input to [`UserSettings::persist_connect`]. +#[derive(Debug, Clone, Copy, Default)] +pub struct ConnectDefaults { + pub connect_timeout: Option, + pub cooldown: Option, + pub retry_count: Option, + pub stability_grace: Option, } #[cfg(test)] @@ -218,4 +268,30 @@ mod tests { assert_eq!(us.retry_count(Some(5)), 5); assert_eq!(us.stability_grace(Some(s(1))), s(1)); } + + #[test] + fn persist_scan_writes_only_passed_fields() { + let mut us = UserSettings::default(); + us.persist_scan(&ScanDefaults { + max_workers: Some(500), + timeout: Some(Duration::from_secs(20)), + ..Default::default() + }); + assert_eq!(us.max_workers, Some(500)); + assert_eq!(us.limit, None); // not passed -> untouched + assert_eq!(us.scan_timeout_secs, Some(20)); + } + + #[test] + fn persist_connect_writes_only_passed_fields() { + let mut us = UserSettings::default(); + us.persist_connect(&ConnectDefaults { + retry_count: Some(4), + cooldown: Some(Duration::from_secs(45)), + ..Default::default() + }); + assert_eq!(us.retry_count, Some(4)); + assert_eq!(us.cooldown_secs, Some(45)); + assert_eq!(us.connect_timeout_secs, None); + } } diff --git a/crates/vmate-core/src/system/killer.rs b/crates/vmate-core/src/system/killer.rs index b5e442b..e004339 100644 --- a/crates/vmate-core/src/system/killer.rs +++ b/crates/vmate-core/src/system/killer.rs @@ -1,6 +1,6 @@ //! Process killing. //! -//! vmate-cli tracks every OpenVPN process it spawns in a process-global PID +//! vmate-cli tracks every OpenVPN process it spawns in a per-session PID //! registry and, by default, cleans up exactly those processes on connection //! switching and shutdown: SIGTERM the process group, wait a grace period, then //! SIGKILL anything still alive. The exact `killall -9 openvpn` form preserved @@ -11,54 +11,66 @@ use anyhow::Result; use nix::sys::signal::{Signal, kill}; use nix::unistd::Pid; use std::process::Stdio; -use std::sync::Arc; -use std::sync::{LazyLock, Mutex}; +use std::sync::{Arc, Mutex}; use std::time::Duration; /// Grace period between SIGTERM and SIGKILL when killing a process tree. pub const KILL_GRACE: Duration = Duration::from_secs(3); -/// PIDs of every OpenVPN process vmate-cli has spawned this run. +/// PIDs of every OpenVPN process spawned in one session. /// -/// The registry is the default cleanup source: `CleanupGuard` kills exactly -/// these processes on drop, never the user's unrelated OpenVPN instances. -static REGISTRY: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); +/// Owned per run — and per test — so concurrent sessions never interfere: +/// each [`CleanupGuard`] sweeps only the registry it was given, never another +/// session's (or test's) live processes. +#[derive(Default)] +pub struct ProcessRegistry { + pids: Mutex>, +} -/// Record a spawned process pid so the cleanup guard can kill it on the way -/// out. A pid of 0 (unknown) is ignored. -pub fn register_process(pid: u32) { - if pid == 0 { - return; +impl ProcessRegistry { + pub fn new() -> Self { + Self::default() } - REGISTRY.lock().unwrap().push(pid); -} -/// Snapshot of the currently registered pids. -pub fn registered_pids() -> Vec { - REGISTRY.lock().unwrap().clone() -} + /// Record a spawned process pid so the cleanup guard can kill it on the + /// way out. A pid of 0 (unknown) is ignored. + pub fn register(&self, pid: u32) { + if pid == 0 { + return; + } + self.pids + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(pid); + } -/// Forget every registered pid without killing anything. -pub fn clear_registry() { - REGISTRY.lock().unwrap().clear(); -} + /// Snapshot of the currently registered pids. + pub fn registered(&self) -> Vec { + self.pids.lock().unwrap_or_else(|e| e.into_inner()).clone() + } -/// Kill every registered process group: SIGTERM all of them, allow one grace -/// period, then SIGKILL anything still alive, and finally clear the registry. -/// -/// Sync and best-effort: individual kill errors are ignored, and the grace -/// period is a fixed sleep because there is no live handle to wait on. This is -/// the last-resort safety net used by `CleanupGuard::drop`. -pub fn kill_all_spawned() { - let pids = registered_pids(); - for &pid in &pids { - let _ = kill_process_group(pid); + /// Forget every registered pid without killing anything. + pub fn clear(&self) { + self.pids.lock().unwrap_or_else(|e| e.into_inner()).clear(); } - std::thread::sleep(Duration::from_secs(1)); - for &pid in &pids { - let _ = force_kill_process_group(pid); + + /// Kill every registered process group: SIGTERM all of them, allow one + /// grace period, then SIGKILL anything still alive, then clear. + /// + /// Sync and best-effort: individual kill errors are ignored, and the grace + /// period is a fixed sleep because there is no live handle to wait on. This + /// is the last-resort safety net used by [`CleanupGuard::drop`]. + pub fn kill_all_graceful(&self) { + let pids = self.registered(); + for &pid in &pids { + let _ = kill_process_group(pid); + } + std::thread::sleep(Duration::from_secs(1)); + for &pid in &pids { + let _ = force_kill_process_group(pid); + } + self.clear(); } - clear_registry(); } /// SIGTERM the process group whose leader has the given pid. The child is @@ -165,30 +177,39 @@ pub async fn kill_process_tree_graceful( /// RAII guard that cleans up every spawned OpenVPN process on drop. /// /// This is the safety net that guarantees no stale OpenVPN processes survive -/// vmate, even on panic or error paths: the per-process registry is killed +/// vmate, even on panic or error paths: the session's registry is killed /// first, and the opt-in global `killall -9 openvpn` sweep runs afterwards. pub struct CleanupGuard { killer: Arc, + registry: Arc, enabled: bool, } impl CleanupGuard { - pub fn new(killer: Arc, enabled: bool) -> Self { - Self { killer, enabled } + pub fn new( + killer: Arc, + registry: Arc, + enabled: bool, + ) -> Self { + Self { + killer, + registry, + enabled, + } } /// Prevent the guard from killing anything on drop. The process registry is /// cleared so the guard leaves no stale pids behind. pub fn disarm(&mut self) { self.enabled = false; - clear_registry(); + self.registry.clear(); } } impl Drop for CleanupGuard { fn drop(&mut self) { if self.enabled { - kill_all_spawned(); + self.registry.kill_all_graceful(); let _ = self.killer.killall_openvpn(); } } @@ -215,7 +236,7 @@ mod tests { #[test] fn disabled_guard_does_nothing() { let killer: Arc = Arc::new(NoopKiller); - let mut guard = CleanupGuard::new(killer, false); + let mut guard = CleanupGuard::new(killer, Arc::new(ProcessRegistry::new()), false); guard.disarm(); // Dropping should be a no-op (would panic only if enabled flag is wrong). } @@ -228,16 +249,16 @@ mod tests { #[test] fn registry_registers_and_clears() { - clear_registry(); - register_process(0); // ignored - register_process(1234); - register_process(5678); - let pids = registered_pids(); + let reg = ProcessRegistry::new(); + reg.register(0); // ignored + reg.register(1234); + reg.register(5678); + let pids = reg.registered(); assert!(!pids.contains(&0)); assert!(pids.contains(&1234)); assert!(pids.contains(&5678)); - clear_registry(); - assert!(registered_pids().is_empty()); + reg.clear(); + assert!(reg.registered().is_empty()); } /// A real (now-reaped) pid and a large pid that no process group can own @@ -245,7 +266,7 @@ mod tests { /// registry is cleared, and nothing is left behind. #[tokio::test] async fn kill_all_spawned_ignores_dead_pids_and_clears() { - clear_registry(); + let reg = ProcessRegistry::new(); let mut child = tokio::process::Command::new("sh") .arg("-c") .arg("exit 0") @@ -253,12 +274,12 @@ mod tests { .unwrap(); let dead_pid = child.id().unwrap(); child.wait().await.unwrap(); - register_process(dead_pid); - register_process(999_999_999); + reg.register(dead_pid); + reg.register(999_999_999); - kill_all_spawned(); + reg.kill_all_graceful(); - assert!(registered_pids().is_empty()); + assert!(reg.registered().is_empty()); } #[test] diff --git a/crates/vmate-core/src/system/mod.rs b/crates/vmate-core/src/system/mod.rs index 55c4530..7b6cdff 100644 --- a/crates/vmate-core/src/system/mod.rs +++ b/crates/vmate-core/src/system/mod.rs @@ -5,9 +5,8 @@ pub mod root; pub mod signal; pub use killer::{ - CleanupGuard, ProcessKiller, RealProcessKiller, clear_registry, force_kill_process_group, - kill_all_spawned, kill_process_group, kill_process_tree_graceful, killall_openvpn, - register_process, + CleanupGuard, ProcessKiller, ProcessRegistry, RealProcessKiller, force_kill_process_group, + kill_process_group, kill_process_tree_graceful, killall_openvpn, }; pub use root::{elevate_with_sudo, is_root, require_root_for}; pub use signal::{ShutdownReason, shutdown_signal};