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

//! `wai.text.zstd` and `wai.text.xz` decode within the reader limits of SPEC §2: a
//! payload that decodes to more than the length limit, or whose stream asks its decoder
//! to keep more history than the window limit, is refused by name, and one at either
//! limit decodes.
//!
//! - The length limit, at a small limit the caller gives and at the limit the
//!   specification sets, for a frame that declares its length and for one that does not.
//! - A zstd frame that declares a length past the limit is refused from the declaration:
//!   the same frame cut off after its header is refused the same way.
//! - The window limit, on streams made here that declare one at the limit and one past
//!   it. Each holds a byte or a line of content: what is refused is the declaration.
//! - The writers refuse what the readers may refuse.
//!
//! The memory the output buffer takes is measured in `text_decode_memory.rs`.
//!
//! Exercises the codec wrappers, so it requires the `codecs` feature.
#![cfg(feature = "codecs")]

use std::io::Write;

use wai::codecs::text::{
    xz_decode, xz_decode_within, xz_encode, zstd_decode, zstd_decode_within, zstd_encode,
    TEXT_DECODED_MAX_LEN, TEXT_WINDOW_MAX_LEN,
};

fn refusal<T: std::fmt::Debug>(r: Result<T, wai::codecs::text::TextError>) -> String {
    r.expect_err("the payload must be refused").0
}

/// `n` bytes that compress well and are not all one value.
fn pattern(n: usize) -> Vec<u8> {
    (0..n).map(|i| (i / 4096) as u8).collect()
}

/// A zstd stream of `n` zero bytes whose frame declares no content size, written in
/// pieces so that making it takes no buffer of `n` bytes.
fn zstd_zeros(n: usize) -> Vec<u8> {
    let piece = vec![0u8; 1 << 20];
    let mut enc = zstd::stream::write::Encoder::new(Vec::new(), 1).unwrap();
    let mut left = n;
    while left > 0 {
        let k = left.min(piece.len());
        enc.write_all(&piece[..k]).unwrap();
        left -= k;
    }
    enc.finish().unwrap()
}

/// The same for xz.
fn xz_zeros(n: usize) -> Vec<u8> {
    let piece = vec![0u8; 1 << 20];
    let mut enc = xz2::write::XzEncoder::new(Vec::new(), 0);
    let mut left = n;
    while left > 0 {
        let k = left.min(piece.len());
        enc.write_all(&piece[..k]).unwrap();
        left -= k;
    }
    enc.finish().unwrap()
}

#[test]
fn a_payload_at_the_callers_limit_decodes_and_one_byte_past_it_is_refused() {
    const LIMIT: usize = 1 << 20;
    let at = pattern(LIMIT);
    let past = pattern(LIMIT + 1);

    // zstd, the frame declaring no content size (a streamed encode).
    assert_eq!(zstd_decode_within(&zstd_encode(&at, 3).unwrap(), LIMIT).unwrap(), at);
    let why = refusal(zstd_decode_within(&zstd_encode(&past, 3).unwrap(), LIMIT));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");

    // zstd, the frame declaring its content size (a one-shot encode).
    let declared_at = zstd::bulk::compress(&at, 3).unwrap();
    let declared_past = zstd::bulk::compress(&past, 3).unwrap();
    assert_eq!(zstd::zstd_safe::get_frame_content_size(&declared_past).ok().flatten(), Some(LIMIT as u64 + 1));
    assert_eq!(zstd_decode_within(&declared_at, LIMIT).unwrap(), at);
    let why = refusal(zstd_decode_within(&declared_past, LIMIT));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");

    // xz.
    assert_eq!(xz_decode_within(&xz_encode(&at, 1).unwrap(), LIMIT).unwrap(), at);
    let why = refusal(xz_decode_within(&xz_encode(&past, 1).unwrap(), LIMIT));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");
}

#[test]
fn a_declared_length_past_the_limit_is_refused_from_the_declaration() {
    const LIMIT: usize = 1 << 16;
    let frame = zstd::bulk::compress(&pattern(LIMIT + 1), 3).unwrap();
    // The frame's first 18 bytes: its header, which is never longer, and at most a few
    // bytes of its first block. Nothing can be decoded from them.
    assert!(frame.len() > 18);
    let header_only = &frame[..18];
    assert_eq!(zstd::zstd_safe::get_frame_content_size(header_only).ok().flatten(), Some(LIMIT as u64 + 1));
    let why = refusal(zstd_decode_within(header_only, LIMIT));
    assert!(why.starts_with("text_decoded_too_long:"), "refused for another reason: {why}");
    // The same cut, within the limit, is the truncated frame it is.
    let why = refusal(zstd_decode_within(header_only, LIMIT + 1));
    assert!(!why.starts_with("text_decoded_too_long:"), "{why}");
}

#[test]
fn the_limit_without_an_argument_is_the_one_the_specification_sets() {
    assert_eq!(TEXT_DECODED_MAX_LEN, 1 << 27);

    let out = zstd_decode(&zstd_zeros(TEXT_DECODED_MAX_LEN)).unwrap();
    assert_eq!(out.len(), TEXT_DECODED_MAX_LEN);
    assert!(out.iter().all(|&b| b == 0));
    drop(out);
    let past = zstd_zeros(TEXT_DECODED_MAX_LEN + 1);
    let why = refusal(zstd_decode(&past));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");
    // A caller cannot raise it.
    let why = refusal(zstd_decode_within(&past, usize::MAX));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");

    let out = xz_decode(&xz_zeros(TEXT_DECODED_MAX_LEN)).unwrap();
    assert_eq!(out.len(), TEXT_DECODED_MAX_LEN);
    drop(out);
    let past = xz_zeros(TEXT_DECODED_MAX_LEN + 1);
    let why = refusal(xz_decode(&past));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");
    let why = refusal(xz_decode_within(&past, usize::MAX));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");
}

/// A zstd frame of one byte of content that declares a window of `2^window_log` bytes
/// (RFC 8878 §3.1.1): the magic number; a frame header descriptor of 0, which says a
/// window descriptor follows and nothing else does; the window descriptor, whose high
/// five bits are `window_log − 10`; and one last block, of type RLE and size 1.
fn zstd_frame_declaring_window(window_log: u32) -> Vec<u8> {
    let mut f = vec![0x28, 0xB5, 0x2F, 0xFD, 0x00, ((window_log - 10) << 3) as u8];
    f.extend_from_slice(&[0x0B, 0x00, 0x00, b'A']);
    f
}

#[test]
fn a_zstd_window_at_the_limit_decodes_and_one_past_it_is_refused() {
    assert_eq!(TEXT_WINDOW_MAX_LEN, 1 << 27);
    assert_eq!(zstd_decode(&zstd_frame_declaring_window(27)).unwrap(), b"A");
    let why = refusal(zstd_decode(&zstd_frame_declaring_window(28)));
    assert!(why.starts_with("text_window_too_large:"), "{why}");
}

/// CRC-32 as the xz format uses it (IEEE 802.3, reflected).
fn crc32(bytes: &[u8]) -> u32 {
    let mut c = !0u32;
    for &b in bytes {
        c ^= b as u32;
        for _ in 0..8 {
            c = (c >> 1) ^ (0xEDB8_8320 & (!(c & 1)).wrapping_add(1));
        }
    }
    !c
}

/// An xz stream of a line of text whose one block declares the LZMA2 dictionary size
/// coded as `code` (the xz file format, §5.3.1: 30 is 2^27 bytes, 31 is 3·2^26). Made by
/// encoding the line, then rewriting the dictionary size in the block header and the
/// header's CRC-32. The stream header is 12 bytes; the block header follows it: its
/// size, the block flags, then the filter (id 0x21, one byte of properties, which is the
/// dictionary size), padding, and the CRC-32 of all of it.
fn xz_stream_declaring_dictionary(code: u8) -> (Vec<u8>, Vec<u8>) {
    let line = b"one line of text, long enough to be worth a block\n".to_vec();
    let mut s = xz_encode(&line, 0).unwrap();
    let start = 12;
    let size = (s[start] as usize + 1) * 4;
    let header = start..start + size;
    let at = (header.start + 2..header.end - 5)
        .find(|&i| s[i] == 0x21 && s[i + 1] == 0x01)
        .expect("the block header names the LZMA2 filter with one byte of properties");
    s[at + 2] = code;
    let crc = crc32(&s[header.start..header.end - 4]);
    s[header.end - 4..header.end].copy_from_slice(&crc.to_le_bytes());
    (s, line)
}

#[test]
fn an_xz_dictionary_at_the_limit_decodes_and_one_past_it_is_refused() {
    let (at, line) = xz_stream_declaring_dictionary(30);
    assert_eq!(xz_decode(&at).unwrap(), line);
    let (past, _) = xz_stream_declaring_dictionary(31);
    let why = refusal(xz_decode(&past));
    assert!(why.starts_with("text_window_too_large:"), "{why}");
}

#[test]
fn the_writers_refuse_what_a_reader_may_refuse() {
    let over = vec![0u8; TEXT_DECODED_MAX_LEN + 1];
    let why = refusal(zstd_encode(&over, 1));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");
    let why = refusal(xz_encode(&over, 0));
    assert!(why.starts_with("text_decoded_too_long:"), "{why}");
}