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

use crate::TypeGenerator;
use rand_core::RngCore;

pub trait Driver: Sized {
    fn gen<T: TypeGenerator>(&mut self) -> Option<T> {
        T::generate(self)
    }

    fn mode(&self) -> DriverMode;

    fn fill_bytes(&mut self, bytes: &mut [u8]) -> Option<()>;
}

/// Byte exhaustion strategy for the driver
#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Eq, Ord)]
pub enum DriverMode {
    /// When the driver bytes are exhausted, the driver will fail fill input bytes.
    /// This is useful for fuzz engines that want accurate mapping of inputs to coverage.
    Direct,

    /// When the driver bytes are exhausted, the driver will continue to fill input bytes with 0.
    /// This is useful for fuzz engines that want to maximize the amount of time spent fuzzing.
    Forced,
}

#[derive(Debug)]
pub struct FuzzDriver<'a> {
    mode: DriverMode,
    input: &'a [u8],
}

impl<'a> FuzzDriver<'a> {
    pub fn new(input: &'a [u8], mode: Option<DriverMode>) -> Self {
        let randomized = mode.is_none();

        let mut driver = Self {
            input,
            mode: mode.unwrap_or(DriverMode::Direct),
        };

        // randomize the driver mode for better coverage
        if randomized {
            driver.mode = if driver.gen().unwrap_or_default() {
                DriverMode::Forced
            } else {
                DriverMode::Direct
            };
        }

        driver
    }
}

impl<'a> Driver for FuzzDriver<'a> {
    fn mode(&self) -> DriverMode {
        self.mode
    }

    fn fill_bytes(&mut self, bytes: &mut [u8]) -> Option<()> {
        match self.mode {
            DriverMode::Forced => {
                let offset = self.input.len().min(bytes.len());
                let (current, remaining) = self.input.split_at(offset);
                let (bytes_to_fill, bytes_to_zero) = bytes.split_at_mut(offset);
                bytes_to_fill.copy_from_slice(current);
                for byte in bytes_to_zero.iter_mut() {
                    *byte = 0;
                }
                self.input = remaining;
                Some(())
            }
            DriverMode::Direct => {
                if bytes.len() > self.input.len() {
                    return None;
                }
                let (current, remaining) = self.input.split_at(bytes.len());
                bytes.copy_from_slice(current);
                self.input = remaining;
                Some(())
            }
        }
    }
}

#[derive(Debug)]
pub struct DirectRng<R: RngCore>(R);

impl<R: RngCore> DirectRng<R> {
    pub fn new(rng: R) -> Self {
        Self(rng)
    }
}

impl<R: RngCore> Driver for DirectRng<R> {
    fn mode(&self) -> DriverMode {
        DriverMode::Direct
    }

    fn fill_bytes(&mut self, bytes: &mut [u8]) -> Option<()> {
        RngCore::try_fill_bytes(&mut self.0, bytes).ok()
    }
}

#[derive(Debug)]
pub struct ForcedRng<R: RngCore>(R);

impl<R: RngCore> ForcedRng<R> {
    pub fn new(rng: R) -> Self {
        Self(rng)
    }
}

impl<R: RngCore> Driver for ForcedRng<R> {
    fn mode(&self) -> DriverMode {
        DriverMode::Forced
    }

    fn fill_bytes(&mut self, bytes: &mut [u8]) -> Option<()> {
        if RngCore::try_fill_bytes(&mut self.0, bytes).is_err() {
            for byte in bytes.iter_mut() {
                *byte = 0;
            }
        }
        Some(())
    }
}