jxl_rs_fuzz.rs

jxl_rs_fuzz.rs

use afl::fuzz;
use jxl::api::{
    Endianness, JxlColorType, JxlDataFormat, JxlDecoder, JxlDecoderOptions, JxlPixelFormat,
    ProcessingResult,
};
use jxl::headers::extra_channels::ExtraChannel;
use jxl::image::JxlOutputBuffer;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum WeaverError {
    #[error("Data is not a JPEG XL image")]
    InvalidJxl,
    #[error("JPEG XL decoder failed: {0}")]
    FailedToDecodeJxl(String),
    #[error("Failed to allocate memory with size {0}")]
    FailedToAllocateMemory(u64),
    #[error("Pixel format is not supported: {0}")]
    PixelFormatIsNotSupported(String),
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub(crate) struct BitDepth(u8);

impl BitDepth {
    pub(crate) fn new(bits: u32) -> Result<Self, WeaverError> {
        let bits = u8::try_from(bits).map_err(|_| {
            WeaverError::PixelFormatIsNotSupported(format!("{bits}-bit JPEG XL samples"))
        })?;
        if !(1..=16).contains(&bits) {
            return Err(WeaverError::PixelFormatIsNotSupported(format!(
                "{bits}-bit JPEG XL samples"
            )));
        }
        Ok(Self(bits))
    }

    pub(crate) fn bits(self) -> u8 {
        self.0
    }
}

pub(crate) fn check_image_size_overflow(width: u64, height: u64, chan: u64, t_size: isize) -> bool {
    let Ok(w) = isize::try_from(width) else {
        return true;
    };
    let Ok(h) = isize::try_from(height) else {
        return true;
    };
    let Ok(n) = isize::try_from(chan) else {
        return true;
    };
    let Some(stride) = w.checked_mul(n) else {
        return true;
    };

    // stride * (height - 1) + width * N
    let Some(h_minus_1) = h.checked_sub(1) else {
        return true;
    };
    let Some(lhs) = stride.checked_mul(h_minus_1) else {
        return true;
    };
    lhs.checked_add(stride)
        .and_then(|x| x.checked_mul(t_size))
        .is_none()
}

pub(crate) struct DecodedJxlPacket<T> {
    pub(crate) data: Vec<T>,
    pub(crate) width: usize,
    pub(crate) height: usize,
    pub(crate) icc: Option<Vec<u8>>,
    pub(crate) bit_depth: BitDepth,
    pub(crate) has_real_alpha: bool,
}

fn decode_error(error: impl std::fmt::Display) -> WeaverError {
    WeaverError::FailedToDecodeJxl(error.to_string())
}

pub(crate) enum PackedJxl {
    Regular(DecodedJxlPacket<u8>),
    HighBitDepth(DecodedJxlPacket<u16>),
}

pub(crate) fn decode_packed_jxl(data: &[u8]) -> Result<PackedJxl, WeaverError> {
    let mut input = data;
    let mut decoder = match JxlDecoder::new(JxlDecoderOptions::default())
        .process(&mut input)
        .map_err(decode_error)?
    {
        ProcessingResult::Complete { result } => result,
        ProcessingResult::NeedsMoreInput { .. } => {
            return Err(WeaverError::FailedToDecodeJxl(
                "truncated before image metadata".into(),
            ));
        }
    };

    let info = decoder.basic_info();

    let (width, height) = info.size;

    if width > 10000 || height > 10000 {
        return Err(WeaverError::FailedToAllocateMemory(isize::MAX as u64));
    }

    let source_bit_depth = info.bit_depth.bits_per_sample();
    // Android's bitmap paths consume full-range 8- or 16-bit RGBA. Asking
    // jxl-rs to normalize here also handles uncommon JXL depths (for example
    // 5, 9, or 14 bits) without teaching every downstream converter about them.
    let bit_depth = BitDepth::new(if source_bit_depth <= 8 { 8 } else { 16 })?;
    if check_image_size_overflow(width as u64, height as u64, 4, 2) {
        return Err(WeaverError::FailedToAllocateMemory(isize::MAX as u64));
    }

    let has_real_alpha = info
        .extra_channels
        .iter()
        .any(|channel| channel.ec_type == ExtraChannel::Alpha);
    let extra_channels = info.extra_channels.len();

    let format = if bit_depth.bits() <= 8 {
        JxlDataFormat::U8 { bit_depth: 8 }
    } else {
        JxlDataFormat::U16 {
            endianness: Endianness::native(),
            bit_depth: 16,
        }
    };
    decoder.set_pixel_format(JxlPixelFormat {
        color_type: JxlColorType::Rgba,
        color_data_format: Some(format),
        // Alpha is requested interleaved through Rgba; all other extra channels
        // are deliberately ignored by this still-image Android API.
        extra_channel_format: vec![None; extra_channels],
    });
    let icc = decoder
        .output_color_profile()
        .try_as_icc()
        .map(|profile| profile.into_owned());

    let decoder = match decoder.process(&mut input).map_err(decode_error)? {
        ProcessingResult::Complete { result } => result,
        ProcessingResult::NeedsMoreInput { .. } => {
            return Err(WeaverError::FailedToDecodeJxl(
                "truncated before frame metadata".into(),
            ));
        }
    };

    let samples = width
        .checked_mul(height)
        .and_then(|value| value.checked_mul(4))
        .ok_or(WeaverError::FailedToAllocateMemory(isize::MAX as u64))?;

    if width > 16600 || height > 16600 {
        return Err(WeaverError::FailedToAllocateMemory(isize::MAX as u64));
    }

    if bit_depth.bits() <= 8 {
        let mut pixels = Vec::new();
        pixels.try_reserve_exact(samples).map_err(decode_error)?;
        pixels.resize(samples, 0);
        let stride = width * 4;
        let mut output = [JxlOutputBuffer::new(&mut pixels, height, stride)];
        match decoder
            .process(&mut input, &mut output)
            .map_err(decode_error)?
        {
            ProcessingResult::Complete { .. } => {}
            ProcessingResult::NeedsMoreInput { .. } => {
                return Err(WeaverError::FailedToDecodeJxl(
                    "truncated while decoding pixels".into(),
                ));
            }
        }
        Ok(PackedJxl::Regular(DecodedJxlPacket {
            data: pixels,
            width,
            height,
            icc,
            bit_depth,
            has_real_alpha,
        }))
    } else {
        let mut pixels: Vec<u16> = Vec::new();
        pixels.try_reserve_exact(samples).map_err(decode_error)?;
        pixels.resize(samples, 0);
        let stride = width * 4 * size_of::<u16>();
        let mut output = [JxlOutputBuffer::new(
            bytemuck::cast_slice_mut(&mut pixels),
            height,
            stride,
        )];
        match decoder
            .process(&mut input, &mut output)
            .map_err(decode_error)?
        {
            ProcessingResult::Complete { .. } => {}
            ProcessingResult::NeedsMoreInput { .. } => {
                return Err(WeaverError::FailedToDecodeJxl(
                    "truncated while decoding pixels".into(),
                ));
            }
        }
        Ok(PackedJxl::HighBitDepth(DecodedJxlPacket {
            data: pixels,
            width,
            height,
            icc,
            bit_depth,
            has_real_alpha,
        }))
    }
}

#[macro_use]
extern crate afl;

fn main() {
    fuzz!(|data: &[u8]| {
        _ = decode_packed_jxl(data);
    });
}