Add client stats to Events (#3116)

* add stats alongside Event over the wire

Author: rmalmain

fuzzers/forkserver/libafl-fuzz/src/hooks.rs

@@ -1,5 +1,5 @@
 use libafl::{
-    events::{Event, EventManagerHook},
+    events::{Event, EventManagerHook, EventWithStats},
     state::Stoppable,
     Error,
 };
@@ -18,9 +18,9 @@ where
         &mut self,
         state: &mut S,
         _client_id: ClientId,
-        event: &Event<I>,
+        event: &EventWithStats<I>,
     ) -> Result<bool, Error> {
-        if self.exit_on_solution && matches!(event, Event::Objective { .. }) {
+        if self.exit_on_solution && matches!(event.event(), Event::Objective { .. }) {
             // TODO: dump state
             state.request_stop();
         }

libafl/src/corpus/minimizer.rs

@@ -6,7 +6,7 @@ use core::{hash::Hash, marker::PhantomData};
 
 use hashbrown::{HashMap, HashSet};
 use libafl_bolts::{
-    AsIter, Named, current_time,
+    AsIter, Named,
     tuples::{Handle, Handled},
 };
 use num_traits::ToPrimitive;
@@ -15,7 +15,7 @@ use z3::{Config, Context, Optimize, ast::Bool};
 use crate::{
     Error, HasMetadata, HasScheduler,
     corpus::Corpus,
-    events::{Event, EventFirer, LogSeverity},
+    events::{Event, EventFirer, EventWithStats, LogSeverity},
     executors::{Executor, HasObservers},
     inputs::Input,
     monitors::stats::{AggregatorOps, UserStats, UserStatsValue},
@@ -123,20 +123,17 @@ where
 
             manager.fire(
                 state,
-                Event::UpdateUserStats {
-                    name: Cow::from("minimisation exec pass"),
-                    value: UserStats::new(UserStatsValue::Ratio(curr, total), AggregatorOps::None),
-                    phantom: PhantomData,
-                },
-            )?;
-
-            manager.fire(
-                state,
-                Event::UpdateExecStats {
-                    time: current_time(),
-                    phantom: PhantomData,
+                EventWithStats::with_current_time(
+                    Event::UpdateUserStats {
+                        name: Cow::from("minimisation exec pass"),
+                        value: UserStats::new(
+                            UserStatsValue::Ratio(curr, total),
+                            AggregatorOps::None,
+                        ),
+                        phantom: PhantomData,
+                    },
                     executions,
-                },
+                ),
             )?;
 
             let seed_expr = Bool::fresh_const(&ctx, "seed");

libafl/src/events/broker_hooks/centralized.rs

@@ -11,7 +11,7 @@ use serde::de::DeserializeOwned;
 
 #[cfg(feature = "llmp_compression")]
 use crate::events::COMPRESS_THRESHOLD;
-use crate::events::{_LLMP_TAG_TO_MAIN, BrokerEventResult, Event};
+use crate::events::{_LLMP_TAG_TO_MAIN, BrokerEventResult, Event, EventWithStats};
 
 /// An LLMP-backed event manager for scalable multi-processed fuzzing
 pub struct CentralizedLlmpHook<I> {
@@ -47,7 +47,7 @@ where
             } else {
                 &*msg
             };
-            let event: Event<I> = postcard::from_bytes(event_bytes)?;
+            let event: EventWithStats<I> = postcard::from_bytes(event_bytes)?;
             match Self::handle_in_broker(client_id, &event)? {
                 BrokerEventResult::Forward => Ok(LlmpMsgHookResult::ForwardToClients),
                 BrokerEventResult::Handled => Ok(LlmpMsgHookResult::Handled),
@@ -85,9 +85,9 @@ impl<I> CentralizedLlmpHook<I> {
     #[expect(clippy::unnecessary_wraps)]
     fn handle_in_broker(
         _client_id: ClientId,
-        event: &Event<I>,
+        event: &EventWithStats<I>,
     ) -> Result<BrokerEventResult, Error> {
-        match &event {
+        match event.event() {
             Event::NewTestcase { .. } | Event::Stop => Ok(BrokerEventResult::Forward),
             _ => Ok(BrokerEventResult::Handled),
         }

libafl/src/events/broker_hooks/centralized_multi_machine.rs

@@ -22,7 +22,7 @@ use tokio::{
 
 use crate::{
     events::{
-        Event,
+        EventWithStats,
         centralized::_LLMP_TAG_TO_MAIN,
         multi_machine::{MultiMachineMsg, TcpMultiMachineState},
     },
@@ -128,7 +128,7 @@ where
     #[cfg(feature = "llmp_compression")]
     fn try_compress(
         state_lock: &mut RwLockWriteGuard<TcpMultiMachineState<A>>,
-        event: &Event<I>,
+        event: &EventWithStats<I>,
     ) -> Result<(Flags, Vec<u8>), Error> {
         let serialized = postcard::to_allocvec(&event)?;
 
@@ -141,7 +141,7 @@ where
     #[cfg(not(feature = "llmp_compression"))]
     fn try_compress(
         _state_lock: &mut RwLockWriteGuard<TcpMultiMachineState<A>>,
-        event: &Event<I>,
+        event: &EventWithStats<I>,
     ) -> Result<(Flags, Vec<u8>), Error> {
         Ok((Flags(0), postcard::to_allocvec(&event)?))
     }

libafl/src/events/broker_hooks/mod.rs

@@ -30,6 +30,8 @@ pub mod centralized_multi_machine;
 #[cfg(all(unix, feature = "multi_machine"))]
 pub use centralized_multi_machine::*;
 
+use super::EventWithStats;
+
 /// An LLMP-backed event hook for scalable multi-processed fuzzing
 #[derive(Debug)]
 pub struct StdLlmpEventHook<I, MT> {
@@ -71,7 +73,7 @@ where
             } else {
                 &*msg
             };
-            let event: Event<I> = postcard::from_bytes(event_bytes)?;
+            let event: EventWithStats<I> = postcard::from_bytes(event_bytes)?;
             match Self::handle_in_broker(
                 monitor,
                 &mut self.client_stats_manager,
@@ -117,8 +119,16 @@ where
         monitor: &mut MT,
         client_stats_manager: &mut ClientStatsManager,
         client_id: ClientId,
-        event: &Event<I>,
+        event: &EventWithStats<I>,
     ) -> Result<BrokerEventResult, Error> {
+        let stats = event.stats();
+
+        client_stats_manager.client_stats_insert(ClientId(0));
+        client_stats_manager.update_client_stats_for(ClientId(0), |client_stat| {
+            client_stat.update_executions(stats.executions, stats.time);
+        });
+
+        let event = event.event();
         match &event {
             Event::NewTestcase {
                 corpus_size,
@@ -138,19 +148,7 @@ where
                 monitor.display(client_stats_manager, event.name(), id);
                 Ok(BrokerEventResult::Forward)
             }
-            Event::UpdateExecStats {
-                time,
-                executions,
-                phantom: _,
-            } => {
-                // TODO: The monitor buffer should be added on client add.
-                client_stats_manager.client_stats_insert(client_id);
-                client_stats_manager.update_client_stats_for(client_id, |client_stat| {
-                    client_stat.update_executions(*executions, *time);
-                });
-                monitor.display(client_stats_manager, event.name(), client_id);
-                Ok(BrokerEventResult::Handled)
-            }
+            Event::Heartbeat => Ok(BrokerEventResult::Handled),
             Event::UpdateUserStats { name, value, .. } => {
                 client_stats_manager.client_stats_insert(client_id);
                 client_stats_manager.update_client_stats_for(client_id, |client_stat| {
@@ -162,8 +160,6 @@ where
             }
             #[cfg(feature = "introspection")]
             Event::UpdatePerfMonitor {
-                time,
-                executions,
                 introspection_stats,
                 phantom: _,
             } => {
@@ -172,8 +168,6 @@ where
                 // Get the client for the staterestorer ID
                 client_stats_manager.client_stats_insert(client_id);
                 client_stats_manager.update_client_stats_for(client_id, |client_stat| {
-                    // Update the normal monitor for this client
-                    client_stat.update_executions(*executions, *time);
                     // Update the performance monitor for this client
                     client_stat.update_introspection_stats((**introspection_stats).clone());
                 });

libafl/src/events/centralized.rs

@@ -22,7 +22,7 @@ use libafl_bolts::{
     llmp::{LLMP_FLAG_COMPRESSED, LLMP_FLAG_INITIALIZED},
 };
 
-use super::AwaitRestartSafe;
+use super::{AwaitRestartSafe, EventWithStats};
 #[cfg(feature = "llmp_compression")]
 use crate::events::llmp::COMPRESS_THRESHOLD;
 use crate::{
@@ -166,18 +166,18 @@ where
     }
 
     #[expect(clippy::match_same_arms)]
-    fn fire(&mut self, state: &mut S, mut event: Event<I>) -> Result<(), Error> {
+    fn fire(&mut self, state: &mut S, mut event: EventWithStats<I>) -> Result<(), Error> {
         if !self.is_main {
             // secondary node
             let mut is_tc = false;
             // Forward to main only if new tc, heartbeat, or optionally, a new objective
-            let should_be_forwarded = match &mut event {
+            let should_be_forwarded = match event.event_mut() {
                 Event::NewTestcase { forward_id, .. } => {
                     *forward_id = Some(ClientId(self.inner.mgr_id().0 as u32));
                     is_tc = true;
                     true
                 }
-                Event::UpdateExecStats { .. } => true, // send UpdateExecStats but this guy won't be handled. the only purpose is to keep this client alive else the broker thinks it is dead and will dc it
+                Event::Heartbeat => true, // the only purpose is to keep this client alive else the broker thinks it is dead and will dc it
                 Event::Objective { .. } => true,
                 Event::Stop => true,
                 _ => false,
@@ -201,7 +201,10 @@ where
         state: &mut S,
         severity_level: LogSeverity,
         message: String,
-    ) -> Result<(), Error> {
+    ) -> Result<(), Error>
+    where
+        S: HasExecutions,
+    {
         self.inner.log(state, severity_level, message)
     }
 
@@ -261,7 +264,7 @@ where
     SHM: ShMem,
     SP: ShMemProvider<ShMem = SHM>,
 {
-    fn try_receive(&mut self, state: &mut S) -> Result<Option<(Event<I>, bool)>, Error> {
+    fn try_receive(&mut self, state: &mut S) -> Result<Option<(EventWithStats<I>, bool)>, Error> {
         if self.is_main {
             // main node
             self.receive_from_secondary(state)
@@ -272,7 +275,7 @@ where
         }
     }
 
-    fn on_interesting(&mut self, state: &mut S, event: Event<I>) -> Result<(), Error> {
+    fn on_interesting(&mut self, state: &mut S, event: EventWithStats<I>) -> Result<(), Error> {
         self.inner.fire(state, event)
     }
 }
@@ -343,7 +346,7 @@ where
     SP: ShMemProvider<ShMem = SHM>,
 {
     #[cfg(feature = "llmp_compression")]
-    fn forward_to_main(&mut self, event: &Event<I>) -> Result<(), Error> {
+    fn forward_to_main(&mut self, event: &EventWithStats<I>) -> Result<(), Error> {
         let serialized = postcard::to_allocvec(event)?;
         let flags = LLMP_FLAG_INITIALIZED;
 
@@ -363,13 +366,16 @@ where
     }
 
     #[cfg(not(feature = "llmp_compression"))]
-    fn forward_to_main(&mut self, event: &Event<I>) -> Result<(), Error> {
+    fn forward_to_main(&mut self, event: &EventWithStats<I>) -> Result<(), Error> {
         let serialized = postcard::to_allocvec(event)?;
         self.client.send_buf(_LLMP_TAG_TO_MAIN, &serialized)?;
         Ok(())
     }
 
-    fn receive_from_secondary(&mut self, state: &mut S) -> Result<Option<(Event<I>, bool)>, Error> {
+    fn receive_from_secondary(
+        &mut self,
+        state: &mut S,
+    ) -> Result<Option<(EventWithStats<I>, bool)>, Error> {
         // TODO: Get around local event copy by moving handle_in_client
         let self_id = self.client.sender().id();
         while let Some((client_id, tag, _flags, msg)) = self.client.recv_buf_with_flags()? {
@@ -392,15 +398,18 @@ where
             } else {
                 msg
             };
-            let event: Event<I> = postcard::from_bytes(event_bytes)?;
-            log::debug!("Processor received message {}", event.name_detailed());
+            let event: EventWithStats<I> = postcard::from_bytes(event_bytes)?;
+            log::debug!(
+                "Processor received message {}",
+                event.event().name_detailed()
+            );
 
-            let event_name = event.name_detailed();
+            let event_name = event.event().name_detailed();
 
-            match event {
+            match event.event() {
                 Event::NewTestcase {
                     client_config,
-                    ref observers_buf,
+                    observers_buf,
                     forward_id,
                     ..
                 } => {
@@ -425,7 +434,7 @@ where
                 _ => {
                     return Err(Error::illegal_state(format!(
                         "Received illegal message that message should not have arrived: {:?}.",
-                        event.name()
+                        event.event().name()
                     )));
                 }
             }

libafl/src/events/events_hooks/mod.rs

@@ -4,7 +4,7 @@
 //! other clients
 use libafl_bolts::ClientId;
 
-use crate::{Error, events::Event};
+use crate::{Error, events::EventWithStats};
 
 /// The `broker_hooks` that are run before and after the event manager calls `try_receive`
 pub trait EventManagerHook<I, S> {
@@ -14,7 +14,7 @@ pub trait EventManagerHook<I, S> {
         &mut self,
         state: &mut S,
         client_id: ClientId,
-        event: &Event<I>,
+        event: &EventWithStats<I>,
     ) -> Result<bool, Error>;
 }
 
@@ -25,7 +25,7 @@ pub trait EventManagerHooksTuple<I, S> {
         &mut self,
         state: &mut S,
         client_id: ClientId,
-        event: &Event<I>,
+        event: &EventWithStats<I>,
     ) -> Result<bool, Error>;
 }
 
@@ -35,7 +35,7 @@ impl<I, S> EventManagerHooksTuple<I, S> for () {
         &mut self,
         _state: &mut S,
         _client_id: ClientId,
-        _event: &Event<I>,
+        _event: &EventWithStats<I>,
     ) -> Result<bool, Error> {
         Ok(true)
     }
@@ -51,7 +51,7 @@ where
         &mut self,
         state: &mut S,
         client_id: ClientId,
-        event: &Event<I>,
+        event: &EventWithStats<I>,
     ) -> Result<bool, Error> {
         let first = self.0.pre_receive(state, client_id, event)?;
         let second = self.1.pre_receive_all(state, client_id, event)?;