bolero/test/
input.rs

1#![cfg_attr(fuzzing_random, allow(dead_code))]
2
3use bolero_engine::{rng::Recommended as Rng, Seed};
4use bolero_generator::{driver, TypeGenerator};
5use rand::SeedableRng;
6use std::{io::Read, path::PathBuf};
7
8pub use bolero_engine::input::*;
9
10pub enum Test {
11    File(FileTest),
12    Rng(RngTest),
13}
14
15impl Test {
16    pub fn seed(&self) -> Option<Seed> {
17        match self {
18            Test::File(_) => None,
19            Test::Rng(t) => Some(t.seed),
20        }
21    }
22}
23
24pub struct FileTest {
25    pub path: PathBuf,
26}
27
28impl FileTest {
29    pub fn read_into(&self, input: &mut Vec<u8>) {
30        std::fs::File::open(&self.path)
31            .unwrap()
32            .read_to_end(input)
33            .unwrap();
34    }
35}
36
37pub struct RngTest {
38    pub seed: Seed,
39}
40
41impl RngTest {
42    #[inline]
43    pub fn rng(&self) -> Rng {
44        Rng::from_seed(self.seed.to_le_bytes())
45    }
46
47    #[inline]
48    pub fn driver(&self, options: &driver::Options) -> driver::Rng<Rng> {
49        let rng = Rng::from_seed(self.seed.to_le_bytes());
50        driver::Rng::new(rng, options)
51    }
52
53    #[inline]
54    pub fn input<'a>(
55        &self,
56        buffer: &'a mut Vec<u8>,
57        cache: &'a mut driver::cache::Cache,
58        options: &'a driver::Options,
59    ) -> cache::Driver<'a, driver::Rng<Rng>> {
60        let driver = self.driver(options);
61        cache::Driver::new(driver, cache, buffer)
62    }
63
64    #[inline]
65    pub fn buffered_input<'a>(
66        &self,
67        buffer: &'a mut Vec<u8>,
68        options: &'a driver::Options,
69    ) -> RngBufferedInput<'a> {
70        let rng = self.rng();
71        let driver = BufferedRng { rng, buffer };
72        let driver = driver::Rng::new(driver, options);
73        RngBufferedInput {
74            driver,
75            slice: vec![],
76        }
77    }
78}
79
80pub struct BufferedRng<'a> {
81    rng: Rng,
82    buffer: &'a mut Vec<u8>,
83}
84
85impl rand::RngCore for BufferedRng<'_> {
86    fn next_u32(&mut self) -> u32 {
87        let mut data = [0; 4];
88        self.fill_bytes(&mut data);
89        u32::from_le_bytes(data)
90    }
91
92    fn next_u64(&mut self) -> u64 {
93        let mut data = [0; 8];
94        self.fill_bytes(&mut data);
95        u64::from_le_bytes(data)
96    }
97
98    fn fill_bytes(&mut self, bytes: &mut [u8]) {
99        self.rng.fill_bytes(bytes);
100        self.buffer.extend_from_slice(bytes);
101    }
102}
103
104pub struct RngBufferedInput<'a> {
105    driver: driver::Rng<BufferedRng<'a>>,
106    slice: Vec<u8>,
107}
108
109impl<'a, Output> Input<Output> for RngBufferedInput<'a> {
110    type Driver = driver::Rng<BufferedRng<'a>>;
111
112    fn with_slice<F: FnMut(&[u8]) -> Output>(&mut self, f: &mut F) -> Output {
113        self.slice.mutate(&mut self.driver);
114        f(&self.slice)
115    }
116
117    fn with_driver<F: FnMut(&mut Self::Driver) -> Output>(&mut self, f: &mut F) -> Output {
118        f(&mut self.driver)
119    }
120}
121
122pub struct ReplayRng<'a> {
123    buffer: &'a [u8],
124}
125
126impl rand::RngCore for ReplayRng<'_> {
127    fn next_u32(&mut self) -> u32 {
128        let mut data = [0; 4];
129        self.fill_bytes(&mut data);
130        u32::from_le_bytes(data)
131    }
132
133    fn next_u64(&mut self) -> u64 {
134        let mut data = [0; 8];
135        self.fill_bytes(&mut data);
136        u64::from_le_bytes(data)
137    }
138
139    fn fill_bytes(&mut self, bytes: &mut [u8]) {
140        let len = self.buffer.len().min(bytes.len());
141        let (copy_from, remaining) = self.buffer.split_at(len);
142        let (copy_to, fill_to) = bytes.split_at_mut(len);
143        copy_to.copy_from_slice(copy_from);
144        fill_to.fill(0);
145        self.buffer = remaining;
146    }
147}
148
149pub struct RngReplayInput<'a> {
150    pub buffer: &'a mut Vec<u8>,
151}
152
153impl bolero_engine::shrink::Input for RngReplayInput<'_> {
154    type Driver<'d>
155        = driver::Rng<ReplayRng<'d>>
156    where
157        Self: 'd;
158
159    #[inline]
160    fn driver(&self, len: usize, options: &driver::Options) -> Self::Driver<'_> {
161        let buffer = &self.buffer[..len];
162        let rng = ReplayRng { buffer };
163        driver::Rng::new(rng, options)
164    }
165}
166
167impl AsRef<Vec<u8>> for RngReplayInput<'_> {
168    #[inline]
169    fn as_ref(&self) -> &Vec<u8> {
170        self.buffer
171    }
172}
173
174impl AsMut<Vec<u8>> for RngReplayInput<'_> {
175    #[inline]
176    fn as_mut(&mut self) -> &mut Vec<u8> {
177        self.buffer
178    }
179}
180
181pub struct ExhastiveInput<'a> {
182    pub buffer: &'a mut Vec<u8>,
183    pub driver: &'a mut driver::exhaustive::Driver,
184}
185
186impl<Output> Input<Output> for ExhastiveInput<'_> {
187    type Driver = driver::exhaustive::Driver;
188
189    fn with_slice<F: FnMut(&[u8]) -> Output>(&mut self, f: &mut F) -> Output {
190        self.buffer.mutate(&mut self.driver);
191        f(self.buffer)
192    }
193
194    fn with_driver<F: FnMut(&mut Self::Driver) -> Output>(&mut self, f: &mut F) -> Output {
195        f(self.driver)
196    }
197}
198
199impl AsRef<Vec<u8>> for ExhastiveInput<'_> {
200    #[inline]
201    fn as_ref(&self) -> &Vec<u8> {
202        self.buffer
203    }
204}
205
206impl AsMut<Vec<u8>> for ExhastiveInput<'_> {
207    #[inline]
208    fn as_mut(&mut self) -> &mut Vec<u8> {
209        self.buffer
210    }
211}