diff --git a/src/coordinator.rs b/src/coordinator.rs index b1012f1..62d2430 100644 --- a/src/coordinator.rs +++ b/src/coordinator.rs @@ -113,7 +113,7 @@ pub async fn run_coordinator( } maybe_finished = handlers.join_next(), if !handlers.is_empty() => { if let Some(join_result) = maybe_finished { - join_result??; + log_worker_result(join_result); } } } @@ -124,7 +124,7 @@ pub async fn run_coordinator( } while let Some(join_result) = handlers.join_next().await { - join_result??; + log_worker_result(join_result); } let duration = start.elapsed().as_secs_f64(); @@ -175,6 +175,14 @@ pub async fn run_coordinator( Ok(()) } +fn log_worker_result(result: std::result::Result, tokio::task::JoinError>) { + match result { + Ok(Ok(())) => {} + Ok(Err(err)) => eprintln!("worker handler ended with error: {err:#}"), + Err(err) => eprintln!("worker task failed: {err}"), + } +} + async fn handle_worker( stream: TcpStream, peer: std::net::SocketAddr, @@ -194,7 +202,8 @@ async fn handle_worker( let worker_name = match lines.next_line().await? { Some(line) if line.starts_with("HELLO ") => line[6..].trim().to_string(), Some(line) => { - return Err(anyhow!("unexpected worker handshake from {peer}: {line}")); + eprintln!("rejected non-worker connection from {peer}: {line}"); + return Ok(()); } None => return Ok(()), }; diff --git a/src/worker.rs b/src/worker.rs index afe1491..2bfb17c 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -1,9 +1,13 @@ -use crate::{config::Config, scan::check_port_open, terminal::TerminalUi}; +use crate::{ + config::Config, + scan::check_port_open, + terminal::{ScreenState, TerminalUi}, +}; use anyhow::{anyhow, Context, Result}; use std::{net::Ipv4Addr, sync::Arc}; use tokio::{ io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, - net::TcpStream, + net::{tcp::OwnedWriteHalf, TcpStream}, task::JoinSet, time::sleep, }; @@ -13,6 +17,36 @@ pub async fn run_worker(config: Config, ui: Arc) -> Result<()> { ui.print_banner(); } + loop { + render_worker_status(&ui, "conectando", 0, 0); + + match run_worker_session(&config, &ui).await? { + WorkerSession::Done => break, + WorkerSession::Disconnected => { + eprintln!("coordinator connection lost; reconnecting..."); + sleep(std::time::Duration::from_secs(1)).await; + } + } + } + + if ui.is_enabled() { + ui.render(&ScreenState { + current_ip: "concluĂ­do".to_string(), + tested: 0, + total: 0, + }); + ui.finish(); + } + + Ok(()) +} + +enum WorkerSession { + Done, + Disconnected, +} + +async fn run_worker_session(config: &Config, ui: &TerminalUi) -> Result { let stream = connect_with_retry(&config.coordinator_addr).await?; let (read_half, mut write_half) = stream.into_split(); let mut lines = BufReader::new(read_half).lines(); @@ -20,34 +54,50 @@ pub async fn run_worker(config: Config, ui: Arc) -> Result<()> { .write_all(format!("HELLO {}\n", config.worker_name).as_bytes()) .await?; + render_worker_status(ui, "aguardando job", 0, 0); + loop { let Some(line) = lines.next_line().await? else { - break; + return Ok(WorkerSession::Disconnected); }; - if line == "DONE" { - break; - } - - if line == "SHUTDOWN" { - break; + if line == "DONE" || line == "SHUTDOWN" { + return Ok(WorkerSession::Done); } let job = parse_job_line(&line)?; - let result = - execute_job(&job.ips, config.port, config.timeout_ms, config.concurrency).await?; + render_job_start(ui, &job); + let result = execute_job( + &job.ips, + config.port, + config.timeout_ms, + config.concurrency, + ui, + ) + .await?; - let result_line = format!( - "RESULT {} {} {} {} {}\n", - job.chunk_start, - job.chunk_end, - result.tested, - result.open_ips.len(), - format_ips(&result.open_ips) - ); - write_half.write_all(result_line.as_bytes()).await?; + if write_result(&mut write_half, &job, &result).await.is_err() { + return Ok(WorkerSession::Disconnected); + } + + render_worker_status(ui, "aguardando job", result.tested, job.ips.len()); } +} +async fn write_result( + write_half: &mut OwnedWriteHalf, + job: &Job, + result: &JobOutcome, +) -> Result<()> { + let result_line = format!( + "RESULT {} {} {} {} {}\n", + job.chunk_start, + job.chunk_end, + result.tested, + result.open_ips.len(), + format_ips(&result.open_ips) + ); + write_half.write_all(result_line.as_bytes()).await?; Ok(()) } @@ -114,11 +164,36 @@ fn parse_job_line(line: &str) -> Result { }) } +fn render_job_start(ui: &TerminalUi, job: &Job) { + let current_ip = job + .ips + .first() + .map(ToString::to_string) + .unwrap_or_else(|| "sem IP".to_string()); + render_worker_status(ui, ¤t_ip, 0, job.ips.len()); +} + +fn render_worker_status(ui: &TerminalUi, current_ip: &str, tested: usize, total: usize) { + if ui.is_enabled() { + let display_total = if total == 0 && current_ip != "concluĂ­do" { + 1 + } else { + total + }; + ui.render(&ScreenState { + current_ip: current_ip.to_string(), + tested, + total: display_total, + }); + } +} + async fn execute_job( ips: &[Ipv4Addr], port: u16, timeout_ms: u64, concurrency: usize, + ui: &TerminalUi, ) -> Result { let mut tested = 0usize; let mut open_ips = Vec::new(); @@ -139,6 +214,12 @@ async fn execute_job( if is_open { open_ips.push(ip); } + + ui.render(&ScreenState { + current_ip: ip.to_string(), + tested, + total: ips.len(), + }); } index = end;