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
136
137

//! The `GeneralizedInput` is an input that ca be generalized to represent a rule, used by Grimoire

use alloc::vec::Vec;

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

use crate::{
    corpus::{CorpusId, Testcase},
    inputs::BytesInput,
    stages::mutational::{MutatedTransform, MutatedTransformPost},
    state::{HasCorpus, HasMetadata},
    Error,
};

/// An item of the generalized input
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash)]
pub enum GeneralizedItem {
    /// Real bytes
    Bytes(Vec<u8>),
    /// An insertion point
    Gap,
}

/// Metadata regarding the generalised content of an input
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(
    any(not(feature = "serdeany_autoreg"), miri),
    allow(clippy::unsafe_derive_deserialize)
)] // for SerdeAny
pub struct GeneralizedInputMetadata {
    generalized: Vec<GeneralizedItem>,
}

impl_serdeany!(GeneralizedInputMetadata);

impl GeneralizedInputMetadata {
    /// Fill the generalized vector from a slice of option (None -> Gap)
    #[must_use]
    pub fn generalized_from_options(v: &[Option<u8>]) -> Self {
        let mut generalized = vec![];
        let mut bytes = vec![];
        if v.first() != Some(&None) {
            generalized.push(GeneralizedItem::Gap);
        }
        for e in v {
            match e {
                None => {
                    if !bytes.is_empty() {
                        generalized.push(GeneralizedItem::Bytes(bytes.clone()));
                        bytes.clear();
                    }
                    generalized.push(GeneralizedItem::Gap);
                }
                Some(b) => {
                    bytes.push(*b);
                }
            }
        }
        if !bytes.is_empty() {
            generalized.push(GeneralizedItem::Bytes(bytes));
        }
        if generalized.last() != Some(&GeneralizedItem::Gap) {
            generalized.push(GeneralizedItem::Gap);
        }
        Self { generalized }
    }

    /// Get the size of the generalized
    #[must_use]
    pub fn generalized_len(&self) -> usize {
        let mut size = 0;
        for item in &self.generalized {
            match item {
                GeneralizedItem::Bytes(b) => size += b.len(),
                GeneralizedItem::Gap => size += 1,
            }
        }
        size
    }

    /// Convert generalized to bytes
    #[must_use]
    pub fn generalized_to_bytes(&self) -> Vec<u8> {
        self.generalized
            .iter()
            .filter_map(|item| match item {
                GeneralizedItem::Bytes(bytes) => Some(bytes),
                GeneralizedItem::Gap => None,
            })
            .flatten()
            .copied()
            .collect()
    }

    /// Get the generalized input
    #[must_use]
    pub fn generalized(&self) -> &[GeneralizedItem] {
        &self.generalized
    }

    /// Get the generalized input (mutable)
    pub fn generalized_mut(&mut self) -> &mut Vec<GeneralizedItem> {
        &mut self.generalized
    }
}

impl<S> MutatedTransform<BytesInput, S> for GeneralizedInputMetadata
where
    S: HasCorpus,
{
    type Post = Self;

    fn try_transform_from(
        base: &mut Testcase<BytesInput>,
        _state: &S,
        corpus_idx: CorpusId,
    ) -> Result<Self, Error> {
        let meta = base
            .metadata_map()
            .get::<GeneralizedInputMetadata>()
            .ok_or_else(|| {
                Error::key_not_found(format!(
                    "Couldn't find the GeneralizedInputMetadata for corpus entry {corpus_idx}",
                ))
            })
            .cloned()?;

        Ok(meta)
    }

    fn try_transform_into(self, _state: &S) -> Result<(BytesInput, Self::Post), Error> {
        Ok((BytesInput::from(self.generalized_to_bytes()), self))
    }
}

impl<S> MutatedTransformPost<S> for GeneralizedInputMetadata where S: HasCorpus {}