libafl/stages/
mod.rs

1/*!
2A [`Stage`] is a technique used during fuzzing, working on one [`crate::corpus::Corpus`] entry, and potentially altering it or creating new entries.
3A well-known [`Stage`], for example, is the mutational stage, running multiple [`crate::mutators::Mutator`]s against a [`crate::corpus::Testcase`], potentially storing new ones, according to [`crate::feedbacks::Feedback`].
4Other stages may enrich [`crate::corpus::Testcase`]s with metadata.
5*/
6
7use alloc::{
8    borrow::{Cow, ToOwned},
9    boxed::Box,
10    string::ToString,
11    vec::Vec,
12};
13use core::{fmt, marker::PhantomData};
14
15#[cfg(feature = "std")]
16pub use afl_stats::{AflStatsStage, CalibrationTime, FuzzTime, SyncTime};
17pub use calibrate::CalibrationStage;
18pub use colorization::*;
19#[cfg(all(feature = "std", unix))]
20pub use concolic::ConcolicTracingStage;
21#[cfg(all(feature = "std", feature = "concolic_mutation", unix))]
22pub use concolic::SimpleConcolicMutationalStage;
23#[cfg(feature = "std")]
24pub use dump::*;
25pub use generalization::GeneralizationStage;
26use hashbrown::HashSet;
27use libafl_bolts::{
28    Named, impl_serdeany,
29    tuples::{HasConstLen, IntoVec},
30};
31pub use logics::*;
32pub use mutational::{MutationalStage, StdMutationalStage};
33pub use power::{PowerMutationalStage, StdPowerMutationalStage};
34use serde::{Deserialize, Serialize};
35#[cfg(feature = "std")]
36pub use sync::*;
37#[cfg(feature = "std")]
38pub use time_tracker::TimeTrackingStageWrapper;
39pub use tmin::{ObserverEqualityFactory, ObserverEqualityFeedback, StdTMinMutationalStage};
40pub use tracing::TracingStage;
41pub use tuneable::*;
42use tuple_list::NonEmptyTuple;
43#[cfg(feature = "unicode")]
44pub use unicode::*;
45#[cfg(feature = "std")]
46pub use verify_timeouts::{TimeoutsToVerify, VerifyTimeoutsStage};
47
48use crate::{
49    Error, HasNamedMetadata,
50    corpus::{CorpusId, HasCurrentCorpusId},
51    events::SendExiting,
52    state::{HasCurrentStageId, HasExecutions, MaybeHasClientPerfMonitor, Stoppable},
53};
54
55/// Mutational stage is the normal fuzzing stage.
56pub mod mutational;
57pub mod push;
58pub mod tmin;
59
60pub mod shadow;
61pub use shadow::*;
62
63pub mod replay;
64pub use replay::*;
65
66#[cfg(feature = "std")]
67pub mod afl_stats;
68pub mod calibrate;
69pub mod colorization;
70#[cfg(all(feature = "std", unix))]
71pub mod concolic;
72#[cfg(feature = "std")]
73pub mod dump;
74pub mod dynamic;
75pub mod generalization;
76pub mod generation;
77pub mod logics;
78pub mod nop;
79pub mod power;
80#[cfg(feature = "std")]
81pub mod sync;
82#[cfg(feature = "std")]
83pub mod time_tracker;
84pub mod tracing;
85pub mod tuneable;
86#[cfg(feature = "unicode")]
87pub mod unicode;
88#[cfg(feature = "std")]
89pub mod verify_timeouts;
90
91/// A stage is one step in the fuzzing process.
92/// Multiple stages will be scheduled one by one for each input.
93pub trait Stage<E, EM, S, Z> {
94    /// Run the stage.
95    ///
96    /// If you want this stage to restart, then
97    /// Before a call to perform, [`Restartable::should_restart`] will be (must be!) called.
98    /// After returning (so non-target crash or timeout in a restarting case), [`Restartable::clear_progress`] gets called.
99    fn perform(
100        &mut self,
101        fuzzer: &mut Z,
102        executor: &mut E,
103        state: &mut S,
104        manager: &mut EM,
105    ) -> Result<(), Error>;
106}
107
108/// Restartable trait takes care of stage restart.
109pub trait Restartable<S> {
110    /// This method will be called before every call to [`Stage::perform`].
111    /// Initialize the restart tracking for this stage, _if it is not yet initialized_.
112    /// On restart, this will be called again.
113    /// As long as [`Restartable::clear_progress`], all subsequent calls happen on restart.
114    /// Returns `true`, if the stage's [`Stage::perform`] method should run, else `false`.
115    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error>;
116
117    /// Clear the current status tracking of the associated stage
118    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error>;
119}
120
121/// A tuple holding all `Stages` used for fuzzing.
122pub trait StagesTuple<E, EM, S, Z> {
123    /// Performs all `Stages` in this tuple.
124    fn perform_all(
125        &mut self,
126        fuzzer: &mut Z,
127        executor: &mut E,
128        state: &mut S,
129        manager: &mut EM,
130    ) -> Result<(), Error>;
131}
132
133impl<E, EM, S, Z> StagesTuple<E, EM, S, Z> for ()
134where
135    S: HasCurrentStageId,
136{
137    fn perform_all(
138        &mut self,
139        _: &mut Z,
140        _: &mut E,
141        stage: &mut S,
142        _: &mut EM,
143    ) -> Result<(), Error> {
144        if stage.current_stage_id()?.is_some() {
145            Err(Error::illegal_state(
146                "Got to the end of the tuple without completing resume.",
147            ))
148        } else {
149            Ok(())
150        }
151    }
152}
153
154impl<Head, Tail, E, EM, S, Z> StagesTuple<E, EM, S, Z> for (Head, Tail)
155where
156    Head: Stage<E, EM, S, Z> + Restartable<S>,
157    Tail: StagesTuple<E, EM, S, Z> + HasConstLen,
158    S: HasCurrentStageId + Stoppable + MaybeHasClientPerfMonitor,
159    EM: SendExiting,
160{
161    /// Performs all stages in the tuple,
162    /// Checks after every stage if state wants to stop
163    /// and returns an [`Error::ShuttingDown`] if so
164    fn perform_all(
165        &mut self,
166        fuzzer: &mut Z,
167        executor: &mut E,
168        state: &mut S,
169        manager: &mut EM,
170    ) -> Result<(), Error> {
171        match state.current_stage_id()? {
172            Some(idx) if idx < StageId(Self::LEN) => {
173                // do nothing; we are resuming
174            }
175            Some(idx) if idx == StageId(Self::LEN) => {
176                // perform the stage, but don't set it
177
178                let stage = &mut self.0;
179
180                stage.perform_restartable(fuzzer, executor, state, manager)?;
181
182                state.clear_stage_id()?;
183            }
184            Some(idx) if idx > StageId(Self::LEN) => {
185                unreachable!("We should clear the stage index before we get here...");
186            }
187            // this is None, but the match can't deduce that
188            _ => {
189                state.set_current_stage_id(StageId(Self::LEN))?;
190
191                let stage = &mut self.0;
192
193                stage.perform_restartable(fuzzer, executor, state, manager)?;
194
195                state.clear_stage_id()?;
196            }
197        }
198
199        // Mark the elapsed time for the scheduler
200        #[cfg(feature = "introspection")]
201        state.introspection_stats_mut().finish_stage();
202
203        if state.stop_requested() {
204            state.discard_stop_request();
205            manager.on_shutdown()?;
206            return Err(Error::shutting_down());
207        }
208
209        // Execute the remaining stages
210        self.1.perform_all(fuzzer, executor, state, manager)
211    }
212}
213
214impl<Head, Tail, E, EM, S, Z> IntoVec<Box<dyn Stage<E, EM, S, Z>>> for (Head, Tail)
215where
216    Head: Stage<E, EM, S, Z> + 'static,
217    Tail: StagesTuple<E, EM, S, Z> + HasConstLen + IntoVec<Box<dyn Stage<E, EM, S, Z>>>,
218    S: HasCurrentStageId,
219{
220    fn into_vec_reversed(self) -> Vec<Box<dyn Stage<E, EM, S, Z>>> {
221        let (head, tail) = self.uncons();
222        let mut ret = tail.0.into_vec_reversed();
223        ret.push(Box::new(head));
224        ret
225    }
226
227    fn into_vec(self) -> Vec<Box<dyn Stage<E, EM, S, Z>>> {
228        let mut ret = self.into_vec_reversed();
229        ret.reverse();
230        ret
231    }
232}
233
234impl<Tail, E, EM, S, Z> IntoVec<Box<dyn Stage<E, EM, S, Z>>> for (Tail,)
235where
236    Tail: IntoVec<Box<dyn Stage<E, EM, S, Z>>>,
237{
238    fn into_vec(self) -> Vec<Box<dyn Stage<E, EM, S, Z>>> {
239        self.0.into_vec()
240    }
241}
242
243impl<E, EM, S, Z> IntoVec<Box<dyn Stage<E, EM, S, Z>>> for Vec<Box<dyn Stage<E, EM, S, Z>>> {
244    fn into_vec(self) -> Vec<Box<dyn Stage<E, EM, S, Z>>> {
245        self
246    }
247}
248
249trait RestartableStage<E, EM, S, Z>: Stage<E, EM, S, Z> + Restartable<S> {
250    fn perform_restartable(
251        &mut self,
252        fuzzer: &mut Z,
253        executor: &mut E,
254        state: &mut S,
255        manager: &mut EM,
256    ) -> Result<(), Error>;
257}
258
259impl<E, EM, S, ST, Z> RestartableStage<E, EM, S, Z> for ST
260where
261    ST: Stage<E, EM, S, Z> + Restartable<S>,
262{
263    /// Run the stage, calling [`Stage::should_restart`] and [`Stage::clear_progress`] appropriately
264    fn perform_restartable(
265        &mut self,
266        fuzzer: &mut Z,
267        executor: &mut E,
268        state: &mut S,
269        manager: &mut EM,
270    ) -> Result<(), Error> {
271        if self.should_restart(state)? {
272            self.perform(fuzzer, executor, state, manager)?;
273        }
274        self.clear_progress(state)
275    }
276}
277
278impl<E, EM, S, Z> StagesTuple<E, EM, S, Z> for Vec<Box<dyn RestartableStage<E, EM, S, Z>>>
279where
280    EM: SendExiting,
281    S: HasCurrentStageId + Stoppable,
282{
283    /// Performs all stages in the `Vec`
284    /// Checks after every stage if state wants to stop
285    /// and returns an [`Error::ShuttingDown`] if so
286    fn perform_all(
287        &mut self,
288        fuzzer: &mut Z,
289        executor: &mut E,
290        state: &mut S,
291        manager: &mut EM,
292    ) -> Result<(), Error> {
293        self.iter_mut().try_for_each(|stage| {
294            if state.stop_requested() {
295                state.discard_stop_request();
296                manager.on_shutdown()?;
297                return Err(Error::shutting_down());
298            }
299            stage.perform_restartable(fuzzer, executor, state, manager)
300        })
301    }
302}
303
304static mut CLOSURE_STAGE_ID: usize = 0;
305/// The name for closure stage
306pub static CLOSURE_STAGE_NAME: &str = "closure";
307
308/// A [`Stage`] that will call a closure
309#[derive(Debug)]
310pub struct ClosureStage<CB, E, EM, Z> {
311    name: Cow<'static, str>,
312    closure: CB,
313    phantom: PhantomData<(E, EM, Z)>,
314}
315
316impl<CB, E, EM, Z> Named for ClosureStage<CB, E, EM, Z> {
317    fn name(&self) -> &Cow<'static, str> {
318        &self.name
319    }
320}
321
322impl<CB, E, EM, S, Z> Stage<E, EM, S, Z> for ClosureStage<CB, E, EM, Z>
323where
324    CB: FnMut(&mut Z, &mut E, &mut S, &mut EM) -> Result<(), Error>,
325    S: HasNamedMetadata + HasCurrentCorpusId,
326{
327    fn perform(
328        &mut self,
329        fuzzer: &mut Z,
330        executor: &mut E,
331        state: &mut S,
332        manager: &mut EM,
333    ) -> Result<(), Error> {
334        (self.closure)(fuzzer, executor, state, manager)
335    }
336}
337
338impl<CB, E, EM, S, Z> Restartable<S> for ClosureStage<CB, E, EM, Z>
339where
340    S: HasNamedMetadata + HasCurrentCorpusId,
341{
342    #[inline]
343    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
344        // There's no restart safety in the content of the closure.
345        // don't restart
346        RetryCountRestartHelper::no_retry(state, &self.name)
347    }
348
349    #[inline]
350    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
351        RetryCountRestartHelper::clear_progress(state, &self.name)
352    }
353}
354
355/// A stage that takes a closure
356impl<CB, E, EM, Z> ClosureStage<CB, E, EM, Z> {
357    /// Create a new [`ClosureStage`]
358    #[must_use]
359    pub fn new(closure: CB) -> Self {
360        // unsafe but impossible that you create two threads both instantiating this instance
361        let stage_id = unsafe {
362            let ret = CLOSURE_STAGE_ID;
363            CLOSURE_STAGE_ID += 1;
364            ret
365        };
366        Self {
367            name: Cow::Owned(CLOSURE_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_ref()),
368            closure,
369            phantom: PhantomData,
370        }
371    }
372}
373
374/// Progress which permits a fixed amount of resumes per round of fuzzing. If this amount is ever
375/// exceeded, the input will no longer be executed by this stage.
376#[derive(Clone, Deserialize, Serialize, Debug)]
377pub struct RetryCountRestartHelper {
378    tries_remaining: Option<usize>,
379    skipped: HashSet<CorpusId>,
380}
381
382impl_serdeany!(RetryCountRestartHelper);
383
384impl RetryCountRestartHelper {
385    /// Don't allow restart
386    pub fn no_retry<S>(state: &mut S, name: &str) -> Result<bool, Error>
387    where
388        S: HasNamedMetadata + HasCurrentCorpusId,
389    {
390        Self::should_restart(state, name, 1)
391    }
392
393    /// Initializes (or counts down in) the progress helper, giving it the amount of max retries
394    ///
395    /// Returns `true` if the stage should run
396    pub fn should_restart<S>(state: &mut S, name: &str, max_retries: usize) -> Result<bool, Error>
397    where
398        S: HasNamedMetadata + HasCurrentCorpusId,
399    {
400        let corpus_id = state.current_corpus_id()?.ok_or_else(|| {
401            Error::illegal_state(
402                "No current_corpus_id set in State, but called RetryCountRestartHelper::should_skip",
403            )
404        })?;
405
406        let initial_tries_remaining = max_retries + 1;
407        let metadata = state.named_metadata_or_insert_with(name, || Self {
408            tries_remaining: Some(initial_tries_remaining),
409            skipped: HashSet::new(),
410        });
411        let tries_remaining = metadata
412            .tries_remaining
413            .unwrap_or(initial_tries_remaining)
414            .checked_sub(1)
415            .ok_or_else(|| {
416                Error::illegal_state(
417                    "Attempted further retries after we had already gotten to none remaining.",
418                )
419            })?;
420
421        metadata.tries_remaining = Some(tries_remaining);
422
423        Ok(if tries_remaining == 0 {
424            metadata.skipped.insert(corpus_id);
425            false
426        } else if metadata.skipped.contains(&corpus_id) {
427            // skip this testcase, we already retried it often enough...
428            false
429        } else {
430            true
431        })
432    }
433
434    /// Clears the progress
435    pub fn clear_progress<S>(state: &mut S, name: &str) -> Result<(), Error>
436    where
437        S: HasNamedMetadata,
438    {
439        state.named_metadata_mut::<Self>(name)?.tries_remaining = None;
440        Ok(())
441    }
442}
443
444/// The index of a stage
445#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
446#[repr(transparent)]
447pub struct StageId(pub(crate) usize);
448
449impl fmt::Display for StageId {
450    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
451        write!(f, "{}", self.0)
452    }
453}
454
455impl_serdeany!(ExecutionCountRestartHelperMetadata);
456
457/// `SerdeAny` metadata used to keep track of executions since start for a given stage.
458#[derive(Debug, Clone, Serialize, Deserialize)]
459pub struct ExecutionCountRestartHelperMetadata {
460    /// How many executions we had when we started this stage initially (this round)
461    started_at_execs: u64,
462}
463
464/// A tool shed of functions to be used for stages that try to run for `n` iterations.
465///
466/// # Note
467/// This helper assumes resumable mutational stages are not nested.
468/// If you want to nest them, you will have to switch all uses of `metadata` in this helper to `named_metadata` instead.
469#[derive(Debug, Default, Clone)]
470pub struct ExecutionCountRestartHelper {
471    /// At what exec count this Stage was started (cache)
472    /// Only used as cache for the value stored in [`MutationalStageMetadata`].
473    started_at_execs: Option<u64>,
474}
475
476impl ExecutionCountRestartHelper {
477    /// Create a new [`ExecutionCountRestartHelperMetadata`]
478    #[must_use]
479    pub fn new() -> Self {
480        Self {
481            started_at_execs: None,
482        }
483    }
484
485    /// The execs done since start of this [`Stage`]/helper
486    pub fn execs_since_progress_start<S>(&mut self, state: &mut S, name: &str) -> Result<u64, Error>
487    where
488        S: HasNamedMetadata + HasExecutions,
489    {
490        let started_at_execs = if let Some(started_at_execs) = self.started_at_execs {
491            started_at_execs
492        } else {
493            state
494                .named_metadata::<ExecutionCountRestartHelperMetadata>(name)
495                .map(|x| {
496                    self.started_at_execs = Some(x.started_at_execs);
497                    x.started_at_execs
498                })
499                .map_err(|err| {
500                    Error::illegal_state(format!(
501                        "The ExecutionCountRestartHelperMetadata should have been set at this point - {err}"
502                    ))
503                })?
504        };
505        Ok(state.executions() - started_at_execs)
506    }
507
508    /// Initialize progress for the stage this wrapper wraps.
509    pub fn should_restart<S>(&mut self, state: &mut S, name: &str) -> Result<bool, Error>
510    where
511        S: HasNamedMetadata + HasExecutions,
512    {
513        let executions = *state.executions();
514        let metadata =
515            state.named_metadata_or_insert_with(name, || ExecutionCountRestartHelperMetadata {
516                started_at_execs: executions,
517            });
518        self.started_at_execs = Some(metadata.started_at_execs);
519        Ok(true)
520    }
521
522    /// Clear progress for the stage this wrapper wraps.
523    pub fn clear_progress<S>(&mut self, state: &mut S, name: &str) -> Result<(), Error>
524    where
525        S: HasNamedMetadata,
526    {
527        self.started_at_execs = None;
528        let _metadata = state.remove_named_metadata::<ExecutionCountRestartHelperMetadata>(name);
529        debug_assert!(
530            _metadata.is_some(),
531            "Called clear_progress, but should_restart was not called before (or did mutational stages get nested?)"
532        );
533        Ok(())
534    }
535}
536
537#[cfg(test)]
538mod test {
539    use alloc::borrow::Cow;
540
541    use libafl_bolts::{Error, Named};
542
543    use crate::{
544        corpus::{Corpus, HasCurrentCorpusId, Testcase},
545        inputs::NopInput,
546        stages::RetryCountRestartHelper,
547        state::{HasCorpus, StdState},
548    };
549
550    /// Test to test retries in stages
551    #[test]
552    fn test_tries_progress() -> Result<(), Error> {
553        struct StageWithOneTry;
554
555        impl Named for StageWithOneTry {
556            fn name(&self) -> &Cow<'static, str> {
557                static NAME: Cow<'static, str> = Cow::Borrowed("TestStage");
558                &NAME
559            }
560        }
561
562        // # Safety
563        // No concurrency per testcase
564        #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
565        unsafe {
566            RetryCountRestartHelper::register();
567        }
568
569        let mut state = StdState::nop()?;
570        let stage = StageWithOneTry;
571
572        let corpus_id = state.corpus_mut().add(Testcase::new(NopInput {}))?;
573
574        state.set_corpus_id(corpus_id)?;
575
576        for _ in 0..10 {
577            // used normally, no retries means we never skip
578            assert!(RetryCountRestartHelper::should_restart(
579                &mut state,
580                stage.name(),
581                1
582            )?);
583            RetryCountRestartHelper::clear_progress(&mut state, stage.name())?;
584        }
585
586        for _ in 0..10 {
587            // used normally, only one retry means we never skip
588            assert!(RetryCountRestartHelper::should_restart(
589                &mut state,
590                stage.name(),
591                2
592            )?);
593            assert!(RetryCountRestartHelper::should_restart(
594                &mut state,
595                stage.name(),
596                2
597            )?);
598            RetryCountRestartHelper::clear_progress(&mut state, stage.name())?;
599        }
600
601        assert!(RetryCountRestartHelper::should_restart(
602            &mut state,
603            stage.name(),
604            2
605        )?);
606        // task failed, let's resume
607        // we still have one more try!
608        assert!(RetryCountRestartHelper::should_restart(
609            &mut state,
610            stage.name(),
611            2
612        )?);
613
614        // task failed, let's resume
615        // out of retries, so now we skip
616        assert!(!RetryCountRestartHelper::should_restart(
617            &mut state,
618            stage.name(),
619            2
620        )?);
621        RetryCountRestartHelper::clear_progress(&mut state, stage.name())?;
622
623        // we previously exhausted this testcase's retries, so we skip
624        assert!(!RetryCountRestartHelper::should_restart(
625            &mut state,
626            stage.name(),
627            2
628        )?);
629        RetryCountRestartHelper::clear_progress(&mut state, stage.name())?;
630
631        Ok(())
632    }
633}
libafl/stages/mod.rs

libafl/stages/
mod.rs
