From 165b9e098689096e23f1cee1ae79cb2d5308c16b Mon Sep 17 00:00:00 2001 From: Kevin Valerio Date: Fri, 27 Sep 2024 18:08:04 +0200 Subject: [PATCH] Removing gauge and adding chart bar --- src/cli/ui/configtable.rs | 2 +- src/cli/ui/custom.rs | 29 ++++++-- src/cli/ui/ratatui.rs | 78 +++++++++----------- src/cli/ui/seed.rs | 87 +++++++++++++---------- src/fuzzer/fuzz.rs | 6 +- src/fuzzer/parser.rs | 10 +++ tests/fixtures/phink/logs/last_seed.phink | 2 + 7 files changed, 123 insertions(+), 91 deletions(-) create mode 100644 tests/fixtures/phink/logs/last_seed.phink diff --git a/src/cli/ui/configtable.rs b/src/cli/ui/configtable.rs index 796d02f..4acd008 100644 --- a/src/cli/ui/configtable.rs +++ b/src/cli/ui/configtable.rs @@ -130,7 +130,7 @@ impl Paint for Configuration { .title("Configuration"), ) .highlight_style(selected_style) - .widths([Constraint::Percentage(25), Constraint::Percentage(60)]) + .widths([Constraint::Percentage(40), Constraint::Percentage(60)]) .column_spacing(1) .highlight_style(Style::default().add_modifier(Modifier::BOLD)) .highlight_symbol("> ") diff --git a/src/cli/ui/custom.rs b/src/cli/ui/custom.rs index ede6e7b..50235e3 100644 --- a/src/cli/ui/custom.rs +++ b/src/cli/ui/custom.rs @@ -10,7 +10,10 @@ use crate::cli::{ AFL_FORKSRV_INIT_TMOUT, }, }; -use anyhow::Context; +use anyhow::{ + bail, + Context, +}; use std::{ process::{ Child, @@ -19,6 +22,11 @@ use std::{ }, sync::mpsc, thread, + thread::sleep, + time::{ + Duration, + Instant, + }, }; #[derive(Clone, Debug)] @@ -73,10 +81,23 @@ impl CustomManager { let child: Child = rx.recv()??; - let mut ratatui = - CustomUI::new(&cloned_config).context("Couldn't create the custom UI ")?; + let mut ratatui = CustomUI::new(&cloned_config); + let start_time = Instant::now(); + + loop { + if start_time.elapsed() > Duration::new(30, 0) { + bail!("Couldn't instantiate the custom UI within 30 seconds..."); + } + if ratatui.is_err() { + println!("Waiting for AFL++ to finish the dry run "); + ratatui = CustomUI::new(&cloned_config); + sleep(Duration::from_millis(100)); + } else { + break; + } + } - ratatui.initialize_tui(child)?; + ratatui.unwrap().initialize_tui(child)?; Ok(()) } } diff --git a/src/cli/ui/ratatui.rs b/src/cli/ui/ratatui.rs index 4ae0982..31add6b 100644 --- a/src/cli/ui/ratatui.rs +++ b/src/cli/ui/ratatui.rs @@ -30,6 +30,7 @@ use ratatui::{ Rect, }, style::{ + palette::tailwind, Color, Modifier, Style, @@ -40,10 +41,13 @@ use ratatui::{ Span, }, widgets::{ + block::Title, Block, Borders, Gauge, + Padding, Paragraph, + Sparkline, }, Frame, }; @@ -59,12 +63,14 @@ use std::{ }, time::Duration, }; +use tailwind::SLATE; #[derive(Clone, Debug)] pub struct CustomUI { ziggy_config: ZiggyConfig, afl_dashboard: AFLDashboard, corpus_watcher: CorpusWatcher, + fuzzing_speed: Vec, } impl CustomUI { @@ -77,6 +83,7 @@ impl CustomUI { .context("Couldn't create AFL dashboard")?, corpus_watcher: CorpusWatcher::from_output(output) .context("Couldn't create the corpus watcher")?, + fuzzing_speed: vec![], }) } @@ -144,10 +151,11 @@ impl CustomUI { f.render_widget(title, area); } - fn render_stats(&self, f: &mut Frame, area: Rect) { + fn render_stats(&mut self, f: &mut Frame, area: Rect) { let data = self.afl_dashboard.read_properties(); - if let Ok(afl) = data { + self.fuzzing_speed.push(afl.exec_speed.into()); + let chunks = Layout::default() .direction(Direction::Horizontal) .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref()) @@ -164,7 +172,7 @@ impl CustomUI { .split(chunks[1]); self.stats_left(f, afl.borrow(), left_chunks[0]); - self.metrics_right(f, afl.borrow(), right_chunk[0]); + self.speed_right(f, right_chunk[0]); } } @@ -174,6 +182,7 @@ impl CustomUI { .margin(1) .constraints([Constraint::Percentage(70), Constraint::Min(1)].as_ref()) .split(area); + let crash_style = Self::if_crash(data); let paragraph = Paragraph::new(Vec::from([ Line::from(vec![ @@ -197,45 +206,6 @@ impl CustomUI { Style::default().add_modifier(Modifier::BOLD), ), ]), - ])) - .block(Block::default().borders(Borders::ALL).title("Statistics")); - - frame.render_widget(paragraph, chunks[0]); - - self.create_stability_display(frame, chunks[1], data); - } - - fn create_stability_display(&self, frame: &mut Frame, area: Rect, data: &AFLProperties) { - let label = format!("{:.2}%", data.stability * 100.0); - let gauge = Gauge::default() - .gauge_style(Style::default().fg(Color::DarkGray).bg(Color::White)) - .use_unicode(true) - .label(label) - .bold() - .ratio(data.stability); - - let block = Block::default() - .title("System Stability") - .borders(Borders::ALL) - .style(Style::default().fg(Color::LightCyan).bg(Color::Black)); - - let paragraph = Paragraph::new(vec![Line::raw("Fuzzing stability")]) - .block(block) - .wrap(ratatui::widgets::Wrap { trim: true }); - - frame.render_widget(paragraph, area); - frame.render_widget(gauge, area); - } - fn metrics_right(&self, frame: &mut Frame, data: &AFLProperties, area: Rect) { - let chunks = Layout::default() - .direction(Direction::Vertical) - .margin(1) - .constraints([Constraint::Percentage(100)].as_ref()) - .split(area); - - let crash_style = Self::if_crash(data); - - let paragraph = Paragraph::new(Vec::from([ Line::from(vec![ Span::raw("Corpus count: "), Span::styled( @@ -252,11 +222,30 @@ impl CustomUI { ), ]), ])) - .block(Block::default().borders(Borders::ALL).title("Metrics")); + .block(Block::default().borders(Borders::ALL).title("Statistics")); frame.render_widget(paragraph, chunks[0]); } + fn speed_right(&self, frame: &mut Frame, area: Rect) { + let chunks = Layout::default() + .direction(Direction::Vertical) + .margin(1) + .constraints([Constraint::Percentage(90)].as_ref()) + .split(area); + + let sparkline = Sparkline::default() + .block( + Block::new() + .borders(Borders::LEFT | Borders::RIGHT) + .title("Execution speed evolution"), + ) + .data(&self.fuzzing_speed) + .style(Style::default().fg(Color::Red)); + + frame.render_widget(sparkline, chunks[0]); + } + fn if_crash(data: &AFLProperties) -> Span { let crash_style = if data.saved_crashes > 0 { Span::styled( @@ -294,8 +283,7 @@ impl CustomUI { if let Some(seeds) = seed_displayer.load() { seed_info_text = seeds .iter() - .enumerate() - .map(|(i, seed)| format!("Seed {}: {}", i + 1, seed)) + .map(|seed| seed.to_string()) .collect::>() .join("\n"); } diff --git a/src/cli/ui/seed.rs b/src/cli/ui/seed.rs index 5cc6c22..1ac98a1 100644 --- a/src/cli/ui/seed.rs +++ b/src/cli/ui/seed.rs @@ -15,24 +15,25 @@ use std::{ io, io::{ BufRead, + Read, Write, }, path::PathBuf, time::{ Duration, Instant, + SystemTime, + UNIX_EPOCH, }, }; -const EVERY_N_SECONDS: u32 = 1; pub const LAST_SEED_FILENAME: &str = "last_seed.phink"; -pub const DELIMITER: &str = "{}\n--------"; +pub const DELIMITER: &str = "-X------X-"; pub struct SeedWriter { input: OneInput, coverage: InputCoverage, responses: Vec, - last_save_time: Instant, } impl SeedWriter { @@ -45,26 +46,29 @@ impl SeedWriter { input, coverage, responses, - last_save_time: Instant::now(), } } - pub fn should_save(&mut self) -> bool { - let now = Instant::now(); - if now.duration_since(self.last_save_time) >= Duration::new(EVERY_N_SECONDS.into(), 0) { - self.last_save_time = now; - true - } else { - false - } + + pub fn should_save() -> bool { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() + % 2 + == 0 } pub fn save(&self, output: PathBuf) { let to = PhinkFiles::new(output).path(PFiles::LastSeed); - - if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(&to) { - for (message, response) in self.input.messages.iter().zip(self.responses.iter()) { - let msg = message.display_with_reply(response.get_response()); - writeln!(file, "{DELIMITER} {msg}").expect("Failed to write to file"); + if let Ok(mut file) = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&to) + { + for message in self.input.messages.iter() { + let msg = message.to_string(); + writeln!(file, "{DELIMITER} {msg}").expect("Failed to save the fuzzed seed"); } } else { eprintln!("Failed to open the file for writing"); @@ -81,7 +85,8 @@ impl SeedDisplayer { Self { output } } pub fn load(&self) -> Option> { - let maybe_file = File::open(PhinkFiles::new(self.output.clone()).path(PFiles::LastSeed)); + let buf = PhinkFiles::new(self.output.clone()).path(PFiles::LastSeed); + let maybe_file = File::open(buf.clone()); if let Ok(file) = maybe_file { return Some(Self::parse(file)) } @@ -89,27 +94,33 @@ impl SeedDisplayer { } fn parse(file: File) -> Vec { - let reader = io::BufReader::new(file); - let mut result = Vec::new(); - let mut current_seed = String::new(); + let mut reader = io::BufReader::new(file); - for line in reader.lines() { - let line = line.unwrap(); - if line == DELIMITER { - result.push(current_seed.clone()); - current_seed.clear(); - } else { - if !current_seed.is_empty() { - current_seed.push('\n'); - } - current_seed.push_str(&line); - } - } + let mut contents = String::new(); - // Push the last seed if the file doesn't end with the delimiter - if !current_seed.is_empty() { - result.push(current_seed); - } - result + reader + .read_to_string(&mut contents) + .expect("Failed to read file"); + + let sections: Vec = contents + .split(DELIMITER) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect(); + + sections + } +} +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_seed_displayer_load() { + let seed_displayer = SeedDisplayer::new(PathBuf::from("tests/fixtures")); + let seeds = seed_displayer.load().unwrap(); + assert_eq!(seeds.len(), 2); + assert_eq!(seeds[0], "crash_with_invariant { data: ' }"); + assert_eq!(seeds[1], "crash_with_invariant { data: }"); } } diff --git a/src/fuzzer/fuzz.rs b/src/fuzzer/fuzz.rs index 559d0de..01932f9 100644 --- a/src/fuzzer/fuzz.rs +++ b/src/fuzzer/fuzz.rs @@ -168,16 +168,16 @@ impl Fuzzer { responses.push(result); } }); - // panic!("xxxxxxxx"); // If the user has `show_ui` turned on, we save the fuzzed seed to display it on the UI if self.ziggy_config.config.show_ui { - let mut seeder = SeedWriter::new( + let seeder = SeedWriter::new( decoded_msgs.to_owned(), coverage.to_owned(), responses.clone(), ); - if seeder.should_save() { + + if SeedWriter::should_save() { seeder.save(self.clone().ziggy_config.fuzz_output()); } } diff --git a/src/fuzzer/parser.rs b/src/fuzzer/parser.rs index 7177cb8..a06c7ab 100644 --- a/src/fuzzer/parser.rs +++ b/src/fuzzer/parser.rs @@ -13,6 +13,10 @@ use crate::{ use contract_transcode::Value; use serde_derive::Serialize; use sp_core::crypto::AccountId32; +use std::fmt::{ + Display, + Formatter, +}; use OriginFuzzingOption::{ DisableOriginFuzzing, EnableOriginFuzzing, @@ -64,6 +68,12 @@ impl Message { } } +impl Display for Message { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(self.message_metadata.to_string().as_str()) + } +} + #[derive(Debug, Clone, Serialize)] pub struct OneInput { pub messages: Vec, diff --git a/tests/fixtures/phink/logs/last_seed.phink b/tests/fixtures/phink/logs/last_seed.phink new file mode 100644 index 0000000..0645a22 --- /dev/null +++ b/tests/fixtures/phink/logs/last_seed.phink @@ -0,0 +1,2 @@ +-X------X- crash_with_invariant { data: ' } +-X------X- crash_with_invariant { data: }