openmina_fuzzer/
lib.rs

1use rand::{rngs::SmallRng, Rng, SeedableRng};
2use rand_distr::{Distribution, Exp};
3
4#[derive(Debug, Clone, Deserialize)]
5pub struct FuzzerConf {
6    pub rng_seed: u64,
7    pub max_extend_size: usize,
8    // rates are in the 0-1000 range or none if the specific mutation is disabled
9    pub pnet_mutation_rate: Option<u32>,
10    pub noise_mutation_rate: Option<u32>,
11    pub select_authentication_mutation_rate: Option<u32>,
12    pub select_multiplexing_mutation_rate: Option<u32>,
13    pub select_stream_mutation_rate: Option<u32>,
14    pub yamux_frame_mutation_rate: Option<u32>,
15    pub yamux_flags_mutation_rate: Option<u32>,
16    pub identify_msg_mutation_rate: Option<u32>,
17    pub kad_data_mutation_rate: Option<u32>,
18    pub rpc_data_mutation_rate: Option<u32>,
19    pub pubsub_mutation_rate: Option<u32>,
20}
21
22#[derive(Debug, Clone)]
23pub struct FuzzerState {
24    pub rng: SmallRng,
25    pub conf: FuzzerConf,
26}
27
28#[derive(Debug)]
29pub enum MutationStrategy {
30    FlipBits = 0,
31    ExtendRandom = 1,
32    ExtendCopy = 2,
33    Shrink = 3,
34}
35
36impl MutationStrategy {
37    pub fn range() -> std::ops::RangeInclusive<u32> {
38        MutationStrategy::FlipBits as u32..=MutationStrategy::Shrink as u32
39    }
40}
41
42impl From<u32> for MutationStrategy {
43    fn from(item: u32) -> Self {
44        match item {
45            0 => MutationStrategy::FlipBits,
46            1 => MutationStrategy::ExtendRandom,
47            2 => MutationStrategy::ExtendCopy,
48            3 => MutationStrategy::Shrink,
49            _ => unreachable!(),
50        }
51    }
52}
53
54impl FuzzerState {
55    pub fn gen_ratio(&mut self, numerator: Option<u32>) -> bool {
56        if let Some(numerator) = numerator {
57            self.rng.gen_ratio(numerator, 1000)
58        } else {
59            false
60        }
61    }
62
63    pub fn gen_flips(&mut self, max_size: usize) -> usize {
64        // Favor small number of flips
65        let exp = Exp::new(2.0).unwrap();
66        let v = exp.sample(&mut self.rng);
67
68        min(max(1, (v * 10.0) as usize), max_size)
69    }
70
71    pub fn flip_bits(&mut self, data: &mut u8) {
72        let num_bit_flips = self.gen_flips(8);
73
74        for _ in 0..num_bit_flips {
75            *data ^= 1u8 << self.rng.gen_range(0..8);
76        }
77    }
78
79    pub fn flip_bytes(&mut self, data: &mut [u8]) {
80        let max_size = data.len();
81        let num_byte_flips = self.gen_flips(max_size);
82
83        for _ in 0..num_byte_flips {
84            let rnd_index = self.rng.gen_range(0..max_size);
85
86            self.flip_bits(&mut data[rnd_index]);
87        }
88    }
89
90    pub fn extend_random(&mut self, data: &[u8]) -> Vec<u8> {
91        let extend_size = self.rng.gen_range(1..self.conf.max_extend_size);
92        let random_pos = self.rng.gen_range(0..=data.len());
93
94        let mut random_bytes = vec![0u8; extend_size];
95        self.rng.fill(&mut random_bytes[..]);
96
97        let mut result = Vec::with_capacity(data.len() + extend_size);
98
99        result.extend(&data[0..random_pos]);
100        result.extend(random_bytes);
101        result.extend(&data[random_pos..]);
102        result
103    }
104
105    pub fn extend_copy(&mut self, data: &[u8]) -> Vec<u8> {
106        let random_pos = self.rng.gen_range(0..=data.len());
107        let copy_pos = self.rng.gen_range(0..=data.len() - 1);
108        let copy_size = self.rng.gen_range(1..=data.len() - copy_pos);
109        let mut extend_size = self.rng.gen_range(copy_size..self.conf.max_extend_size);
110        let mut result = Vec::with_capacity(data.len() + extend_size);
111
112        result.extend(&data[0..random_pos]);
113
114        while extend_size > copy_size {
115            result.extend(&data[copy_pos..copy_pos + copy_size]);
116            //self.flip_bytes(result.as_mut_slice());
117            extend_size -= copy_size;
118        }
119
120        result.extend(&data[random_pos..]);
121        result
122    }
123
124    pub fn shrink(&mut self, data: &[u8]) -> Vec<u8> {
125        let pos = self.rng.gen_range(0..=data.len() - 1);
126        let size = self.rng.gen_range(1..=data.len() - pos);
127        let mut result = data[..pos].to_vec();
128
129        result.extend_from_slice(&data[pos + size..]);
130        result
131    }
132}
133
134use lazy_static::lazy_static;
135use serde::Deserialize;
136use std::{
137    cmp::{max, min},
138    env,
139    fs::File,
140    sync::Mutex,
141};
142
143lazy_static! {
144    pub static ref FUZZER: Option<Mutex<FuzzerState>> = {
145        if let Ok(conf) = env::var("FUZZER_CONF") {
146            let file = File::open(conf).expect("Failed to open fuzzer conf file");
147            let conf: FuzzerConf =
148                serde_json::from_reader(file).expect("Failed to parse fuzzer configuration");
149
150            println!("[+] FUZZER conf: {:?}", conf);
151
152            Some(Mutex::new(FuzzerState {
153                rng: SmallRng::seed_from_u64(conf.rng_seed),
154                conf,
155            }))
156        } else {
157            println!("[-] no FUZZER conf specified");
158            None
159        }
160    };
161}
162
163#[macro_export]
164macro_rules! fuzz {
165    ($expr:expr, $mutator:expr) => {
166        if let Some(fuzzer) = $crate::FUZZER.as_ref() {
167            let mut fuzzer = fuzzer.lock().expect("can't lock FUZZER");
168            $mutator(&mut fuzzer, $expr);
169        }
170    };
171}
172
173#[macro_export]
174macro_rules! fuzzed {
175    ($expr:expr, $mutator:expr) => {
176        if let Some(fuzzer) = $crate::FUZZER.as_ref() {
177            let mut fuzzer = fuzzer.lock().expect("can't lock FUZZER");
178            let mut fuzzed = $expr;
179            $mutator(&mut fuzzer, &mut fuzzed);
180            fuzzed
181        } else {
182            $expr
183        }
184    };
185}
openmina_fuzzer/lib.rs

openmina_fuzzer/
lib.rs
