diff --git a/libafl/src/monitors/logics.rs b/libafl/src/monitors/logics.rs new file mode 100644 index 0000000000..fefd254a39 --- /dev/null +++ b/libafl/src/monitors/logics.rs @@ -0,0 +1,179 @@ +//! Monitor wrappers that add logics to monitor + +use libafl_bolts::ClientId; + +use crate::monitors::{stats::ClientStatsManager, Monitor}; + +/// The wrapped monitor will keep displaying until the closure evaluates to false +#[derive(Debug)] +pub struct WhileMonitor { + closure: CB, + monitor: M, +} + +impl Monitor for WhileMonitor +where + CB: FnMut(&mut ClientStatsManager, &str, ClientId) -> bool, + M: Monitor, +{ + fn display( + &mut self, + client_stats_manager: &mut ClientStatsManager, + event_msg: &str, + sender_id: ClientId, + ) { + while (self.closure)(client_stats_manager, event_msg, sender_id) { + self.monitor + .display(client_stats_manager, event_msg, sender_id); + } + } +} + +impl WhileMonitor +where + CB: FnMut(&mut ClientStatsManager, &str, ClientId) -> bool, + M: Monitor, +{ + /// Create a new [`WhileMonitor`]. + /// + /// The `closure` will be evaluated at each `display` call + #[must_use] + pub fn new(closure: CB, monitor: M) -> Self { + Self { closure, monitor } + } +} + +/// The wrapped monitor will display if the closure evaluates to true +#[derive(Debug)] +pub struct IfMonitor { + closure: CB, + monitor: M, +} + +impl Monitor for IfMonitor +where + CB: FnMut(&mut ClientStatsManager, &str, ClientId) -> bool, + M: Monitor, +{ + fn display( + &mut self, + client_stats_manager: &mut ClientStatsManager, + event_msg: &str, + sender_id: ClientId, + ) { + if (self.closure)(client_stats_manager, event_msg, sender_id) { + self.monitor + .display(client_stats_manager, event_msg, sender_id); + } + } +} + +impl IfMonitor +where + CB: FnMut(&mut ClientStatsManager, &str, ClientId) -> bool, + M: Monitor, +{ + /// Create a new [`IfMonitor`]. + /// + /// The `closure` will be evaluated at each `display` call + #[must_use] + pub fn new(closure: CB, monitor: M) -> Self { + Self { closure, monitor } + } +} + +/// Either wrapped monitors will display based on the closure evaluation +#[derive(Debug)] +pub struct IfElseMonitor { + closure: CB, + if_monitor: M1, + else_monitor: M2, +} + +impl Monitor for IfElseMonitor +where + CB: FnMut(&mut ClientStatsManager, &str, ClientId) -> bool, + M1: Monitor, + M2: Monitor, +{ + fn display( + &mut self, + client_stats_manager: &mut ClientStatsManager, + event_msg: &str, + sender_id: ClientId, + ) { + if (self.closure)(client_stats_manager, event_msg, sender_id) { + self.if_monitor + .display(client_stats_manager, event_msg, sender_id); + } else { + self.else_monitor + .display(client_stats_manager, event_msg, sender_id); + } + } +} + +impl IfElseMonitor +where + CB: FnMut(&mut ClientStatsManager, &str, ClientId) -> bool, + M1: Monitor, + M2: Monitor, +{ + /// Create a new [`IfElseMonitor`]. + /// + /// The `closure` will be evaluated at each `display` call + #[must_use] + pub fn new(closure: CB, if_monitor: M1, else_monitor: M2) -> Self { + Self { + closure, + if_monitor, + else_monitor, + } + } +} + +/// A monitor wrapper where the monitor does not need to be initialized, but can be [`None`]. +#[derive(Debug)] +pub struct OptionalMonitor { + monitor: Option, +} + +impl Monitor for OptionalMonitor +where + M: Monitor, +{ + fn display( + &mut self, + client_stats_manager: &mut ClientStatsManager, + event_msg: &str, + sender_id: ClientId, + ) { + if let Some(monitor) = self.monitor.as_mut() { + monitor.display(client_stats_manager, event_msg, sender_id); + } + } +} + +impl OptionalMonitor +where + M: Monitor, +{ + /// Create a new [`OptionalMonitor`] + #[must_use] + pub fn new(monitor: Option) -> Self { + Self { monitor } + } + + /// Create a new [`OptionalMonitor`] with a monitor + #[must_use] + pub fn some(monitor: M) -> Self { + Self { + monitor: Some(monitor), + } + } + + /// Create a new [`OptionalMonitor`] without a monitor + #[must_use] + pub fn none() -> Self { + Self { monitor: None } + } +} diff --git a/libafl/src/monitors/mod.rs b/libafl/src/monitors/mod.rs index 51055c69ea..226920432b 100644 --- a/libafl/src/monitors/mod.rs +++ b/libafl/src/monitors/mod.rs @@ -5,6 +5,9 @@ pub use multi::MultiMonitor; pub mod stats; +pub mod logics; +pub use logics::{IfElseMonitor, IfMonitor, OptionalMonitor, WhileMonitor}; + #[cfg(feature = "std")] pub mod disk; #[cfg(feature = "std")]