Added scaffoded project
This commit is contained in:
@@ -0,0 +1,398 @@
|
||||
use crate::{
|
||||
config::Config,
|
||||
db::Database,
|
||||
models::{ChunkRecord, PendingRange},
|
||||
scan::{build_pending_ranges, load_targets},
|
||||
terminal::{ScreenState, TerminalUi},
|
||||
};
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use std::{
|
||||
collections::VecDeque,
|
||||
net::Ipv4Addr,
|
||||
sync::{
|
||||
atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||
Arc, Mutex,
|
||||
},
|
||||
time::Instant,
|
||||
};
|
||||
use tokio::{
|
||||
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
|
||||
net::{TcpListener, TcpStream},
|
||||
task::JoinSet,
|
||||
};
|
||||
|
||||
pub async fn run_coordinator(
|
||||
config: Config,
|
||||
database: Arc<Database>,
|
||||
ui: Arc<TerminalUi>,
|
||||
interrupted: Arc<AtomicBool>,
|
||||
) -> Result<()> {
|
||||
let targets = Arc::new(load_targets(&config.cidrs)?);
|
||||
let target_count = targets.len();
|
||||
let chunk_records = database.load_chunk_records().await?;
|
||||
let (tested_before_run, _open_before_run) = database.summarize_chunks().await?;
|
||||
let pending_ranges = build_pending_ranges(target_count, config.chunk_size, &chunk_records);
|
||||
let pending_count: usize = pending_ranges.iter().map(|range| range.len()).sum();
|
||||
|
||||
if ui.is_enabled() {
|
||||
ui.print_banner();
|
||||
ui.render(&ScreenState {
|
||||
current_ip: if pending_count > 0 {
|
||||
targets[pending_ranges[0].start].to_string()
|
||||
} else {
|
||||
"concluído".to_string()
|
||||
},
|
||||
tested: tested_before_run,
|
||||
total: target_count,
|
||||
});
|
||||
} else {
|
||||
println!(
|
||||
"Coordinator ready on {} for {} hosts across {} CIDRs. Pending {}, already tested {}.",
|
||||
config.coordinator_bind,
|
||||
target_count,
|
||||
config.cidrs.len(),
|
||||
pending_count,
|
||||
tested_before_run
|
||||
);
|
||||
}
|
||||
|
||||
if pending_ranges.is_empty() {
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "concluído".to_string(),
|
||||
tested: target_count,
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let listener = TcpListener::bind(config.coordinator_bind)
|
||||
.await
|
||||
.with_context(|| format!("failed to bind {}", config.coordinator_bind))?;
|
||||
let queue = Arc::new(Mutex::new(VecDeque::from(pending_ranges)));
|
||||
let tested_this_run = Arc::new(AtomicUsize::new(0));
|
||||
let open_this_run = Arc::new(AtomicUsize::new(0));
|
||||
let start = Instant::now();
|
||||
let mut handlers = JoinSet::new();
|
||||
|
||||
loop {
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
break;
|
||||
}
|
||||
|
||||
tokio::select! {
|
||||
accept_result = listener.accept() => {
|
||||
let (stream, peer) = accept_result?;
|
||||
let queue = Arc::clone(&queue);
|
||||
let targets = Arc::clone(&targets);
|
||||
let database = Arc::clone(&database);
|
||||
let ui = Arc::clone(&ui);
|
||||
let interrupted = Arc::clone(&interrupted);
|
||||
let tested_this_run = Arc::clone(&tested_this_run);
|
||||
let open_this_run = Arc::clone(&open_this_run);
|
||||
let config = config.clone();
|
||||
|
||||
handlers.spawn(async move {
|
||||
handle_worker(
|
||||
stream,
|
||||
peer,
|
||||
queue,
|
||||
targets,
|
||||
database,
|
||||
ui,
|
||||
interrupted,
|
||||
tested_before_run,
|
||||
tested_this_run,
|
||||
open_this_run,
|
||||
config,
|
||||
)
|
||||
.await
|
||||
});
|
||||
}
|
||||
maybe_finished = handlers.join_next(), if !handlers.is_empty() => {
|
||||
if let Some(join_result) = maybe_finished {
|
||||
join_result??;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if queue.lock().expect("queue mutex poisoned").is_empty() && handlers.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
while let Some(join_result) = handlers.join_next().await {
|
||||
join_result??;
|
||||
}
|
||||
|
||||
let duration = start.elapsed().as_secs_f64();
|
||||
let total_open_count = database.count_open_endpoints(config.port).await?;
|
||||
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "interrompido".to_string(),
|
||||
tested: tested_before_run + tested_this_run.load(Ordering::SeqCst),
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
} else {
|
||||
println!("Interrupted after {:.2}s.", duration);
|
||||
println!(
|
||||
"Hosts tested in this run: {}",
|
||||
tested_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!(
|
||||
"Open endpoints found: {}",
|
||||
open_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if ui.is_enabled() {
|
||||
ui.render(&ScreenState {
|
||||
current_ip: "concluído".to_string(),
|
||||
tested: target_count,
|
||||
total: target_count,
|
||||
});
|
||||
ui.finish();
|
||||
} else {
|
||||
println!("Coordinator finished in {:.2}s", duration);
|
||||
println!(
|
||||
"Hosts tested in this run: {}",
|
||||
tested_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!(
|
||||
"Open endpoints found in this run: {}",
|
||||
open_this_run.load(Ordering::SeqCst)
|
||||
);
|
||||
println!("Total open endpoints in DB: {}", total_open_count);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_worker(
|
||||
stream: TcpStream,
|
||||
peer: std::net::SocketAddr,
|
||||
queue: Arc<Mutex<VecDeque<PendingRange>>>,
|
||||
targets: Arc<Vec<Ipv4Addr>>,
|
||||
database: Arc<Database>,
|
||||
ui: Arc<TerminalUi>,
|
||||
interrupted: Arc<AtomicBool>,
|
||||
tested_before_run: usize,
|
||||
tested_this_run: Arc<AtomicUsize>,
|
||||
open_this_run: Arc<AtomicUsize>,
|
||||
config: Config,
|
||||
) -> Result<()> {
|
||||
let (read_half, mut write_half) = stream.into_split();
|
||||
let mut lines = BufReader::new(read_half).lines();
|
||||
|
||||
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}"));
|
||||
}
|
||||
None => return Ok(()),
|
||||
};
|
||||
|
||||
if !worker_name.is_empty() {
|
||||
println!("worker connected from {peer}: {worker_name}");
|
||||
} else {
|
||||
println!("worker connected from {peer}");
|
||||
}
|
||||
|
||||
loop {
|
||||
if interrupted.load(Ordering::SeqCst) {
|
||||
let _ = write_half.write_all(b"SHUTDOWN\n").await;
|
||||
break;
|
||||
}
|
||||
|
||||
let Some(range) = next_range(&queue) else {
|
||||
let _ = write_half.write_all(b"DONE\n").await;
|
||||
break;
|
||||
};
|
||||
|
||||
if let Err(err) = process_range(
|
||||
&mut write_half,
|
||||
&mut lines,
|
||||
range,
|
||||
&targets,
|
||||
&database,
|
||||
&ui,
|
||||
tested_before_run,
|
||||
&tested_this_run,
|
||||
&open_this_run,
|
||||
&config,
|
||||
)
|
||||
.await
|
||||
{
|
||||
requeue_range(&queue, range);
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process_range(
|
||||
write_half: &mut tokio::net::tcp::OwnedWriteHalf,
|
||||
lines: &mut tokio::io::Lines<BufReader<tokio::net::tcp::OwnedReadHalf>>,
|
||||
range: PendingRange,
|
||||
targets: &Arc<Vec<Ipv4Addr>>,
|
||||
database: &Arc<Database>,
|
||||
ui: &Arc<TerminalUi>,
|
||||
tested_before_run: usize,
|
||||
tested_this_run: &Arc<AtomicUsize>,
|
||||
open_this_run: &Arc<AtomicUsize>,
|
||||
config: &Config,
|
||||
) -> Result<()> {
|
||||
let ips = targets[range.start..range.end].to_vec();
|
||||
let job_line = format!(
|
||||
"JOB {} {} {} {} {}\n",
|
||||
range.start,
|
||||
range.end,
|
||||
config.port,
|
||||
config.timeout_ms,
|
||||
format_ips(&ips)
|
||||
);
|
||||
write_half.write_all(job_line.as_bytes()).await?;
|
||||
|
||||
let Some(result_line) = lines.next_line().await? else {
|
||||
return Err(anyhow!("worker disconnected before returning job"));
|
||||
};
|
||||
|
||||
let result = parse_result_line(&result_line)?;
|
||||
if result.chunk_start != range.start || result.chunk_end != range.end {
|
||||
return Err(anyhow!(
|
||||
"worker returned mismatched range: expected {}..{}, got {}..{}",
|
||||
range.start,
|
||||
range.end,
|
||||
result.chunk_start,
|
||||
result.chunk_end
|
||||
));
|
||||
}
|
||||
|
||||
for ip in result.open_ips {
|
||||
database.upsert_open_endpoint(ip, config.port).await?;
|
||||
}
|
||||
|
||||
let previous = database.load_chunk_record(range.chunk_start).await?;
|
||||
let tested_before = previous
|
||||
.as_ref()
|
||||
.map(|record| record.tested_count)
|
||||
.unwrap_or(0);
|
||||
let open_before = previous
|
||||
.as_ref()
|
||||
.map(|record| record.open_count)
|
||||
.unwrap_or(0);
|
||||
let tested_count = tested_before + result.tested;
|
||||
let open_count = open_before + result.open_count;
|
||||
database
|
||||
.save_chunk_record(&ChunkRecord {
|
||||
chunk_start: range.chunk_start,
|
||||
chunk_end: range.chunk_end,
|
||||
tested_count,
|
||||
open_count,
|
||||
is_complete: result.tested >= range.len(),
|
||||
})
|
||||
.await?;
|
||||
|
||||
let tested_now = tested_this_run.fetch_add(result.tested, Ordering::SeqCst) + result.tested;
|
||||
open_this_run.fetch_add(result.open_count, Ordering::SeqCst);
|
||||
|
||||
if ui.is_enabled() {
|
||||
let current_ip = ips
|
||||
.last()
|
||||
.map(|ip| ip.to_string())
|
||||
.unwrap_or_else(|| "concluído".to_string());
|
||||
ui.render(&ScreenState {
|
||||
current_ip,
|
||||
tested: tested_before_run + tested_now,
|
||||
total: targets.len(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn requeue_range(queue: &Arc<Mutex<VecDeque<PendingRange>>>, range: PendingRange) {
|
||||
queue
|
||||
.lock()
|
||||
.expect("queue mutex poisoned")
|
||||
.push_front(range);
|
||||
}
|
||||
|
||||
fn next_range(queue: &Arc<Mutex<VecDeque<PendingRange>>>) -> Option<PendingRange> {
|
||||
queue.lock().expect("queue mutex poisoned").pop_front()
|
||||
}
|
||||
|
||||
struct JobResult {
|
||||
chunk_start: usize,
|
||||
chunk_end: usize,
|
||||
tested: usize,
|
||||
open_count: usize,
|
||||
open_ips: Vec<Ipv4Addr>,
|
||||
}
|
||||
|
||||
fn parse_result_line(line: &str) -> Result<JobResult> {
|
||||
let mut parts = line.splitn(6, ' ');
|
||||
let kind = parts.next().unwrap_or("");
|
||||
if kind != "RESULT" {
|
||||
return Err(anyhow!("unexpected worker message: {line}"));
|
||||
}
|
||||
|
||||
let chunk_start = parse_usize_part(parts.next(), "chunk_start", line)?;
|
||||
let chunk_end = parse_usize_part(parts.next(), "chunk_end", line)?;
|
||||
let tested = parse_usize_part(parts.next(), "tested", line)?;
|
||||
let open_count = parse_usize_part(parts.next(), "open_count", line)?;
|
||||
let ips_part = parts.next().unwrap_or("");
|
||||
let open_ips = if ips_part.trim().is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
ips_part
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|entry| !entry.is_empty())
|
||||
.map(|entry| {
|
||||
entry
|
||||
.parse::<Ipv4Addr>()
|
||||
.with_context(|| format!("invalid IP in worker result: {entry}"))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
};
|
||||
|
||||
if open_ips.len() != open_count {
|
||||
return Err(anyhow!(
|
||||
"open_count mismatch in worker result: declared {}, received {}",
|
||||
open_count,
|
||||
open_ips.len()
|
||||
));
|
||||
}
|
||||
|
||||
Ok(JobResult {
|
||||
chunk_start,
|
||||
chunk_end,
|
||||
tested,
|
||||
open_count,
|
||||
open_ips,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_usize_part(part: Option<&str>, name: &str, line: &str) -> Result<usize> {
|
||||
let value = part.ok_or_else(|| anyhow!("missing {name} in worker message: {line}"))?;
|
||||
value
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| anyhow!("invalid {name} in worker message: {line}"))
|
||||
}
|
||||
|
||||
fn format_ips(ips: &[Ipv4Addr]) -> String {
|
||||
ips.iter()
|
||||
.map(ToString::to_string)
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
}
|
||||
Reference in New Issue
Block a user