Structured Aware Harness for tinybmp Rust crate

structured.rs

// Harness by https://twitter.com/addisoncrump_vr
#![cfg_attr(fuzzing, no_main)]

use arbitrary::{Arbitrary, Unstructured};
use embedded_graphics_core::geometry::Point;
use libfuzzer_sys::fuzz_target;
use rand::rngs::StdRng;
use rand::{RngCore, SeedableRng};
use std::num::{NonZeroI8, NonZeroU8};
#[cfg(not(fuzzing))]
use std::{
    fs::{read_dir, File},
    io::Read,
};
use tinybmp::Bpp::*;
use tinybmp::{Bpp, RawBmp};

#[allow(non_camel_case_types)]
#[derive(Debug, Copy, Clone, PartialOrd, PartialEq)]
#[repr(u8)]
enum DibType {
    DIB_INFO_HEADER_SIZE = 40,
    DIB_V3_HEADER_SIZE = 56,
    DIB_V4_HEADER_SIZE = 108,
    DIB_V5_HEADER_SIZE = 124,
}

#[derive(Debug, Copy, Clone, PartialOrd, PartialEq)]
#[repr(u8)]
enum CompressionMethod {
    Rgb = 0,
    Bitfields = 3,
}

#[derive(Debug, Clone)]
struct FuzzyBmp {
    dib_type: DibType,
    bpp: Bpp,
    compress: CompressionMethod,
    colours_used: u32,
    width: u32,
    height: i32,
    mask_red: u32,
    mask_green: u32,
    mask_blue: u32,
    mask_alpha: u32,
    // TODO compute at arb time
    colour_table: Box<[u8]>,
    image_data: Box<[u8]>,
}

impl From<FuzzyBmp> for Vec<u8> {
    fn from(bmp: FuzzyBmp) -> Self {
        let mut res = Vec::new();

        // --- HEADER ---

        // header magic
        res.extend_from_slice(b"BM");

        // size; calculated later
        let size_marker = res.len();

        // reserved data
        res.extend_from_slice(&0u16.to_le_bytes());
        res.extend_from_slice(&0u16.to_le_bytes());

        // image data offset; calculated later
        let data_ptr_marker = res.len();

        // --- DIB Header ---

        // length
        res.extend_from_slice(&(bmp.dib_type as u32).to_le_bytes());

        // width, height
        res.extend_from_slice(&bmp.width.to_le_bytes());
        res.extend_from_slice(&bmp.height.to_le_bytes());

        // color planes -- unused
        res.extend_from_slice(&0u16.to_le_bytes());

        // bpp
        res.extend_from_slice(&(bmp.bpp.bits()).to_le_bytes());

        // compression method
        res.extend_from_slice(&(bmp.compress as u32).to_le_bytes());

        // image data len
        res.extend_from_slice(&u32::try_from(bmp.image_data.len()).unwrap().to_le_bytes());

        // pels per meter x + y -- unused
        res.extend_from_slice(&0u32.to_le_bytes());
        res.extend_from_slice(&0u32.to_le_bytes());

        // colours used
        res.extend_from_slice(&bmp.colours_used.to_le_bytes());

        // colours important -- unused
        res.extend_from_slice(&0u32.to_le_bytes());

        // channel masks, if we meet the described conditions
        if bmp.dib_type >= DibType::DIB_V3_HEADER_SIZE
            && bmp.compress == CompressionMethod::Bitfields
        {
            res.extend_from_slice(&bmp.mask_red.to_le_bytes());
            res.extend_from_slice(&bmp.mask_green.to_le_bytes());
            res.extend_from_slice(&bmp.mask_blue.to_le_bytes());
            res.extend_from_slice(&bmp.mask_alpha.to_le_bytes());
        }

        // colour table
        res.extend_from_slice(&bmp.colour_table);

        // insert data pointer marker, with corrections
        res.splice(
            data_ptr_marker..data_ptr_marker,
            u32::try_from(res.len() + 2 * core::mem::size_of::<u32>())
                .unwrap()
                .to_le_bytes(),
        );

        // image data
        res.extend_from_slice(&bmp.image_data);

        // insert total length, with corrections
        res.splice(
            size_marker..size_marker,
            u32::try_from(res.len() + core::mem::size_of::<u32>())
                .unwrap()
                .to_le_bytes(),
        );

        res
    }
}

impl<'a> Arbitrary<'a> for FuzzyBmp {
    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
        let dib_type = match u8::arbitrary(u)? % 4 {
            0 => DibType::DIB_INFO_HEADER_SIZE,
            1 => DibType::DIB_V3_HEADER_SIZE,
            2 => DibType::DIB_V4_HEADER_SIZE,
            3 => DibType::DIB_V5_HEADER_SIZE,
            _ => unreachable!(),
        };
        let bpp = match u8::arbitrary(u)? % 6 {
            0 => Bits1,
            1 => Bits4,
            2 => Bits8,
            3 => Bits16,
            4 => Bits24,
            5 => Bits32,
            _ => unreachable!(),
        };
        let compress = match bool::arbitrary(u)? {
            false => CompressionMethod::Rgb,
            true => CompressionMethod::Bitfields,
        };

        // simply restrict it to 8 bits to keep it within reasonable bounds
        let width: u32 = NonZeroU8::arbitrary(u)?.get().into();
        let height: i32 = NonZeroI8::arbitrary(u)?.get().into();

        let mask_red = u.arbitrary()?;
        let mask_green = u.arbitrary()?;
        let mask_blue = u.arbitrary()?;
        let mask_alpha = u.arbitrary()?;

        let mut colours_used = u8::arbitrary(u)? as u32;
        // bias towards special case
        let colour_table_num_entries = if colours_used % 4 == 0 && bpp.bits() < 16 {
            colours_used = 0;
            (1 << bpp.bits()) as usize
        } else {
            colours_used as usize
        };

        // image data doesn't actually matter; seed and fill
        let seed = u64::arbitrary(u)?;
        // rows must be a multiple of 4 bytes, so we do some wacky corrections here
        let baseline = width * height.unsigned_abs();
        let data_len = (match bpp {
            Bits1 => baseline / 8,
            Bits4 => baseline / 2,
            Bits8 => baseline,
            Bits16 => baseline * 2,
            Bits24 => baseline * 3,
            Bits32 => baseline * 4,
            _ => unreachable!(),
        } as u32
            + (width * 4 - 1))
            * (width * 4)
            / (width * 4);
        let mut rng = StdRng::seed_from_u64(seed);
        let mut image_data = vec![0u8; data_len as usize];
        rng.fill_bytes(image_data.as_mut_slice());
        let image_data = image_data.into_boxed_slice();

        let colour_table = u
            .arbitrary_iter()?
            .flatten()
            .take(colour_table_num_entries * 4)
            .collect::<Vec<_>>()
            .into_boxed_slice();
        if colour_table.len() != colour_table_num_entries * 4 {
            return Err(arbitrary::Error::NotEnoughData);
        }

        Ok(Self {
            dib_type,
            bpp,
            compress,
            colours_used,
            width,
            height,
            mask_red,
            mask_green,
            mask_blue,
            mask_alpha,
            colour_table,
            image_data,
        })
    }

    fn size_hint(_: usize) -> (usize, Option<usize>) {
        (25, Some(65536))
    }
}

fn do_fuzz(bmp: FuzzyBmp) {
    let bmp = Vec::from(bmp);
    if let Ok(bmp) = RawBmp::from_slice(bmp.as_slice()) {
        bmp.pixels().for_each(|p| drop(p));
        for x in 0.. {
            if bmp.pixel(Point::new(x, 0)).is_none() {
                break;
            }
            for y in 1.. {
                if bmp.pixel(Point::new(x, y)).is_none() {
                    break;
                }
            }
        }
    }
}

fuzz_target!(|data: FuzzyBmp| {
    do_fuzz(data);
});

#[cfg(not(fuzzing))]
extern "C" {
    fn rust_fuzzer_test_input(bytes: &[u8]) -> i32;
}

#[cfg(not(fuzzing))]
fn main() {
    let dir = std::env::args().skip(1).next().unwrap();
    let mut content = Vec::new();
    for entry in read_dir(dir).unwrap() {
        let entry = entry.unwrap();
        if !entry.file_type().unwrap().is_file() {
            continue;
        }
        let mut file = File::open(entry.path()).unwrap();
        content.clear();
        content.reserve(file.metadata().unwrap().len() as usize);
        file.read_to_end(&mut content).unwrap();
        unsafe {
            rust_fuzzer_test_input(content.as_slice());
        }
    }
}