use arbitrary::Unstructured;
use commonware_utils::from_hex_formatted;
use rand_chacha::ChaCha8Rng;
use rand_core::{RngCore as _, SeedableRng};
use std::{
panic::{catch_unwind, AssertUnwindSafe, UnwindSafe},
time::{Duration, Instant},
};
enum Error {
NoDisplay,
String(String),
}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NoDisplay => write!(f, "<not displayable>"),
Self::String(s) => write!(f, "{}", s),
}
}
}
fn try_catch<T>(f: impl FnOnce() -> T + UnwindSafe) -> Result<T, Error> {
catch_unwind(f).map_err(|e| {
e.downcast_ref::<String>()
.map(|s| Error::String(s.to_string()))
.or_else(|| {
e.downcast_ref::<&'static str>()
.map(|s| Error::String((*s).to_string()))
})
.unwrap_or(Error::NoDisplay)
})
}
#[derive(Copy, Clone)]
struct Branch {
seed: u32,
thread: u32,
size: u32,
}
impl Branch {
const fn new(seed: u64) -> Self {
Self {
seed: (seed >> 32) as u32,
thread: seed as u32,
size: 0,
}
}
fn try_from_hex(s: &str) -> Option<Self> {
let bytes: [u8; 12] = from_hex_formatted(s)?.try_into().ok()?;
let seed = u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
let thread = u32::from_be_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
let size = u32::from_be_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]);
Some(Self { seed, thread, size })
}
fn from_hex(s: &str) -> Self {
Self::try_from_hex(s).expect("invalid MINIFUZZ_BRANCH hex format (expected 24 hex chars)")
}
const fn next(self) -> Self {
Self {
seed: self.seed,
thread: self.thread.wrapping_add(1),
size: self.size.saturating_add(1),
}
}
const fn rng_seed(self) -> u64 {
(self.seed as u64) << 32 | self.thread as u64
}
}
impl std::fmt::Display for Branch {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "0x{:08x}{:08x}{:08x}", self.seed, self.thread, self.size)
}
}
const ENV_VAR: &str = "MINIFUZZ_BRANCH";
fn branch_from_env() -> Option<Branch> {
std::env::var(ENV_VAR)
.ok()
.and_then(|s| Branch::try_from_hex(&s))
}
use std::num::Saturating;
const DIVISOR: usize = 1000;
const ADD_BYTES_COEFFS: [Saturating<usize>; 3] =
[Saturating(100), Saturating(1000), Saturating(8000)];
const MODIFY_PREFIX_NUMERATOR: usize = 100;
const MODIFY_PREFIX_MIN_BYTES: usize = 8;
const COPY_PORTION_NUMERATOR: usize = 50;
const STRATEGY_WEIGHTS: [u32; 6] = [200, 20, 10, 50, 10, 200];
struct Sampler {
rng: ChaCha8Rng,
buf: Vec<u8>,
count: i64,
last_bytes_used: usize,
}
impl Sampler {
fn new_with_buf(branch: Branch, mut buf: Vec<u8>) -> Self {
buf.clear();
let rng = ChaCha8Rng::seed_from_u64(branch.rng_seed());
Self {
rng,
buf,
count: branch.size.into(),
last_bytes_used: 0,
}
}
const fn set_bytes_used(&mut self, used: usize) {
self.last_bytes_used = used;
}
fn new(branch: Branch) -> Self {
Self::new_with_buf(branch, Vec::with_capacity(1 << 16))
}
fn switch(&mut self, branch: Branch) {
let buf = std::mem::take(&mut self.buf);
*self = Self::new_with_buf(branch, buf)
}
fn strategy_add_bytes(&mut self, force: bool) {
const MAX_BUF_SIZE: usize = 1 << 13; if !force && self.buf.len() >= MAX_BUF_SIZE {
return;
}
let size = self.buf.len();
let size = Saturating(size);
let num_bytes =
(ADD_BYTES_COEFFS[0] * size * size + ADD_BYTES_COEFFS[1] * size + ADD_BYTES_COEFFS[2])
.0
/ DIVISOR;
let new_size = if force {
self.buf.len() + num_bytes
} else {
(self.buf.len() + num_bytes).min(MAX_BUF_SIZE)
};
let start = self.buf.len();
self.buf.resize(new_size, 0);
self.rng.fill_bytes(&mut self.buf[start..]);
}
fn strategy_modify_prefix(&mut self) {
if self.buf.is_empty() {
return;
}
let max_by_proportion = self.buf.len() * MODIFY_PREFIX_NUMERATOR / DIVISOR;
let prefix_len = max_by_proportion
.max(MODIFY_PREFIX_MIN_BYTES)
.min(self.buf.len());
self.rng.fill_bytes(&mut self.buf[..prefix_len]);
}
fn strategy_copy_portion(&mut self) {
if self.buf.len() < 2 {
return;
}
let portion_size = (self.buf.len() * COPY_PORTION_NUMERATOR / DIVISOR).max(1);
let portion_size = portion_size.min(self.buf.len() / 2);
let src_start = (self.rng.next_u32() as usize) % (self.buf.len() - portion_size + 1);
let dst_start = (self.rng.next_u32() as usize) % (self.buf.len() - portion_size + 1);
self.buf
.copy_within(src_start..src_start + portion_size, dst_start);
}
fn strategy_clear_non_prefix(&mut self) {
let prefix_len = MODIFY_PREFIX_MIN_BYTES.min(self.buf.len());
if prefix_len >= self.buf.len() {
return;
}
let start = prefix_len + (self.rng.next_u32() as usize) % (self.buf.len() - prefix_len);
let len = ((self.rng.next_u32() as usize) % (self.buf.len() - start)).max(1);
self.buf[start..start + len].fill(0);
}
fn strategy_arithmetic_non_prefix(&mut self) {
let prefix_len = MODIFY_PREFIX_MIN_BYTES.min(self.buf.len());
if prefix_len >= self.buf.len() {
return;
}
let start = prefix_len + (self.rng.next_u32() as usize) % (self.buf.len() - prefix_len);
let len = ((self.rng.next_u32() as usize) % (self.buf.len() - start)).max(1);
let delta = self.rng.next_u32() as u8;
for b in &mut self.buf[start..start + len] {
*b = b.wrapping_add(delta);
}
}
fn pick_strategy(&mut self) -> u32 {
let mut weights = STRATEGY_WEIGHTS;
let unused = self.buf.len().saturating_sub(self.last_bytes_used);
if unused > 64 {
weights[4] = 0; weights[5] = 0;
} else if unused > 16 {
weights[4] /= 4; weights[5] /= 4;
}
if self.buf.len() <= 16 {
weights[0] = weights[0].saturating_mul(2);
}
let total: u32 = weights.iter().sum();
if total == 0 {
return 0; }
let mut choice = self.rng.next_u32() % total;
for (i, &w) in weights.iter().enumerate() {
if choice < w {
return i as u32;
}
choice -= w;
}
0
}
fn next(&mut self) -> Option<&[u8]> {
if self.count < 0 {
return None;
}
self.count -= 1;
if self.buf.is_empty() {
self.strategy_add_bytes(false);
} else {
match self.pick_strategy() {
0 => self.strategy_modify_prefix(),
1 => self.strategy_copy_portion(),
2 => self.strategy_clear_non_prefix(),
3 => self.strategy_arithmetic_non_prefix(),
_ => self.strategy_add_bytes(false),
}
}
Some(self.buf.as_slice())
}
}
enum SearchBound {
Limit(u64),
Time(Duration),
}
pub struct Builder {
search_bound: SearchBound,
min_iterations: u64,
seed: Option<u64>,
reproduce: Option<Branch>,
}
impl Default for Builder {
fn default() -> Self {
Self {
search_bound: SearchBound::Time(Duration::from_secs(10)),
min_iterations: 100,
seed: None,
reproduce: None,
}
}
}
impl Builder {
pub const fn with_seed(self, seed: u64) -> Self {
Self {
seed: Some(seed),
..self
}
}
pub const fn with_search_limit(self, search_limit: u64) -> Self {
let min_iterations = if self.min_iterations > search_limit {
search_limit
} else {
self.min_iterations
};
Self {
search_bound: SearchBound::Limit(search_limit),
min_iterations,
..self
}
}
pub const fn with_search_time(self, duration: Duration) -> Self {
Self {
search_bound: SearchBound::Time(duration),
..self
}
}
pub const fn with_min_iterations(self, min_iterations: u64) -> Self {
Self {
min_iterations,
..self
}
}
pub fn with_reproduce(self, hex: &str) -> Self {
Self {
reproduce: Some(Branch::from_hex(hex)),
..self
}
}
pub fn test(
self,
mut s: impl FnMut(&mut arbitrary::Unstructured<'_>) -> Result<(), arbitrary::Error>,
) {
let mut branch = match (self.reproduce, self.seed) {
(Some(b), _) => b,
(None, Some(seed)) => Branch::new(seed),
(None, None) => branch_from_env().unwrap_or_else(|| Branch::new(rand::random())),
};
let mut sampler = Sampler::new(branch);
let mut tries: u64 = 0;
let deadline = match self.search_bound {
SearchBound::Time(d) => Some(Instant::now() + d),
SearchBound::Limit(_) => None,
};
let limit = match self.search_bound {
SearchBound::Limit(l) => l,
SearchBound::Time(_) => u64::MAX,
};
'search: loop {
while let Some(sample) = sampler.next() {
let sample_len = sample.len();
let result = try_catch(AssertUnwindSafe(|| {
let mut u = Unstructured::new(sample);
let res = s(&mut u);
(res, u.len())
}));
match result {
Err(e) => {
panic!("failure ({ENV_VAR} = {branch}):\n{e}")
}
Ok((Err(arbitrary::Error::NotEnoughData), _)) => {
sampler.strategy_add_bytes(true);
}
Ok((Err(arbitrary::Error::IncorrectFormat), remaining)) => {
sampler.set_bytes_used(sample_len - remaining);
}
Ok((Err(e), _)) => {
panic!("failure ({ENV_VAR} = {branch}) (while sampling):\n{e}")
}
Ok((Ok(()), remaining)) => {
sampler.set_bytes_used(sample_len - remaining);
tries += 1;
let past_min = tries >= self.min_iterations;
let past_limit =
tries >= limit || deadline.is_some_and(|d| Instant::now() >= d);
if past_min && past_limit {
break 'search;
}
}
}
}
branch = branch.next();
sampler.switch(branch);
}
eprintln!("failed to find, final: {}", branch);
}
}
pub fn test(s: impl FnMut(&mut arbitrary::Unstructured<'_>) -> Result<(), arbitrary::Error>) {
Builder::default().test(s)
}
#[cfg(test)]
mod tests {
use arbitrary::Unstructured;
#[derive(Debug)]
enum Plan {
Leaf(u8),
Branch(bool, Box<Self>),
}
impl Plan {
fn generate(u: &mut Unstructured<'_>, depth: usize) -> arbitrary::Result<Self> {
if depth == 0 {
Ok(Self::Leaf(u.arbitrary()?))
} else {
let b: bool = u.arbitrary()?;
let child = Self::generate(u, depth - 1)?;
Ok(Self::Branch(b, Box::new(child)))
}
}
fn follow_path(&self, path: &mut impl Iterator<Item = bool>) -> Option<u8> {
match self {
Self::Leaf(v) => Some(*v),
Self::Branch(b, child) => {
if *b == path.next()? {
child.follow_path(path)
} else {
None
}
}
}
}
}
fn search_haystack(depth: usize) {
super::Builder::default()
.with_search_limit(1_000_000)
.with_seed(0)
.test(|u| {
let plan = Plan::generate(u, depth)?;
let mut path = [true, false].into_iter().cycle();
if let Some(leaf) = plan.follow_path(&mut path) {
assert_ne!(leaf, 77);
}
Ok(())
});
}
#[test]
#[should_panic]
fn search_haystack_depth_0() {
search_haystack(0);
}
#[test]
#[should_panic]
fn search_haystack_depth_1() {
search_haystack(1);
}
#[test]
#[should_panic]
fn search_haystack_depth_2() {
search_haystack(2);
}
#[test]
#[should_panic]
fn search_haystack_depth_4() {
search_haystack(4);
}
#[test]
#[should_panic]
fn search_haystack_depth_6() {
search_haystack(6);
}
#[test]
#[should_panic]
fn search_haystack_depth_8() {
search_haystack(8);
}
#[test]
#[should_panic]
fn search_haystack_depth_10() {
search_haystack(10);
}
#[test]
#[should_panic(expected = "MINIFUZZ_BRANCH = 0x")]
fn reproduce_failure() {
super::Builder::default()
.with_reproduce("0x0000000000000000000000a0")
.test(|u| {
let v: u8 = u.arbitrary()?;
assert_ne!(v, 42);
Ok(())
});
}
#[test]
#[should_panic(expected = "<not displayable>")]
fn panic_non_displayable() {
struct NonDisplayable;
super::Builder::default()
.with_search_limit(1)
.with_seed(0)
.test(|_u| {
std::panic::panic_any(NonDisplayable);
});
}
#[test]
fn search_limit_reduces_min_iterations() {
let mut calls = 0u64;
super::Builder::default()
.with_min_iterations(1000)
.with_search_limit(1)
.with_seed(0)
.test(|_u| {
calls += 1;
Ok(())
});
assert_eq!(calls, 1);
}
#[test]
fn incorrect_format_does_not_count_as_try() {
let mut calls = 0u64;
super::Builder::default()
.with_search_limit(1)
.with_seed(0)
.test(|_u| {
calls += 1;
if calls == 1 {
Err(arbitrary::Error::IncorrectFormat)
} else {
Ok(())
}
});
assert_eq!(calls, 2);
}
#[test]
#[should_panic(expected = "MINIFUZZ_BRANCH = 0x")]
fn min_iterations_overrides_search_time() {
super::Builder::default()
.with_search_time(std::time::Duration::ZERO)
.with_min_iterations(1000)
.with_seed(0)
.test(|u| {
let v: u8 = u.arbitrary()?;
assert_ne!(v, 42);
Ok(())
});
}
}