From 747f1ea2b791cdbf2c711c7559a54dca5a7faeed Mon Sep 17 00:00:00 2001 From: Kevin Valerio Date: Fri, 27 Sep 2024 16:32:46 +0200 Subject: [PATCH] Killing child when exiting --- Cargo.toml | 3 +- src/cli/ui/custom.rs | 7 +-- src/cli/ui/monitor/corpus.rs | 11 +--- src/cli/ui/monitor/logs.rs | 5 +- src/cli/ui/ratatui.rs | 87 +++++++++++++++--------------- tests/cli_fuzz_integration_test.rs | 5 +- 6 files changed, 58 insertions(+), 60 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 9ed2b98..353cd03 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,7 +66,8 @@ quote = { version = "1.0.36" } proc-macro2 = { version = "1.0.81", features = ["span-locations"] } rand = { version = "0.8" } walkdir = { version = "2.5.0" } -regex = "1.10.4" +ctrlc = { version = "3.4.5" } +regex = { version = "1.10.4" } ratatui = { version = "0.28.0", features = ["all-widgets"] } crossterm = { version = "0.28.1" } chrono = { version = "0.4" } diff --git a/src/cli/ui/custom.rs b/src/cli/ui/custom.rs index e99b5be..ede6e7b 100644 --- a/src/cli/ui/custom.rs +++ b/src/cli/ui/custom.rs @@ -71,11 +71,12 @@ impl CustomManager { tx.send(fuzzer_pid).unwrap(); }); - let pid: Child = rx.recv()??; + let child: Child = rx.recv()??; - let ratatui = CustomUI::new(&cloned_config).context("Couldn't create the custom UI ")?; + let mut ratatui = + CustomUI::new(&cloned_config).context("Couldn't create the custom UI ")?; - ratatui.initialize_tui(pid)?; + ratatui.initialize_tui(child)?; Ok(()) } } diff --git a/src/cli/ui/monitor/corpus.rs b/src/cli/ui/monitor/corpus.rs index f721d0d..eb4d930 100644 --- a/src/cli/ui/monitor/corpus.rs +++ b/src/cli/ui/monitor/corpus.rs @@ -1,14 +1,7 @@ use crate::cli::{ - config::{ - PFiles, - PhinkFiles, - }, - ui::{ - monitor::logs::AFLDashboard, - traits::FromPath, - }, + config::PFiles, + ui::traits::FromPath, }; -use anyhow::bail; use std::{ fs, path::PathBuf, diff --git a/src/cli/ui/monitor/logs.rs b/src/cli/ui/monitor/logs.rs index ae3d27d..1d5d9e7 100644 --- a/src/cli/ui/monitor/logs.rs +++ b/src/cli/ui/monitor/logs.rs @@ -1,8 +1,5 @@ use crate::cli::{ - config::{ - PFiles, - PhinkFiles, - }, + config::PFiles, ui::traits::FromPath, }; use anyhow::bail; diff --git a/src/cli/ui/ratatui.rs b/src/cli/ui/ratatui.rs index a008ae4..4ae0982 100644 --- a/src/cli/ui/ratatui.rs +++ b/src/cli/ui/ratatui.rs @@ -17,7 +17,6 @@ use crate::cli::{ ziggy::ZiggyConfig, }; use anyhow::Context; -use crossterm::event::KeyCode; use ratatui::{ crossterm::event::{ self, @@ -51,8 +50,13 @@ use ratatui::{ use std::{ borrow::Borrow, process::Child, - sync::mpsc, - thread, + sync::{ + atomic::{ + AtomicBool, + Ordering, + }, + Arc, + }, time::Duration, }; @@ -76,7 +80,7 @@ impl CustomUI { }) } - fn ui(self, f: &mut Frame) -> anyhow::Result<()> { + fn ui(&mut self, f: &mut Frame) -> anyhow::Result<()> { let chunks = Layout::default() .direction(Direction::Vertical) .margin(1) @@ -91,25 +95,25 @@ impl CustomUI { ) .split(f.area()); - self.clone().render_title(f, chunks[0]); - self.clone().render_stats(f, chunks[1]); - self.clone().render_chart_and_config(f, chunks[2]); - self.clone().render_seed(f, chunks[3])?; + self.render_title(f, chunks[0]); + self.render_stats(f, chunks[1]); + self.render_chart_and_config(f, chunks[2]); + self.render_seed(f, chunks[3])?; Ok(()) } - fn render_chart_and_config(self, f: &mut Frame, area: Rect) { + fn render_chart_and_config(&mut self, f: &mut Frame, area: Rect) { let chunks = Layout::default() .direction(Direction::Horizontal) .margin(1) .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref()) .split(area); - self.clone().render_chart(f, chunks[0]); + self.render_chart(f, chunks[0]); self.ziggy_config.config.render(f, chunks[1]); } - fn render_octopus(self, f: &mut Frame, area: Rect) { + fn render_octopus(&self, f: &mut Frame, area: Rect) { let ascii_art = r#" ,---. ( @ @ ) @@ -128,7 +132,7 @@ impl CustomUI { f.render_widget(octopus, area); } - fn render_title(self, f: &mut Frame, area: Rect) { + fn render_title(&self, f: &mut Frame, area: Rect) { self.render_octopus(f, area); let title = Paragraph::new("Phink Fuzzing Dashboard") .style( @@ -140,7 +144,7 @@ impl CustomUI { f.render_widget(title, area); } - fn render_stats(self, f: &mut Frame, area: Rect) { + fn render_stats(&self, f: &mut Frame, area: Rect) { let data = self.afl_dashboard.read_properties(); if let Ok(afl) = data { @@ -271,7 +275,7 @@ impl CustomUI { crash_style } - fn render_chart(mut self, f: &mut Frame, area: Rect) { + fn render_chart(&mut self, f: &mut Frame, area: Rect) { let chunks = Layout::default() .direction(Direction::Vertical) .constraints([Constraint::Percentage(100)].as_ref()) @@ -283,8 +287,8 @@ impl CustomUI { f.render_widget(chart_manager.create_chart(), chunks[0]); } - fn render_seed(self, f: &mut Frame, area: Rect) -> anyhow::Result<()> { - let seed_displayer = SeedDisplayer::new(self.ziggy_config.fuzz_output()); + fn render_seed(&self, f: &mut Frame, area: Rect) -> anyhow::Result<()> { + let seed_displayer = SeedDisplayer::new(self.clone().ziggy_config.fuzz_output()); let mut seed_info_text: String = String::new(); if let Some(seeds) = seed_displayer.load() { @@ -306,45 +310,44 @@ impl CustomUI { Ok(()) } - pub fn initialize_tui(&self, mut child: Child) -> anyhow::Result<()> { + pub fn initialize_tui(&mut self, mut child: Child) -> anyhow::Result<()> { let stdout = std::io::stdout(); let backend = ratatui::backend::CrosstermBackend::new(stdout); let mut terminal = ratatui::Terminal::new(backend)?; terminal.clear()?; - let (tx, rx) = mpsc::channel(); - - let input_handle = thread::spawn(move || { - loop { - if event::poll(Duration::from_millis(100)).unwrap() { - if let Event::Key(key) = event::read().unwrap() { - // If we CTRL+C, we send the exit signal - if key.code == KeyCode::Char('c') - && key.modifiers.contains(event::KeyModifiers::CONTROL) - { - tx.send(()).unwrap(); - break; - } - } - } - } - }); + let running = Arc::new(AtomicBool::new(true)); + let r = running.clone(); - loop { - if rx.try_recv().is_ok() { - break; - } + ctrlc::set_handler(move || { + r.store(false, Ordering::SeqCst); + })?; + + while running.load(Ordering::SeqCst) { terminal.draw(|f| { - if let Err(err) = self.clone().ui(f) { + if let Err(err) = self.ui(f) { eprintln!("{:?}", err); } })?; - thread::sleep(Duration::from_millis(Self::REFRESH_MS)); - } + if event::poll(Duration::from_millis(Self::REFRESH_MS))? { + if let Event::Key(key) = crossterm::event::read()? { + if key.kind == event::KeyEventKind::Press + && key.code == event::KeyCode::Char('q') + { + let _ = child.kill(); + break; + } + } + } + } let _ = child.kill(); - input_handle.join().unwrap(); + + println!( + "👋 It was nice fuzzing with you. Killing PID {}. Bye bye! ", + child.id() + ); terminal.clear()?; Ok(()) diff --git a/tests/cli_fuzz_integration_test.rs b/tests/cli_fuzz_integration_test.rs index 9a4ec28..c975294 100644 --- a/tests/cli_fuzz_integration_test.rs +++ b/tests/cli_fuzz_integration_test.rs @@ -24,7 +24,10 @@ mod tests { instrumenter::path::InstrumentedPath, }; - use phink_lib::cli::ui::monitor::logs::AFLDashboard; + use phink_lib::cli::ui::{ + monitor::logs::AFLDashboard, + traits::FromPath, + }; use std::{ fs, time::Duration,