Skip to content

Commit 373fe03

Browse files
tokatokatokatoka
authored
Use HashMap to manage ClientStat, Fix #3133 (#3142)
* fix * lol * lol * clp fixer * clp fixer * revert cargo.toml
1 parent fa8a576 commit 373fe03

File tree

20 files changed

+234
-195
lines changed

20 files changed

+234
-195
lines changed

‎libafl/Cargo.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ default = [
3939
"regex",
4040
"serdeany_autoreg",
4141
"libafl_bolts/xxh3",
42-
"tui_monitor",
4342
]
4443
document-features = ["dep:document-features"]
4544

‎libafl/examples/tui_mock/main.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,6 @@ pub fn main() {
1717
let _client_stats = ClientStats::default();
1818
let mut client_stats_manager = ClientStatsManager::default();
1919

20-
monitor.display(&mut client_stats_manager, "Test", ClientId(0));
20+
let _ = monitor.display(&mut client_stats_manager, "Test", ClientId(0));
2121
sleep(Duration::from_secs(10));
2222
}

‎libafl/src/corpus/minimizer.rs

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ where
190190
// Perform the optimization!
191191
opt.check(&[]);
192192

193-
let res = if let Some(model) = opt.get_model() {
193+
if let Some(model) = opt.get_model() {
194194
let mut removed = Vec::with_capacity(state.corpus().count());
195195
for (seed, (id, _)) in seed_exprs {
196196
// if the model says the seed isn't there, mark it for deletion
@@ -214,11 +214,8 @@ where
214214
}
215215

216216
*state.corpus_mut().current_mut() = None; //we may have removed the current ID from the corpus
217-
Ok(())
218-
} else {
219-
Err(Error::unknown("Corpus minimization failed; unsat."))
220-
};
221-
222-
res
217+
return Ok(());
218+
}
219+
Err(Error::unknown("Corpus minimization failed; unsat."))
223220
}
224221
}

‎libafl/src/events/broker_hooks/mod.rs

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ where
9393
&mut self.client_stats_manager,
9494
"Broker Heartbeat",
9595
ClientId(0),
96-
);
96+
)?;
9797
Ok(())
9898
}
9999
}
@@ -114,7 +114,6 @@ where
114114
}
115115

116116
/// Handle arriving events in the broker
117-
#[expect(clippy::unnecessary_wraps)]
118117
fn handle_in_broker(
119118
monitor: &mut MT,
120119
client_stats_manager: &mut ClientStatsManager,
@@ -123,10 +122,10 @@ where
123122
) -> Result<BrokerEventResult, Error> {
124123
let stats = event.stats();
125124

126-
client_stats_manager.client_stats_insert(ClientId(0));
127-
client_stats_manager.update_client_stats_for(ClientId(0), |client_stat| {
125+
client_stats_manager.client_stats_insert(client_id)?;
126+
client_stats_manager.update_client_stats_for(client_id, |client_stat| {
128127
client_stat.update_executions(stats.executions, stats.time);
129-
});
128+
})?;
130129

131130
let event = event.event();
132131
match &event {
@@ -141,21 +140,21 @@ where
141140
client_id
142141
};
143142

144-
client_stats_manager.client_stats_insert(id);
143+
client_stats_manager.client_stats_insert(id)?;
145144
client_stats_manager.update_client_stats_for(id, |client_stat| {
146145
client_stat.update_corpus_size(*corpus_size as u64);
147-
});
148-
monitor.display(client_stats_manager, event.name(), id);
146+
})?;
147+
monitor.display(client_stats_manager, event.name(), id)?;
149148
Ok(BrokerEventResult::Forward)
150149
}
151150
Event::Heartbeat => Ok(BrokerEventResult::Handled),
152151
Event::UpdateUserStats { name, value, .. } => {
153-
client_stats_manager.client_stats_insert(client_id);
152+
client_stats_manager.client_stats_insert(client_id)?;
154153
client_stats_manager.update_client_stats_for(client_id, |client_stat| {
155154
client_stat.update_user_stats(name.clone(), value.clone());
156-
});
155+
})?;
157156
client_stats_manager.aggregate(name);
158-
monitor.display(client_stats_manager, event.name(), client_id);
157+
monitor.display(client_stats_manager, event.name(), client_id)?;
159158
Ok(BrokerEventResult::Handled)
160159
}
161160
#[cfg(feature = "introspection")]
@@ -166,24 +165,24 @@ where
166165
// TODO: The monitor buffer should be added on client add.
167166

168167
// Get the client for the staterestorer ID
169-
client_stats_manager.client_stats_insert(client_id);
168+
client_stats_manager.client_stats_insert(client_id)?;
170169
client_stats_manager.update_client_stats_for(client_id, |client_stat| {
171170
// Update the performance monitor for this client
172171
client_stat.update_introspection_stats((**introspection_stats).clone());
173-
});
172+
})?;
174173

175174
// Display the monitor via `.display` only on core #1
176-
monitor.display(client_stats_manager, event.name(), client_id);
175+
monitor.display(client_stats_manager, event.name(), client_id)?;
177176

178177
// Correctly handled the event
179178
Ok(BrokerEventResult::Handled)
180179
}
181180
Event::Objective { objective_size, .. } => {
182-
client_stats_manager.client_stats_insert(client_id);
181+
client_stats_manager.client_stats_insert(client_id)?;
183182
client_stats_manager.update_client_stats_for(client_id, |client_stat| {
184183
client_stat.update_objective_size(*objective_size as u64);
185-
});
186-
monitor.display(client_stats_manager, event.name(), client_id);
184+
})?;
185+
monitor.display(client_stats_manager, event.name(), client_id)?;
187186
Ok(BrokerEventResult::Handled)
188187
}
189188
Event::Log {

‎libafl/src/events/multi_machine.rs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -268,10 +268,10 @@ where
268268
let timeout = current_time() + parent_lock.node_descriptor.timeout;
269269

270270
parent_lock.parent = loop {
271-
log::debug!("Trying to connect to parent @ {}..", parent_addr);
271+
log::debug!("Trying to connect to parent @ {parent_addr}..");
272272
match TcpStream::connect(parent_addr).await {
273273
Ok(stream) => {
274-
log::debug!("Connected to parent @ {}", parent_addr);
274+
log::debug!("Connected to parent @ {parent_addr}");
275275

276276
break Some(stream);
277277
}
@@ -302,10 +302,10 @@ where
302302

303303
// The main listening loop. Should never fail.
304304
'listening: loop {
305-
log::debug!("listening for children on {:?}...", listener);
305+
log::debug!("listening for children on {listener:?}...");
306306
match listener.accept().await {
307307
Ok((mut stream, addr)) => {
308-
log::debug!("{} joined the children.", addr);
308+
log::debug!("{addr} joined the children.");
309309
let mut state_guard = state.write().await;
310310

311311
if let Err(e) = state_guard
@@ -487,7 +487,7 @@ where
487487

488488
// Garbage collect disconnected children
489489
for id_to_remove in &ids_to_remove {
490-
log::debug!("Child {:?} has been garbage collected.", id_to_remove);
490+
log::debug!("Child {id_to_remove:?} has been garbage collected.");
491491
self.children.remove(id_to_remove);
492492
}
493493
}
@@ -596,7 +596,7 @@ where
596596

597597
// Garbage collect disconnected children
598598
for id_to_remove in &ids_to_remove {
599-
log::debug!("Child {:?} has been garbage collected.", id_to_remove);
599+
log::debug!("Child {id_to_remove:?} has been garbage collected.");
600600
self.children.remove(id_to_remove);
601601
}
602602

‎libafl/src/events/simple.rs

Lines changed: 42 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ use alloc::vec::Vec;
55
use core::sync::atomic::{Ordering, compiler_fence};
66
use core::{fmt::Debug, marker::PhantomData, time::Duration};
77

8+
#[cfg(feature = "std")]
9+
use hashbrown::HashMap;
810
use libafl_bolts::ClientId;
911
#[cfg(all(feature = "std", any(windows, not(feature = "fork"))))]
1012
use libafl_bolts::os::startable_self;
@@ -200,40 +202,39 @@ where
200202
}
201203

202204
/// Handle arriving events in the broker
203-
#[expect(clippy::unnecessary_wraps)]
204205
fn handle_in_broker(
205206
monitor: &mut MT,
206207
client_stats_manager: &mut ClientStatsManager,
207208
event: &EventWithStats<I>,
208209
) -> Result<BrokerEventResult, Error> {
209210
let stats = event.stats();
210211

211-
client_stats_manager.client_stats_insert(ClientId(0));
212+
client_stats_manager.client_stats_insert(ClientId(0))?;
212213
client_stats_manager.update_client_stats_for(ClientId(0), |client_stat| {
213214
client_stat.update_executions(stats.executions, stats.time);
214-
});
215+
})?;
215216

216217
let event = event.event();
217218
match event {
218219
Event::NewTestcase { corpus_size, .. } => {
219-
client_stats_manager.client_stats_insert(ClientId(0));
220+
client_stats_manager.client_stats_insert(ClientId(0))?;
220221
client_stats_manager.update_client_stats_for(ClientId(0), |client_stat| {
221222
client_stat.update_corpus_size(*corpus_size as u64);
222-
});
223-
monitor.display(client_stats_manager, event.name(), ClientId(0));
223+
})?;
224+
monitor.display(client_stats_manager, event.name(), ClientId(0))?;
224225
Ok(BrokerEventResult::Handled)
225226
}
226227
Event::Heartbeat => {
227-
monitor.display(client_stats_manager, event.name(), ClientId(0));
228+
monitor.display(client_stats_manager, event.name(), ClientId(0))?;
228229
Ok(BrokerEventResult::Handled)
229230
}
230231
Event::UpdateUserStats { name, value, .. } => {
231-
client_stats_manager.client_stats_insert(ClientId(0));
232+
client_stats_manager.client_stats_insert(ClientId(0))?;
232233
client_stats_manager.update_client_stats_for(ClientId(0), |client_stat| {
233234
client_stat.update_user_stats(name.clone(), value.clone());
234-
});
235+
})?;
235236
client_stats_manager.aggregate(name);
236-
monitor.display(client_stats_manager, event.name(), ClientId(0));
237+
monitor.display(client_stats_manager, event.name(), ClientId(0))?;
237238
Ok(BrokerEventResult::Handled)
238239
}
239240
#[cfg(feature = "introspection")]
@@ -244,16 +245,16 @@ where
244245
// TODO: The monitor buffer should be added on client add.
245246
client_stats_manager.update_client_stats_for(ClientId(0), |client_stat| {
246247
client_stat.update_introspection_stats((**introspection_stats).clone());
247-
});
248-
monitor.display(client_stats_manager, event.name(), ClientId(0));
248+
})?;
249+
monitor.display(client_stats_manager, event.name(), ClientId(0))?;
249250
Ok(BrokerEventResult::Handled)
250251
}
251252
Event::Objective { objective_size, .. } => {
252-
client_stats_manager.client_stats_insert(ClientId(0));
253+
client_stats_manager.client_stats_insert(ClientId(0))?;
253254
client_stats_manager.update_client_stats_for(ClientId(0), |client_stat| {
254255
client_stat.update_objective_size(*objective_size as u64);
255-
});
256-
monitor.display(client_stats_manager, event.name(), ClientId(0));
256+
})?;
257+
monitor.display(client_stats_manager, event.name(), ClientId(0))?;
257258
Ok(BrokerEventResult::Handled)
258259
}
259260
Event::Log {
@@ -525,31 +526,32 @@ where
525526
}
526527

527528
// If we're restarting, deserialize the old state.
528-
let (state, mgr) = match staterestorer.restore::<(S, Duration, Vec<ClientStats>)>()? {
529-
None => {
530-
log::info!("First run. Let's set it all up");
531-
// Mgr to send and receive msgs from/to all other fuzzer instances
532-
(
533-
None,
534-
SimpleRestartingEventManager::launched(monitor, staterestorer),
535-
)
536-
}
537-
// Restoring from a previous run, deserialize state and corpus.
538-
Some((state, start_time, clients_stats)) => {
539-
log::info!("Subsequent run. Loaded previous state.");
540-
// We reset the staterestorer, the next staterestorer and receiver (after crash) will reuse the page from the initial message.
541-
staterestorer.reset();
542-
543-
// reload the state of the monitor to display the correct stats after restarts
544-
let mut this = SimpleRestartingEventManager::launched(monitor, staterestorer);
545-
this.inner.client_stats_manager.set_start_time(start_time);
546-
this.inner
547-
.client_stats_manager
548-
.update_all_client_stats(clients_stats);
549-
550-
(Some(state), this)
551-
}
552-
};
529+
let (state, mgr) =
530+
match staterestorer.restore::<(S, Duration, HashMap<ClientId, ClientStats>)>()? {
531+
None => {
532+
log::info!("First run. Let's set it all up");
533+
// Mgr to send and receive msgs from/to all other fuzzer instances
534+
(
535+
None,
536+
SimpleRestartingEventManager::launched(monitor, staterestorer),
537+
)
538+
}
539+
// Restoring from a previous run, deserialize state and corpus.
540+
Some((state, start_time, clients_stats)) => {
541+
log::info!("Subsequent run. Loaded previous state.");
542+
// We reset the staterestorer, the next staterestorer and receiver (after crash) will reuse the page from the initial message.
543+
staterestorer.reset();
544+
545+
// reload the state of the monitor to display the correct stats after restarts
546+
let mut this = SimpleRestartingEventManager::launched(monitor, staterestorer);
547+
this.inner.client_stats_manager.set_start_time(start_time);
548+
this.inner
549+
.client_stats_manager
550+
.update_all_client_stats(clients_stats);
551+
552+
(Some(state), this)
553+
}
554+
};
553555

554556
/* TODO: Not sure if this is needed
555557
// We commit an empty NO_RESTART message to this buf, against infinite loops,

0 commit comments

Comments
 (0)