open-standards/wai/wai-vector/tests/held_sets.rs

//! §2.7 step 6: a held set is read by digest, at most `length + 1` bytes of it, and
//! hashed before it is parsed.

use wai_vector::refuse::Code;
use wai_vector::sha256::sha256;
use wai_vector::vectors::{corpus_store, find, sets};
use wai_vector::wvg1::{self, Resolver};

/// A resolver that serves one stream for one digest and counts the bytes it hands out.
struct Counting {
    sha: [u8; 32],
    data: Vec<u8>,
    at: usize,
    served: usize,
}

impl Resolver for Counting {
    fn open(&mut self, sha: &[u8; 32]) -> bool {
        self.at = 0;
        *sha == self.sha
    }
    fn read(&mut self, buf: &mut [u8]) -> usize {
        let n = buf.len().min(self.data.len() - self.at);
        buf[..n].copy_from_slice(&self.data[self.at..self.at + n]);
        self.at += n;
        self.served += n;
        n
    }
}

#[test]
fn held_set_one_byte_longer_is_a_mismatch() {
    let v = find("refusals", "s6-held-one-byte-longer").unwrap();
    let e = wvg1::load(&v.bytes, &mut corpus_store()).unwrap_err();
    assert_eq!((e.code, e.step), (Code::OutlineMismatch, "6"));
    let v = find("canonical", "held").unwrap();
    assert!(wvg1::load(&v.bytes, &mut corpus_store()).is_ok());
}

/// A stream far longer than its record's length: the sink reads `length + 1` bytes and
/// stops.
#[test]
fn a_counting_resolver_sees_at_most_length_plus_one_bytes() {
    let b = sets::set_b().write();
    let v = find("canonical", "held").unwrap();
    let d = wvg1::parse(&v.bytes).unwrap();
    assert_eq!(d.sets[0].length as usize, b.len());
    let mut long = b.clone();
    long.extend(std::iter::repeat_n(0xAB, 1 << 20));
    let mut r = Counting { sha: sha256(&b), data: long, at: 0, served: 0 };
    let e = wvg1::load(&v.bytes, &mut r).unwrap_err();
    assert_eq!(e.code, Code::OutlineMismatch);
    assert!(r.served <= b.len() + 1, "the sink read {} bytes of a set of length {}", r.served, b.len());
    // The exact stream: read to its end and accepted.
    let mut r = Counting { sha: sha256(&b), data: b.clone(), at: 0, served: 0 };
    assert!(wvg1::load(&v.bytes, &mut r).is_ok());
    assert_eq!(r.served, b.len());
}