1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

//! The queue corpus scheduler implements an AFL-like queue mechanism
//! The [`TuneableScheduler`] extends the queue scheduler with a method to
//! chose the next corpus entry manually

use alloc::borrow::ToOwned;
use core::marker::PhantomData;

use libafl_bolts::impl_serdeany;
use serde::{Deserialize, Serialize};

use super::RemovableScheduler;
use crate::{
    corpus::{Corpus, CorpusId, HasTestcase},
    schedulers::Scheduler,
    state::{HasCorpus, HasMetadata, State, UsesState},
    Error,
};

#[derive(Default, Clone, Copy, Eq, PartialEq, Debug, Serialize, Deserialize)]
#[cfg_attr(
    any(not(feature = "serdeany_autoreg"), miri),
    allow(clippy::unsafe_derive_deserialize)
)] // for SerdeAny
struct TuneableSchedulerMetadata {
    next: Option<CorpusId>,
}

impl_serdeany!(TuneableSchedulerMetadata);

/// Walk the corpus in a queue-like fashion
/// With the specific `set_next` method, we can chose the next corpus entry manually
#[derive(Debug, Clone)]
pub struct TuneableScheduler<S> {
    phantom: PhantomData<S>,
}

impl<S> TuneableScheduler<S>
where
    S: HasMetadata + HasCorpus,
{
    /// Creates a new `TuneableScheduler`
    #[must_use]
    pub fn new(state: &mut S) -> Self {
        if !state.has_metadata::<TuneableSchedulerMetadata>() {
            state.add_metadata(TuneableSchedulerMetadata::default());
        }
        Self {
            phantom: PhantomData,
        }
    }

    fn metadata_mut(state: &mut S) -> &mut TuneableSchedulerMetadata {
        state
            .metadata_map_mut()
            .get_mut::<TuneableSchedulerMetadata>()
            .unwrap()
    }

    fn metadata(state: &S) -> &TuneableSchedulerMetadata {
        state
            .metadata_map()
            .get::<TuneableSchedulerMetadata>()
            .unwrap()
    }

    /// Sets the next corpus id to be used
    pub fn set_next(state: &mut S, next: CorpusId) {
        Self::metadata_mut(state).next = Some(next);
    }

    /// Gets the next set corpus id
    pub fn get_next(state: &S) -> Option<CorpusId> {
        Self::metadata(state).next
    }

    /// Resets this to a queue scheduler
    pub fn reset(state: &mut S) {
        let metadata = Self::metadata_mut(state);
        metadata.next = None;
    }

    /// Gets the current corpus entry id
    pub fn get_current(state: &S) -> CorpusId {
        state
            .corpus()
            .current()
            .unwrap_or_else(|| state.corpus().first().expect("Empty corpus"))
    }
}

impl<S> UsesState for TuneableScheduler<S>
where
    S: State,
{
    type State = S;
}

impl<S> RemovableScheduler for TuneableScheduler<S> where
    S: HasCorpus + HasMetadata + HasTestcase + State
{
}

impl<S> Scheduler for TuneableScheduler<S>
where
    S: HasCorpus + HasMetadata + HasTestcase + State,
{
    fn on_add(&mut self, state: &mut Self::State, idx: CorpusId) -> Result<(), Error> {
        // Set parent id
        let current_idx = *state.corpus().current();
        state
            .corpus()
            .get(idx)?
            .borrow_mut()
            .set_parent_id_optional(current_idx);

        Ok(())
    }

    /// Gets the next entry in the queue
    fn next(&mut self, state: &mut Self::State) -> Result<CorpusId, Error> {
        if state.corpus().count() == 0 {
            return Err(Error::empty("No entries in corpus".to_owned()));
        }
        let id = if let Some(next) = Self::get_next(state) {
            // next was set
            next
        } else if let Some(next) = state.corpus().next(Self::get_current(state)) {
            next
        } else {
            state.corpus().first().unwrap()
        };
        self.set_current_scheduled(state, Some(id))?;
        Ok(id)
    }
}