open-standards/wai/wai-rs/tests/media_decode_bounds.rs

//! The image, audio and video decoders hold a payload to the reader limits of SPEC §2:
//! a picture past the limits, an audio payload that declares or decodes to more samples
//! than the limit, and a video payload past them are refused by name, from what the
//! payload declares wherever it declares it, and one at a limit is not refused for it.
//!
//! - Pictures: `check_dimensions` at and one past each limit; real PNG, JPEG XL and
//!   AVIF payloads one pixel past the side limit; PNG, JPEG and AVIF headers that
//!   declare more pixels than the limit with sides within it. Each decoder decodes only
//!   the format its capability names, and JPEG XL decodes to three channels whatever
//!   the source holds. The writers refuse what a reader may refuse.
//! - Audio: a FLAC stream whose STREAMINFO declares 2^36 - 1 samples is refused from the
//!   declaration (it used to ask the allocator for 256 GiB up front, which aborts the
//!   process wherever the allocator refuses that); the declared count is
//!   taken over every channel; a 24-bit stream decodes to [-1, 1]; both decoders refuse
//!   output past a limit the caller gives, and not at it.
//! - Video: the framing header past each limit; more temporal units than the header
//!   declares; a frame larger than it declares; the whole-output budget, at and one
//!   past it; monochrome and 10-bit 4:4:4 streams, which the decoder used to panic on
//!   or misread.
//!
//! Exercises the codec wrappers, so it requires the `codecs` feature.
#![cfg(feature = "codecs")]

use wai::codecs::audio::{
    flac_decode, flac_decode_within, flac_encode, opus_decode_within, opus_encode,
    AUDIO_DECODED_MAX_SAMPLES,
};
use wai::codecs::image::{
    avif_decode, avif_encode, check_dimensions, jpeg_decode, jpeg_encode, jxl_decode,
    jxl_encode, png_decode, png_encode, IMAGE_MAX_PIXELS, IMAGE_MAX_SIDE,
};
use wai::codecs::video::{
    av1_decode, av1_decode_each, av1_decode_within, av1_encode, av1_header,
    VIDEO_MAX_TEMPORAL_UNITS,
};

fn why<T: std::fmt::Debug, E: std::fmt::Debug>(r: Result<T, E>) -> String {
    format!("{:?}", r.expect_err("the payload must be refused"))
}

fn grey(h: u32, w: u32) -> Vec<u8> {
    (0..h * w * 3).map(|i| (i % 251) as u8).collect()
}

// ---- pictures ------------------------------------------------------------------

#[test]
fn picture_limits_hold_at_and_one_past_each_bound() {
    let side = IMAGE_MAX_SIDE as u64;
    assert_eq!(IMAGE_MAX_PIXELS, 1 << 27);
    for (w, h) in [(side, 1), (1, side), (16_384, 8_192), (side, IMAGE_MAX_PIXELS / side)] {
        check_dimensions(w, h).unwrap_or_else(|e| panic!("{w} × {h} is within the limits: {e:?}"));
    }
    for (w, h) in [(side + 1, 1), (1, side + 1), (16_384, 8_193), (side, IMAGE_MAX_PIXELS / side + 1)] {
        let e = why(check_dimensions(w, h));
        assert!(e.contains("image_too_large:"), "{w} × {h}: {e}");
    }
}

#[test]
fn png_and_jxl_one_pixel_past_the_side_limit_are_refused() {
    // Made with the libraries themselves: the reference writers refuse such a picture.
    let w = IMAGE_MAX_SIDE + 1;
    let rgb = grey(1, w);
    let mut png = Vec::new();
    image::ImageEncoder::write_image(
        image::codecs::png::PngEncoder::new(&mut png), &rgb, w, 1, image::ExtendedColorType::Rgb8,
    )
    .unwrap();
    assert!(why(png_decode(&png)).contains("image_too_large:"));
    let mut enc = jpegxl_rs::encoder_builder().lossless(true).uses_original_profile(true).build().unwrap();
    let jxl = enc.encode::<u8, u8>(&rgb, w, 1).unwrap().data;
    assert!(why(jxl_decode(&jxl)).contains("image_too_large:"));
    // At the side limit both decode.
    let rgb = grey(1, IMAGE_MAX_SIDE);
    let (out, h, w) = png_decode(&png_encode(&rgb, 1, IMAGE_MAX_SIDE).unwrap()).unwrap();
    assert_eq!((out.len(), h, w), (rgb.len(), 1, IMAGE_MAX_SIDE));
    let (out, h, w) = jxl_decode(&jxl_encode(&rgb, 1, IMAGE_MAX_SIDE, None).unwrap()).unwrap();
    assert_eq!((out, h, w), (rgb, 1, IMAGE_MAX_SIDE));
}

/// CRC-32 (ISO 3309), for PNG chunks.
fn crc32(bytes: &[u8]) -> u32 {
    let mut c = 0xFFFF_FFFFu32;
    for &b in bytes {
        c ^= b as u32;
        for _ in 0..8 {
            c = if c & 1 != 0 { 0xEDB8_8320 ^ (c >> 1) } else { c >> 1 };
        }
    }
    !c
}

#[test]
fn png_and_jpeg_headers_past_the_pixel_limit_are_refused_before_decoding() {
    // 16 384 × 8 193: each side within the limit, one row of pixels past it. Only the
    // header says so; the picture data is an 8 × 8's, so a decoder that read it first
    // would fail for another reason.
    let rgb = grey(8, 8);
    let mut png = png_encode(&rgb, 8, 8).unwrap();
    assert_eq!(&png[12..16], b"IHDR");
    png[16..20].copy_from_slice(&16_384u32.to_be_bytes());
    png[20..24].copy_from_slice(&8_193u32.to_be_bytes());
    let crc = crc32(&png[12..29]);
    png[29..33].copy_from_slice(&crc.to_be_bytes());
    assert!(why(png_decode(&png)).contains("image_too_large:"));

    let mut jpg = jpeg_encode(&rgb, 8, 8, 90).unwrap();
    let sof = jpg.windows(2).position(|m| m == [0xFF, 0xC0]).expect("a baseline SOF0");
    jpg[sof + 5..sof + 7].copy_from_slice(&8_193u16.to_be_bytes());
    jpg[sof + 7..sof + 9].copy_from_slice(&16_384u16.to_be_bytes());
    assert!(why(jpeg_decode(&jpg)).contains("image_too_large:"));
}

#[test]
fn avif_declarations_past_the_limits_are_refused_and_a_small_one_decodes() {
    let rgb = grey(16, 16);
    let avif = avif_encode(&rgb, 16, 16, 80).unwrap();
    let (out, h, w) = avif_decode(&avif).unwrap();
    assert_eq!((out.len(), h, w), (rgb.len(), 16, 16));
    let ispe = avif.windows(4).position(|b| b == b"ispe").expect("an ispe property");
    for (w, h) in [(IMAGE_MAX_SIDE + 1, 1u32), (16_384, 8_193)] {
        let mut a = avif.clone();
        a[ispe + 8..ispe + 12].copy_from_slice(&w.to_be_bytes());
        a[ispe + 12..ispe + 16].copy_from_slice(&h.to_be_bytes());
        let e = why(avif_decode(&a));
        assert!(e.contains("image_too_large:"), "{w} × {h}: {e}");
    }
}

#[test]
fn each_image_decoder_decodes_only_its_own_format() {
    let rgb = grey(8, 8);
    let png = png_encode(&rgb, 8, 8).unwrap();
    let jpg = jpeg_encode(&rgb, 8, 8, 90).unwrap();
    assert!(png_decode(&jpg).is_err(), "png_decode read a JPEG");
    assert!(jpeg_decode(&png).is_err(), "jpeg_decode read a PNG");
    assert!(png_decode(&png).is_ok() && jpeg_decode(&jpg).is_ok());
}

#[test]
fn the_writers_refuse_what_a_reader_may_refuse() {
    let w = IMAGE_MAX_SIDE + 1;
    let rgb = grey(1, w);
    assert!(why(png_encode(&rgb, 1, w)).contains("image_too_large:"));
    assert!(why(jpeg_encode(&rgb, 1, w, 90)).contains("image_too_large:"));
    assert!(why(avif_encode(&rgb, 1, w, 80)).contains("image_too_large:"));
    assert!(why(jxl_encode(&rgb, 1, w, None)).contains("image_too_large:"));
    let long = vec![0f32; AUDIO_DECODED_MAX_SAMPLES + 1];
    assert!(why(flac_encode(&long, 48_000)).contains("audio_decoded_too_long:"));
    assert!(why(opus_encode(&long, 48_000, 64_000)).contains("audio_decoded_too_long:"));
    drop(long);
    let frame = grey(1, w);
    assert!(why(av1_encode(&[frame], 1, w, 30, 1, false, 60)).contains("video_too_large:"));
}

// ---- audio ---------------------------------------------------------------------

fn tone(n: usize, amp: f32) -> Vec<f32> {
    (0..n).map(|i| amp * (i as f32 * 0.05).sin()).collect()
}

/// Rewrites a FLAC stream's STREAMINFO: its channel count and its 36-bit total.
fn declare(flac: &[u8], channels: u8, total: u64) -> Vec<u8> {
    let mut f = flac.to_vec();
    assert_eq!(&f[..4], b"fLaC");
    let si = 8; // after the marker and the STREAMINFO block header
    // Byte 12 of STREAMINFO: the low bit of the sample rate, three bits of channels - 1,
    // the high bit of bits per sample - 1. Byte 13: four bits of bps - 1, then the
    // total's top four bits; bytes 14..18 the rest of it.
    f[si + 12] = (f[si + 12] & 0xF1) | ((channels - 1) << 1);
    f[si + 13] = (f[si + 13] & 0xF0) | ((total >> 32) as u8 & 0x0F);
    f[si + 14..si + 18].copy_from_slice(&(total as u32).to_be_bytes());
    f
}

#[test]
fn a_flac_stream_that_declares_too_many_samples_is_refused_from_the_declaration() {
    let flac = flac_encode(&tone(4_096, 0.5), 48_000).unwrap();
    let max = AUDIO_DECODED_MAX_SAMPLES as u64;
    // 2^36 - 1, the largest total STREAMINFO can carry: the reader used to reserve that
    // many samples, 256 GiB, before decoding one. Now it refuses, having allocated
    // nothing for it.
    for (ch, total) in [(1, (1u64 << 36) - 1), (1, max + 1), (2, max / 2 + 1)] {
        let e = why(flac_decode(&declare(&flac, ch, total)));
        assert!(e.contains("audio_decoded_too_long:") && e.contains("declares"), "{ch} × {total}: {e}");
    }
    // At the limit the declaration is not refused (the stream then decodes what it holds).
    let at = declare(&flac, 1, max);
    let (out, _) = flac_decode(&at).unwrap();
    assert_eq!(out.len(), 4_096);
}

#[test]
fn flac_decodes_a_24_bit_stream_to_unit_scale() {
    use flac_bound::{FlacEncoder, WriteWrapper};
    let x = tone(4_800, 0.5);
    let ints: Vec<i32> = x.iter().map(|&s| (s * 8_388_607.0).round() as i32).collect();
    let mut bytes = Vec::new();
    {
        let mut ww = WriteWrapper(&mut bytes);
        let mut enc = FlacEncoder::new().unwrap().channels(1).bits_per_sample(24)
            .sample_rate(48_000).init_write(&mut ww).unwrap();
        enc.process_interleaved(&ints, ints.len() as u32).unwrap();
        enc.finish().unwrap();
    }
    let (out, _) = flac_decode(&bytes).unwrap();
    let peak = out.iter().fold(0f32, |m, &v| m.max(v.abs()));
    assert!((0.49..=0.51).contains(&peak), "a 0.5 tone decoded to a peak of {peak}");
}

#[test]
fn the_audio_decoders_refuse_output_past_the_callers_limit_and_not_at_it() {
    let x = tone(9_600, 0.5);
    let flac = flac_encode(&x, 48_000).unwrap();
    // A stream that declares its total is refused from the declaration; one whose total
    // is 0 (unknown) is refused from the count as samples decode.
    assert!(why(flac_decode_within(&declare(&flac, 1, 9_600), 9_599)).contains("declares"));
    let unknown = declare(&flac, 1, 0);
    assert!(why(flac_decode_within(&unknown, 9_599)).contains("decodes to"));
    assert_eq!(flac_decode_within(&unknown, 9_600).unwrap().0.len(), 9_600);

    let opus = opus_encode(&x, 48_000, 64_000).unwrap();
    let n = opus_decode_within(&opus, usize::MAX).unwrap().0.len();
    assert!(n > 9_000, "decoded {n} samples of 9 600");
    assert_eq!(opus_decode_within(&opus, n).unwrap().0.len(), n);
    assert!(why(opus_decode_within(&opus, n - 1)).contains("audio_decoded_too_long:"));
}

// ---- video ---------------------------------------------------------------------

fn framing(n: u32, h: u32, w: u32, units: &[Vec<u8>]) -> Vec<u8> {
    let mut out = Vec::new();
    for v in [n, h, w, 30, 1] {
        out.extend_from_slice(&v.to_le_bytes());
    }
    for u in units {
        out.extend_from_slice(&(u.len() as u32).to_le_bytes());
        out.extend_from_slice(u);
    }
    out
}

/// The temporal units of an AV1 payload, after its framing header.
fn units_of(payload: &[u8]) -> Vec<Vec<u8>> {
    let mut at = 20;
    let mut out = Vec::new();
    while at < payload.len() {
        let n = u32::from_le_bytes(payload[at..at + 4].try_into().unwrap()) as usize;
        out.push(payload[at + 4..at + 4 + n].to_vec());
        at += 4 + n;
    }
    out
}

#[test]
fn the_av1_framing_header_is_held_to_the_limits() {
    let n = VIDEO_MAX_TEMPORAL_UNITS;
    let side = IMAGE_MAX_SIDE;
    assert!(av1_header(&framing(n, 2_048, side, &[])).is_ok(), "at every limit");
    for (n, h, w) in [(n + 1, 64, 64), (1, 1, side + 1), (1, 8_193, 16_384)] {
        let e = why(av1_header(&framing(n, h, w, &[])));
        assert!(e.contains("video_too_large:"), "{n} × {w}×{h}: {e}");
    }
}

#[test]
fn av1_refuses_more_units_than_declared_and_frames_larger_than_declared() {
    let clip: Vec<Vec<u8>> = (0..3).map(|_| grey(64, 64)).collect();
    let payload = av1_encode(&clip, 64, 64, 30, 1, false, 60).unwrap();
    let units = units_of(&payload);
    assert_eq!(units.len(), 3, "one temporal unit a frame");
    assert_eq!(av1_decode(&payload).unwrap().0.len(), 3);
    let e = why(av1_decode(&framing(2, 64, 64, &units)));
    assert!(e.contains("more than the 2 temporal units"), "{e}");
    // A frame with more pixels than the header declares never reaches a picture: dav1d
    // is held to h × w pixels (the test below). One with as many pixels, but wider than
    // the header says, is decoded and then refused by its shape.
    assert!(av1_decode(&framing(3, 32, 32, &units)).is_err());
    let wide = av1_unit(128, 16, rav1e::color::ChromaSampling::Cs420, 8, [100, 128, 128]);
    let e = why(av1_decode(&framing(1, 32, 64, &[wide])));
    assert!(e.contains("128×16 frame in a payload whose frames are at most 64×32"), "{e}");
}

#[test]
fn the_whole_output_budget_refuses_from_the_header_and_streaming_does_not_need_one() {
    let clip: Vec<Vec<u8>> = (0..4).map(|_| grey(64, 64)).collect();
    let payload = av1_encode(&clip, 64, 64, 30, 1, false, 60).unwrap();
    let need = 4 * 64 * 64 * 3;
    assert_eq!(av1_decode_within(&payload, need).unwrap().0.len(), 4);
    assert!(why(av1_decode_within(&payload, need - 1)).contains("video_decoded_too_large:"));
    let mut seen = 0;
    let hd = av1_decode_each(&payload, |rgb, w, h| {
        assert_eq!((rgb.len(), w, h), (64 * 64 * 3, 64, 64));
        seen += 1;
        Ok(())
    })
    .unwrap();
    assert_eq!((seen, hd.n), (4, 4));
}

#[test]
fn a_sequence_header_larger_than_the_framing_header_is_refused_before_it_allocates() {
    // A 256×256 stream in a framing header that says 16×16: dav1d is held to 16×16
    // pixels, so it refuses the stream's sequence header instead of allocating its
    // frames and leaving the size check to the decoded picture.
    let unit = av1_unit(256, 256, rav1e::color::ChromaSampling::Cs420, 8, [100, 128, 128]);
    let e = why(av1_decode(&framing(1, 16, 16, &[unit.clone()])));
    assert!(!e.contains("256×256 frame"), "the frame was decoded before it was refused: {e}");
    assert!(av1_decode(&framing(1, 256, 256, &[unit])).is_ok());
}

/// One AV1 temporal unit of a `w`×`h` picture in the given layout and depth, every
/// sample of plane p equal to `fill[p]`, made with rav1e.
fn av1_unit(w: usize, h: usize, cs: rav1e::color::ChromaSampling, depth: usize, fill: [u16; 3]) -> Vec<u8> {
    use rav1e::{Config, EncoderConfig, EncoderStatus};
    let mut cfg = EncoderConfig::with_speed_preset(10);
    cfg.width = w;
    cfg.height = h;
    cfg.chroma_sampling = cs;
    cfg.bit_depth = depth;
    cfg.still_picture = true;
    cfg.quantizer = 0;
    cfg.min_quantizer = 0;
    let mut ctx: rav1e::Context<u16> = Config::new().with_encoder_config(cfg).new_context().unwrap();
    let mut f = ctx.new_frame();
    let planes = if cs == rav1e::color::ChromaSampling::Cs400 { 1 } else { 3 };
    for (p, plane) in f.planes.iter_mut().enumerate().take(planes) {
        for v in plane.data_origin_mut().iter_mut() {
            *v = fill[p];
        }
    }
    ctx.send_frame(f).unwrap();
    ctx.flush();
    let mut out = Vec::new();
    loop {
        match ctx.receive_packet() {
            Ok(p) => out.extend_from_slice(&p.data),
            Err(EncoderStatus::Encoded) => continue,
            Err(_) => break,
        }
    }
    out
}

#[test]
fn av1_decodes_monochrome_and_10_bit_4_4_4() {
    use rav1e::color::ChromaSampling::{Cs400, Cs444};
    // Monochrome, 8 bits: Y = 100 everywhere, so every RGB sample is 100.
    let mono = framing(1, 16, 16, &[av1_unit(16, 16, Cs400, 8, [100, 0, 0])]);
    let (frames, ..) = av1_decode(&mono).unwrap();
    assert!(frames[0].iter().all(|&v| v == 100), "monochrome: {:?}", &frames[0][..6]);
    // 10-bit 4:4:4: Y = 400, Cb = Cr = 512 (neutral), so every sample is 400 / 4 = 100.
    let deep = framing(1, 16, 16, &[av1_unit(16, 16, Cs444, 10, [400, 512, 512])]);
    let (frames, ..) = av1_decode(&deep).unwrap();
    assert!(frames[0].iter().all(|&v| v == 100), "10-bit 4:4:4: {:?}", &frames[0][..6]);
}