From d0796681ef9fd095046ebcd782f225f4494d6775 Mon Sep 17 00:00:00 2001 From: Almaz Date: Tue, 15 Jul 2025 11:01:41 +0000 Subject: [PATCH] feat: support exiting tui --- Cargo.lock | 1 + Cargo.toml | 9 +- src/main.rs | 247 ++++++++++++++++++++++++++++++++++++---------------- src/tui.rs | 37 +++----- 4 files changed, 192 insertions(+), 102 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 88eb207..59c1f57 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1346,6 +1346,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "tokio-util", "toml", "tracing", "tracing-log", diff --git a/Cargo.toml b/Cargo.toml index 3a3d947..a23a314 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ rlimit = "=0.10.2" serde = "=1.0.219" serde_json = "=1.0.140" tokio = { version = "=1.46.1", features = ["full"] } +tokio-util = { version = "=0.7.15", optional = true } toml = "=0.9.2" tracing = "=0.1.41" tracing-log = "=0.2.0" @@ -30,7 +31,13 @@ tui-logger = { version = "=0.17.3", features = [ url = "=2.5.4" [features] -tui = ["dep:crossterm", "dep:futures", "dep:ratatui", "dep:tui-logger"] +tui = [ + "dep:crossterm", + "dep:futures", + "dep:ratatui", + "dep:tokio-util", + "dep:tui-logger", +] [target.'cfg(not(target_os = "android"))'.dependencies] reqwest = { version = "=0.12.22", default-features = false, features = [ diff --git a/src/main.rs b/src/main.rs index 2cef90f..fc15966 100644 --- a/src/main.rs +++ b/src/main.rs @@ -66,67 +66,45 @@ fn create_reqwest_client() -> reqwest::Result { .build() } -#[tokio::main] -async fn main() -> color_eyre::Result<()> { - color_eyre::install().wrap_err("failed to install color_eyre hooks")?; +fn create_logging_filter( + config: &config::Config, +) -> tracing_subscriber::filter::Targets { + let base = tracing_subscriber::filter::Targets::new() + .with_default(tracing::level_filters::LevelFilter::INFO) + .with_target( + "hickory_proto::udp::udp_client_stream", + tracing::level_filters::LevelFilter::ERROR, + ) + .with_target( + // TODO: remove for hickory_proto >= 0.25.0 + "hickory_proto::xfer::dns_exchange", + tracing::level_filters::LevelFilter::ERROR, + ); - let raw_config_path = raw_config::get_config_path(); - let raw_config = raw_config::read_config(std::path::Path::new( - &raw_config_path, - )) - .await - .wrap_err_with(move || format!("failed to read {raw_config_path}"))?; - - let config = Arc::new( - config::Config::from_raw_config(raw_config) - .await - .wrap_err("failed to create Config from RawConfig")?, - ); - - let targets_filter = { - let base = tracing_subscriber::filter::Targets::new() - .with_default(tracing::level_filters::LevelFilter::INFO) - .with_target( - "hickory_proto::udp::udp_client_stream", - tracing::level_filters::LevelFilter::ERROR, - ) - .with_target( - // TODO: remove for hickory_proto >= 0.25.0 - "hickory_proto::xfer::dns_exchange", - tracing::level_filters::LevelFilter::ERROR, - ); - if config.debug { - base.with_target( - "proxy_scraper_checker::checker", - tracing::level_filters::LevelFilter::DEBUG, - ) - } else { - base - } - }; - - #[cfg(feature = "tui")] - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - - #[cfg(feature = "tui")] - let tui_task = - tokio::task::spawn(tui::Tui::new(targets_filter)?.run(tx.clone(), rx)); - - #[cfg(not(feature = "tui"))] - tracing_subscriber::registry() - .with(targets_filter) - .with(tracing_subscriber::fmt::layer()) - .init(); - - let http_client = create_reqwest_client() - .wrap_err("failed to create reqwest HTTP client")?; + if config.debug { + base.with_target( + "proxy_scraper_checker::checker", + tracing::level_filters::LevelFilter::DEBUG, + ) + } else { + base + } +} +fn spawn_ip_database_tasks( + config: &Arc, + http_client: &reqwest::Client, + #[cfg(feature = "tui")] tx: &tokio::sync::mpsc::UnboundedSender< + event::Event, + >, +) -> tokio::task::JoinSet> { let mut output_dependencies_tasks = tokio::task::JoinSet::new(); if config.asn_enabled() { let http_client = http_client.clone(); #[cfg(feature = "tui")] let tx = tx.clone(); + output_dependencies_tasks.spawn(async move { ipdb::DbType::Asn .download( @@ -142,6 +120,7 @@ async fn main() -> color_eyre::Result<()> { let http_client = http_client.clone(); #[cfg(feature = "tui")] let tx = tx.clone(); + output_dependencies_tasks.spawn(async move { ipdb::DbType::Geo .download( @@ -153,6 +132,58 @@ async fn main() -> color_eyre::Result<()> { }); } + output_dependencies_tasks +} + +async fn wait_for_ip_database_tasks( + mut tasks: tokio::task::JoinSet>, +) -> color_eyre::Result<()> { + while let Some(task) = tasks.join_next().await { + task.wrap_err("output dependencies task panicked or was cancelled")??; + } + Ok(()) +} + +async fn process_proxies( + config: Arc, + proxies: std::collections::HashSet, + #[cfg(feature = "tui")] tx: &tokio::sync::mpsc::UnboundedSender< + event::Event, + >, +) -> color_eyre::Result> { + if config.checking.check_url.is_empty() { + Ok(proxies.into_iter().collect()) + } else { + #[cfg(not(feature = "tui"))] + tracing::info!("Started checking {} proxies", proxies.len()); + + checker::check_all( + config, + proxies, + #[cfg(feature = "tui")] + tx.clone(), + ) + .await + .wrap_err("failed to check proxies") + } +} + +async fn main_task( + config: Arc, + #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender< + event::Event, + >, +) -> color_eyre::Result<()> { + let http_client = create_reqwest_client() + .wrap_err("failed to create reqwest HTTP client")?; + + let ip_db_tasks = spawn_ip_database_tasks( + &config, + &http_client, + #[cfg(feature = "tui")] + &tx, + ); + let proxies = scraper::scrape_all( Arc::clone(&config), http_client.clone(), @@ -164,26 +195,15 @@ async fn main() -> color_eyre::Result<()> { drop(http_client); - while let Some(task) = output_dependencies_tasks.join_next().await { - task.wrap_err("output dependencies task panicked or was cancelled")??; - } - drop(output_dependencies_tasks); + wait_for_ip_database_tasks(ip_db_tasks).await?; - let proxies = if config.checking.check_url.is_empty() { - proxies.into_iter().collect() - } else { - #[cfg(not(feature = "tui"))] - tracing::info!("Started checking {} proxies", proxies.len()); - - checker::check_all( - Arc::clone(&config), - proxies, - #[cfg(feature = "tui")] - tx.clone(), - ) - .await - .wrap_err("failed to check proxies")? - }; + let proxies = process_proxies( + Arc::clone(&config), + proxies, + #[cfg(feature = "tui")] + &tx, + ) + .await?; output::save_proxies(config, proxies) .await @@ -194,11 +214,88 @@ async fn main() -> color_eyre::Result<()> { #[cfg(feature = "tui")] tx.send(event::Event::App(event::AppEvent::Done))?; - #[cfg(feature = "tui")] - drop(tx); + Ok(()) +} + +#[cfg(feature = "tui")] +async fn cancellable_main( + config: Arc, + tx: tokio::sync::mpsc::UnboundedSender, + token: tokio_util::sync::CancellationToken, +) -> color_eyre::Result<()> { + #[expect(clippy::integer_division_remainder_used)] + { + tokio::select! { + biased; + () = token.cancelled() => Ok(()), + r = main_task(config, tx) => r + } + } +} + +#[cfg(feature = "tui")] +async fn run_with_tui( + config: Arc, + logging_filter: tracing_subscriber::filter::Targets, +) -> color_eyre::Result<()> { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + + let tui_task = + tokio::task::spawn(tui::Tui::new(logging_filter)?.run(tx.clone(), rx)); + + let token = tokio_util::sync::CancellationToken::new(); + let main_task = + tokio::task::spawn(cancellable_main(config, tx, token.clone())); - #[cfg(feature = "tui")] tui_task.await??; + token.cancel(); + drop(token); + main_task.await??; Ok(()) } + +#[cfg(not(feature = "tui"))] +async fn run_without_tui( + config: Arc, + logging_filter: tracing_subscriber::filter::Targets, +) -> color_eyre::Result<()> { + tracing_subscriber::registry() + .with(logging_filter) + .with(tracing_subscriber::fmt::layer()) + .init(); + + main_task(config).await +} + +async fn load_config() -> color_eyre::Result> { + let raw_config_path = raw_config::get_config_path(); + let raw_config = raw_config::read_config(std::path::Path::new( + &raw_config_path, + )) + .await + .wrap_err_with(move || format!("failed to read {raw_config_path}"))?; + + let config = config::Config::from_raw_config(raw_config) + .await + .wrap_err("failed to create Config from RawConfig")?; + + Ok(Arc::new(config)) +} + +#[tokio::main] +async fn main() -> color_eyre::Result<()> { + color_eyre::install().wrap_err("failed to install color_eyre hooks")?; + + let config = load_config().await?; + let logging_filter = create_logging_filter(&config); + + #[cfg(feature = "tui")] + { + run_with_tui(config, logging_filter).await + } + #[cfg(not(feature = "tui"))] + { + run_without_tui(config, logging_filter).await + } +} diff --git a/src/tui.rs b/src/tui.rs index ce0717a..94706c2 100644 --- a/src/tui.rs +++ b/src/tui.rs @@ -95,9 +95,6 @@ impl Drop for Tui { pub enum AppMode { #[default] Running, - /// Wait for the user confirmation to close the UI - Done, - /// Close the UI Quit, } @@ -301,17 +298,12 @@ fn draw(f: &mut Frame, state: &AppState, logger_state: &TuiWidgetState) { working_proxies_block.inner(layout[2]), ); } - - let done = matches!(state.mode, AppMode::Done); - let mut lines = Vec::with_capacity(usize::from(done).saturating_add(2)); - lines.push(Line::from("Up/PageUp/k - scroll logs up")); - lines.push(Line::from("Down/PageDown/j - scroll logs down")); - if done { - lines.push( - Line::from("Enter/ESC/q/Ctrl-C - exit") - .style(Style::default().fg(Color::Red)), - ); - } + let lines = vec![ + Line::from("Up/PageUp/k - scroll logs up"), + Line::from("Down/PageDown/j - scroll logs down"), + Line::from("ESC/q/Ctrl-C - exit") + .style(Style::default().fg(Color::Red)), + ]; f.render_widget(Text::from(lines).centered(), outer_layout[3]); } @@ -329,16 +321,11 @@ async fn handle_event( Event::Crossterm(crossterm_event) => { match crossterm_event { CrosstermEvent::Key(key_event) => match key_event.code { - KeyCode::Enter - | KeyCode::Esc - | KeyCode::Char('q' | 'Q') - if matches!(state.mode, AppMode::Done) => - { + KeyCode::Esc | KeyCode::Char('q' | 'Q') => { state.mode = AppMode::Quit; } KeyCode::Char('c' | 'C') - if key_event.modifiers == KeyModifiers::CONTROL - && matches!(state.mode, AppMode::Done) => + if key_event.modifiers == KeyModifiers::CONTROL => { state.mode = AppMode::Quit; } @@ -407,11 +394,9 @@ async fn handle_event( .or_insert(1); } AppEvent::Done => { - state.mode = if is_interactive().await { - AppMode::Done - } else { - AppMode::Quit - }; + if !is_interactive().await { + state.mode = AppMode::Quit; + } } } false