feroxfuzz 1.0.0-rc.13

#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

use super::{set_states_corpus_index, CorpusIndex, Scheduler};
use crate::error::FeroxFuzzError;
use crate::state::SharedState;
use crate::std_ext::ops::Len;
use crate::std_ext::tuple::Named;

use tracing::{error, instrument, trace};

/// In-order access of the associated [`Corpus`]
///
/// [`Corpus`]: crate::corpora::Corpus
///
/// # Examples
///
/// if you have a corpus with the following entries:
///
/// `FUZZ_USER`: ["user1", "user2", "user3"]
/// `FUZZ_PASS`: ["pass1", "pass2", "pass3"]
///
/// and a fuzzable url defined as
///
/// `http://example.com/login?username=FUZZ_USER&password=FUZZ_PASS`
///
///
/// then the resultant `OrderedScheduler` scheduling of the two corpora would be:
///
/// `http://example.com/login?username=user1&password=pass1`
/// `http://example.com/login?username=user2&password=pass2`
/// `http://example.com/login?username=user3&password=pass3`
///
#[allow(clippy::doc_link_with_quotes)]
#[derive(Clone, Debug, Default)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct OrderedScheduler {
    current: usize,
    indices: Vec<CorpusIndex>,

    #[cfg_attr(feature = "serde", serde(skip))]
    state: SharedState,
}

impl Scheduler for OrderedScheduler {
    #[instrument(skip(self), fields(%self.current, ?self.indices), level = "trace")]
    fn next(&mut self) -> Result<(), FeroxFuzzError> {
        // iterate through the indices and increment the current index
        for index in &mut self.indices {
            if self.current > 0 && index.should_reset(self.current) {
                // once any of the indices reaches the end of its loop, the entire
                // scheduler has run to completion
                trace!("scheduler has run to completion");
                return Err(FeroxFuzzError::IterationStopped);
            }

            set_states_corpus_index(&self.state, index.name(), index.current())?;

            index.next()?;
        }

        self.current += 1; // update the total number of times .next has been called

        Ok(())
    }

    /// resets all indexes that are tracked by the scheduler as well as their associated atomic
    /// indexes in the [`SharedState`] instance
    #[instrument(skip(self), level = "trace")]
    fn reset(&mut self) {
        self.current = 0;

        self.indices.iter_mut().for_each(|index| {
            // first, we get the corpus associated with the current corpus_index
            let corpus = self.state.corpus_by_name(index.name()).unwrap();

            // and then get its length
            let len = corpus.len();

            // if any items were added to the corpus, we'll need to update the length/expected iterations
            // accordingly
            index.update_length(len);
            index.update_iterations(len);

            // we'll also reset the current index as well
            index.reset();

            // finally, we get the SharedState's view of the index in sync with the Scheduler's
            //
            // i.e. at this point, the state and local indices should all be 0, and any items
            // added to the corpus should be reflected in each index's length/iterations
            set_states_corpus_index(&self.state, index.name(), 0).unwrap();
        });

        trace!("scheduler has been reset");
    }

    fn update_length(&mut self) {
        // basically the same logic as reset, but we don't need to reset the index, nor reset
        // the state's view of the index
        self.indices.iter_mut().for_each(|index| {
            // first, we get the corpus associated with the current corpus_index
            let corpus = self.state.corpus_by_name(index.name()).unwrap();

            // and then get its length
            let len = corpus.len();

            // if any items were added to the corpus, we'll need to update the length/expected iterations
            // accordingly
            index.update_length(len);
            index.update_iterations(len);
        });
    }
}

impl OrderedScheduler {
    /// create a new `OrderedScheduler`
    ///
    /// # Errors
    ///
    /// This function will return an error if any corpus found in the `SharedState`'s
    /// `corpora` map is empty, or if the `SharedState`'s `corpora` map is empty.
    ///
    /// # Examples
    ///
    /// see `examples/cartesian-product.rs` for a more robust example
    /// and explanation
    ///
    /// ```
    /// use feroxfuzz::schedulers::{Scheduler, ProductScheduler};
    /// use feroxfuzz::prelude::*;
    /// use feroxfuzz::corpora::{RangeCorpus, Wordlist};
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// // create two corpora, one with a set of user names, and one with a range of ids
    /// // where only even ids are considered
    /// let users = Wordlist::new().word("user").word("admin").name("users").build();
    /// let ids = RangeCorpus::with_stop(5).name("ids").build()?;
    ///
    /// let state = SharedState::with_corpora([ids, users]);
    ///
    /// let order = ["users", "ids"];
    /// let mut scheduler = ProductScheduler::new(order, state.clone())?;
    ///
    /// let mut counter = 0;
    ///
    /// while Scheduler::next(&mut scheduler).is_ok() {
    ///     counter += 1;
    /// }
    ///
    /// // users.len() * ids.len() = 2 * 5 = 10
    /// assert_eq!(counter, 10);
    ///
    /// # Ok(())
    /// # }
    #[inline]
    #[instrument(skip_all, level = "trace")]
    #[allow(clippy::cast_possible_truncation)]
    #[allow(clippy::cast_sign_loss)]
    pub fn new(state: SharedState) -> Result<Self, FeroxFuzzError> {
        let corpora = state.corpora();

        let mut current = state
            .stats()
            .read()
            .map_or(0, |stats| stats.requests() as usize);

        let mut indices = Vec::with_capacity(corpora.len());

        for (name, corpus) in &*corpora {
            let length = corpus.len();

            if length == 0 {
                // one of the corpora was empty
                error!(%name, "corpus is empty");

                return Err(FeroxFuzzError::EmptyCorpus { name: name.clone() });
            }

            // the total number of expected iterations per corpus is simply
            // the length of the corpus
            indices.push(CorpusIndex::new(name, length, length));
        }

        if indices.is_empty() {
            // empty iterator passed in
            error!("no corpora were found");
            return Err(FeroxFuzzError::EmptyCorpusMap);
        }

        let mut scheduler = Self {
            current,
            indices,
            state,
        };

        while current > 0 {
            Scheduler::next(&mut scheduler)?;
            current -= 1;
        }

        Ok(scheduler)
    }

    /// by default, the [`OrderedScheduler`] will iterate through *every* corpus
    /// and stop iteration when the shortest corpus is exhausted.
    ///
    /// This function allows you to specify which corpora will be updated by the
    /// scheduler. This is useful if you want to only update a subset of corpora
    /// at a time, instead of all of them.
    #[instrument(skip(self), level = "trace")]
    pub fn limit_to_corpora(&mut self, corpora: &[&str]) {
        self.indices.retain(|index| corpora.contains(&index.name()));
    }
}

#[allow(clippy::copy_iterator)]
impl Iterator for OrderedScheduler {
    type Item = ();

    fn next(&mut self) -> Option<Self::Item> {
        Scheduler::next(self).ok()
    }
}

impl Named for OrderedScheduler {
    #[allow(clippy::unnecessary_literal_bound)]
    fn name(&self) -> &str {
        "OrderedScheduler"
    }
}