libafl/inputs/
multi.rs

1//! An input composed of multiple parts identified by a key.
2
3use alloc::{string::String, vec::Vec};
4use core::{fmt::Debug, hash::Hash};
5
6use serde::{Serialize, de::DeserializeOwned};
7
8use crate::{
9    corpus::CorpusId,
10    inputs::{Input, ListInput},
11};
12
13/// An input composed of multiple parts, each identified by a key.
14///
15/// It relies on a list to store the keys and parts. Keys may appear multiple times.
16pub type MultipartInput<I, K> = ListInput<(K, I)>;
17
18impl<I, K> Input for MultipartInput<I, K>
19where
20    I: Input,
21    K: PartialEq + Debug + Serialize + DeserializeOwned + Clone + Hash,
22{
23    fn generate_name(&self, id: Option<CorpusId>) -> String {
24        self.parts()
25            .iter()
26            .map(|(_k, i)| i.generate_name(id))
27            .collect::<Vec<_>>()
28            .join(",")
29    }
30}
31
32/// Trait for types that provide a way to access parts by key.
33pub trait Keyed<K, V> {
34    /// Get the keys of the parts of this input.
35    ///
36    /// Keys may appear multiple times if they are used multiple times in the input.
37    fn keys<'a>(&'a self) -> impl Iterator<Item = &'a K>
38    where
39        K: 'a;
40
41    /// Get a reference to each part with the provided key along with its index.
42    fn with_key<'a, 'b>(&'b self, key: &'a K) -> impl Iterator<Item = (usize, &'b V)> + 'a
43    where
44        'b: 'a,
45        V: 'b;
46
47    /// Gets a mutable reference to each part with the provided key along with its index.
48    fn with_key_mut<'a, 'b>(
49        &'b mut self,
50        key: &'a K,
51    ) -> impl Iterator<Item = (usize, &'b mut V)> + 'a
52    where
53        'b: 'a,
54        V: 'b;
55}
56
57impl<I, K> Keyed<K, I> for MultipartInput<I, K>
58where
59    K: PartialEq,
60{
61    fn keys<'a>(&'a self) -> impl Iterator<Item = &'a K>
62    where
63        K: 'a,
64    {
65        self.parts().iter().map(|(k, _)| k)
66    }
67
68    fn with_key<'a, 'b>(&'b self, key: &'a K) -> impl Iterator<Item = (usize, &'b I)> + 'a
69    where
70        'b: 'a,
71        I: 'b,
72    {
73        self.parts()
74            .iter()
75            .enumerate()
76            .filter_map(move |(i, (k, input))| (key == k).then_some((i, input)))
77    }
78
79    fn with_key_mut<'a, 'b>(
80        &'b mut self,
81        key: &'a K,
82    ) -> impl Iterator<Item = (usize, &'b mut I)> + 'a
83    where
84        'b: 'a,
85        I: 'b,
86    {
87        self.parts_mut()
88            .iter_mut()
89            .enumerate()
90            .filter_map(move |(i, (k, input))| (key == k).then_some((i, input)))
91    }
92}
libafl/inputs/multi.rs

libafl/inputs/
multi.rs
