open-standards/joule-code/joule-code-rs/crates/joule-moq/tests/draft22_roundtrip.rs

//! Generated round trips and hostile input for the `draft-ietf-moq-transport-22`
//! codec.
//!
//! A seeded generator (no randomness from outside the test, so a failure
//! reproduces) builds valid control messages of every type, subgroup and fetch
//! streams in both Group Orders, datagrams and `WAI_RECEIPT` values. For each,
//! encode → decode must give the value back, and encode again the same bytes.
//!
//! Then the encodings, and every vector of the corpus, are mutated (bits
//! flipped, bytes cut, inserted, overwritten, lengths swapped for large ones)
//! and fed to every decoder: control messages, whole streams in both orders,
//! the incremental decoder in random pieces, datagrams, property lists, tokens,
//! the `WAI_RECEIPT` value. Each must return, a value or a refusal; a panic
//! fails the test with the input that caused it.

use joule_moq::carriage::WaiReceiptProperty;
use joule_moq::draft22::control::*;
use joule_moq::draft22::data::{encode_fetch_stream, encode_subgroup_stream, FetchObject, Limits};
use joule_moq::draft22::kvp::{prop, KeyValuePair, Properties, SetupOptions};
use joule_moq::draft22::params::{ctx, param, ParamContext, Range, RangeFilter};
use joule_moq::draft22::{
    decode_uni_stream, varint, Datagram, DatagramContent, FetchEntry, FetchHeader, GroupOrder, Location, LocationFilter, ObjectDatagram, ObjectStatus,
    Parameter, Parameters, PublishDoneCode, RangeEnd, RequestErrorCode, StreamHeader, SubgroupHeader, SubgroupId, SubgroupObject, Token, UniEvent,
    UniStreamDecoder,
};
use joule_moq::track::{TrackName, TrackNamespace};

/// xorshift64*: small, seeded, and the same on every host.
struct Rng(u64);

impl Rng {
    fn next(&mut self) -> u64 {
        let mut x = self.0;
        x ^= x >> 12;
        x ^= x << 25;
        x ^= x >> 27;
        self.0 = x;
        x.wrapping_mul(0x2545_f491_4f6c_dd1d)
    }
    fn below(&mut self, n: u64) -> u64 {
        if n == 0 {
            0
        } else {
            self.next() % n
        }
    }
    fn chance(&mut self, num: u64, den: u64) -> bool {
        self.below(den) < num
    }
    /// A value of a random width, so every varint length appears.
    fn wide(&mut self) -> u64 {
        let bits = self.below(65) as u32;
        if bits == 0 {
            0
        } else if bits == 64 {
            self.next()
        } else {
            self.next() & ((1u64 << bits) - 1)
        }
    }
    fn small(&mut self) -> u64 {
        self.below(1000)
    }
    fn bytes(&mut self, max: u64) -> Vec<u8> {
        let n = self.below(max + 1);
        (0..n).map(|_| self.next() as u8).collect()
    }
    fn nonempty(&mut self, max: u64) -> Vec<u8> {
        let mut b = self.bytes(max.max(1) - 1);
        b.push(self.next() as u8);
        b
    }
}

fn namespace(r: &mut Rng) -> TrackNamespace {
    let n = r.below(5);
    TrackNamespace((0..n).map(|_| r.nonempty(8)).collect())
}

fn token(r: &mut Rng) -> Token {
    match r.below(4) {
        0 => Token::Delete { alias: r.wide() },
        1 => Token::Register { alias: r.wide(), token_type: r.wide(), value: r.bytes(12) },
        2 => Token::UseAlias { alias: r.wide() },
        _ => Token::UseValue { token_type: r.wide(), value: r.bytes(12) },
    }
}

fn location_filter(r: &mut Rng) -> LocationFilter {
    let start = Location::new(r.wide() >> 1, r.wide());
    let delta = r.below(1000);
    match r.below(6) {
        0 => LocationFilter::None,
        1 => LocationFilter::RelativeStart { start_group: r.wide() },
        2 => LocationFilter::AbsoluteStart { start },
        3 => LocationFilter::AbsoluteStartGroupEnd { start, end_group_delta: delta },
        4 => LocationFilter::AbsoluteRange { start, end_group_delta: delta, end_object: r.wide() },
        _ => LocationFilter::NextObject,
    }
}

fn ranges(r: &mut Rng, max_value: u64) -> Vec<Range> {
    let n = r.below(4);
    let mut out = Vec::new();
    let mut at = r.below(max_value.min(50) + 1);
    for i in 0..n {
        let open = i + 1 == n && r.chance(1, 3);
        let end = if open { None } else { Some((at + r.below(20)).min(max_value)) };
        out.push(Range { start: at, end });
        match end {
            Some(e) => at = e + r.below(5),
            None => break,
        }
        if at > max_value {
            break;
        }
    }
    out
}

fn range_filter(r: &mut Rng, ty: u64, set_id: u8, may_remove: bool) -> Parameter {
    if may_remove && r.chance(1, 8) {
        return Parameter::RangeFilter { ty, filter: None };
    }
    let property_type = if ty == param::OBJECT_PROPERTY_FILTER || ty == param::TRACK_PROPERTY_FILTER { Some(r.wide() & !1) } else { None };
    let max = if ty == param::PRIORITY_FILTER { 255 } else { u64::MAX >> 1 };
    Parameter::RangeFilter { ty, filter: Some(RangeFilter { set_id, property_type, ranges: ranges(r, max) }) }
}

/// A valid parameter block for `c`, in wire order.
fn params(r: &mut Rng, c: &ParamContext) -> Parameters {
    let mut out = Vec::new();
    for &ty in c.allowed {
        if !r.chance(1, 3) {
            continue;
        }
        let p = match ty {
            param::OBJECT_DELIVERY_TIMEOUT => Parameter::ObjectDeliveryTimeout(r.wide()),
            param::AUTHORIZATION_TOKEN => {
                for _ in 0..r.below(2) {
                    out.push(Parameter::AuthorizationToken(token(r)));
                }
                Parameter::AuthorizationToken(token(r))
            }
            param::RENDEZVOUS_TIMEOUT => Parameter::RendezvousTimeout(r.wide()),
            param::SUBGROUP_DELIVERY_TIMEOUT => Parameter::SubgroupDeliveryTimeout(r.wide()),
            param::EXPIRES => Parameter::Expires(r.wide()),
            param::LARGEST_OBJECT => Parameter::LargestObject(Location::new(r.wide(), r.wide())),
            param::FILL_TIMEOUT => Parameter::FillTimeout(r.wide()),
            param::FORWARD => Parameter::Forward(r.chance(1, 2)),
            param::SUBSCRIBER_PRIORITY => Parameter::SubscriberPriority(r.next() as u8),
            param::LOCATION_FILTER => Parameter::LocationFilter(location_filter(r)),
            param::GROUP_ORDER => Parameter::GroupOrder(if r.chance(1, 2) { GroupOrder::Ascending } else { GroupOrder::Descending }),
            param::FILL_PARAMETERS => Parameter::FillParameters(params(r, &ctx::FILL_PARAMETERS_INNER)),
            t if param::is_range_filter(t) => {
                // Up to three, each with a SetID of its own.
                let n = 1 + r.below(3);
                for s in 1..n {
                    out.push(range_filter(r, t, s as u8, false));
                }
                // A removal (Length 0) has no SetID, so only one may be in a block.
                range_filter(r, t, 0, true)
            }
            param::NEW_GROUP_REQUEST => Parameter::NewGroupRequest(r.wide()),
            param::TRACK_NAMESPACE_PREFIX => Parameter::TrackNamespacePrefix(namespace(r)),
            param::INCLUDE_PROPERTIES => Parameter::IncludeProperties(r.chance(1, 2)),
            param::RELAY_PATH => Parameter::RelayPath((0..1 + r.below(4)).map(|_| [r.next() as u8; 32]).collect()),
            other => panic!("no generator for {other:#x}"),
        };
        out.push(p);
    }
    out.sort_by_key(|p| p.type_code());
    Parameters(out)
}

/// Pairs that pass the checks of an object's or a track's properties: types
/// from the application ranges and the timeouts, an Immutable Properties pair
/// holding more of the same.
fn safe_pairs(r: &mut Rng, nest: bool) -> Vec<KeyValuePair> {
    let pool = [0x02u64, 0x06, 0x78, 0x79, 0x7e, 0x7f, 0x3800, 0x3a40, 0x3a41, 0x3a45, 0x3fff];
    let mut out: Vec<KeyValuePair> = (0..r.below(4))
        .map(|_| {
            let ty = pool[r.below(pool.len() as u64) as usize];
            if ty.is_multiple_of(2) {
                KeyValuePair::varint(ty, r.wide())
            } else {
                KeyValuePair::bytes(ty, r.bytes(16))
            }
        })
        .collect();
    if nest && r.chance(1, 3) {
        let inner = safe_pairs(r, false);
        let mut b = Vec::new();
        joule_moq::draft22::kvp::encode_pairs(&inner, &mut b).unwrap();
        out.push(KeyValuePair::bytes(prop::IMMUTABLE_PROPERTIES, b));
    }
    out.sort_by_key(|k| k.ty);
    out
}

fn properties(r: &mut Rng) -> Properties {
    Properties(safe_pairs(r, true))
}

fn setup(r: &mut Rng) -> Setup {
    let mut o = SetupOptions::new();
    for &(ty, bytes) in &[(0x01u64, true), (0x04, false), (0x05, true), (0x06, false), (0x07, true), (0x08, false), (0x9d, true), (0x3a46, false)] {
        if r.chance(1, 2) {
            o.push(if bytes { KeyValuePair::bytes(ty, r.bytes(20)) } else { KeyValuePair::varint(ty, r.wide()) });
        }
    }
    for _ in 0..r.below(3) {
        o.push(KeyValuePair::bytes(0x03, token(r).to_bytes()));
    }
    o.0.sort_by_key(|p| p.ty);
    Setup { options: o }
}

fn track_request(r: &mut Rng, c: &ParamContext) -> TrackRequest {
    TrackRequest { request_id: r.wide(), namespace: namespace(r), name: TrackName(r.bytes(8)), params: params(r, c) }
}

fn ns_request(r: &mut Rng, c: &ParamContext) -> NamespaceRequest {
    NamespaceRequest { request_id: r.wide(), prefix: namespace(r), params: params(r, c) }
}

fn control(r: &mut Rng, which: u64) -> ControlMessage {
    match which % 19 {
        0 => ControlMessage::Setup(setup(r)),
        1 => ControlMessage::Goaway(Goaway { new_session_uri: r.bytes(40), timeout_ms: r.wide() }),
        2 => ControlMessage::Subscribe(track_request(r, &ctx::SUBSCRIBE)),
        3 => ControlMessage::SubscribeOk(SubscribeOk { track_alias: r.wide(), params: params(r, &ctx::SUBSCRIBE_OK), track_properties: properties(r) }),
        4 => ControlMessage::PublishStateNotify(params(r, &ctx::PUBLISH_STATE_NOTIFY)),
        5 => ControlMessage::Publish(Publish {
            request_id: r.wide(),
            namespace: namespace(r),
            name: TrackName(r.bytes(8)),
            track_alias: r.wide(),
            params: params(r, &ctx::PUBLISH),
            track_properties: properties(r),
        }),
        6 => ControlMessage::PublishDone(PublishDone { status: PublishDoneCode(r.wide()), stream_count: r.wide(), reason: r.bytes(30) }),
        7 => ControlMessage::Fetch(track_request(r, &ctx::FETCH)),
        8 => ControlMessage::FetchOk(FetchOk { end_of_track: r.chance(1, 2), end_location: Location::new(r.wide(), r.wide()), params: Parameters::new(), track_properties: properties(r) }),
        9 => ControlMessage::TrackStatus(track_request(r, &ctx::TRACK_STATUS)),
        10 => ControlMessage::PublishNamespace(ns_request(r, &ctx::PUBLISH_NAMESPACE)),
        11 => ControlMessage::SubscribeNamespace(ns_request(r, &ctx::SUBSCRIBE_NAMESPACE)),
        12 => ControlMessage::SubscribeTracks(ns_request(r, &ctx::SUBSCRIBE_TRACKS)),
        13 => ControlMessage::Namespace(namespace(r)),
        14 => ControlMessage::NamespaceDone(namespace(r)),
        15 => ControlMessage::PublishSkipped(PublishSkipped { suffix: namespace(r), name: TrackName(r.bytes(8)) }),
        16 => ControlMessage::RequestUpdate(RequestUpdate { request_id: r.wide(), params: params(r, &ctx::REQUEST_UPDATE) }),
        17 => ControlMessage::RequestOk(RequestOk { params: params(r, &ctx::REQUEST_OK), track_properties: properties(r) }),
        _ => {
            let redirect = r.chance(1, 3);
            ControlMessage::RequestError(RequestError {
                code: if redirect { RequestErrorCode::REDIRECT } else { RequestErrorCode(r.wide() & !0x34 | (r.below(2) * 0x100)) },
                retry_interval: r.wide(),
                reason: r.bytes(30),
                redirect: redirect.then(|| {
                    let ns = namespace(r);
                    Redirect { connect_uri: r.bytes(30), namespace: ns, name: TrackName(r.bytes(8)) }
                }),
            })
        }
    }
}

fn subgroup(r: &mut Rng) -> (SubgroupHeader, Vec<SubgroupObject>) {
    let h = SubgroupHeader {
        track_alias: r.wide(),
        group_id: r.wide(),
        subgroup_id: match r.below(3) {
            0 => SubgroupId::Zero,
            1 => SubgroupId::FirstObjectId,
            _ => SubgroupId::Explicit(r.wide()),
        },
        publisher_priority: r.chance(1, 2).then(|| r.next() as u8),
        properties: r.chance(1, 2),
        end_of_group: r.chance(1, 2),
        first_object: r.chance(1, 2),
    };
    let mut id = r.small();
    let objs = (0..r.below(6))
        .map(|i| {
            if i > 0 {
                id += 1 + r.below(4);
            }
            let status = match r.below(5) {
                0 => ObjectStatus::EndOfGroup,
                1 => ObjectStatus::EndOfTrack,
                _ => ObjectStatus::Normal,
            };
            let normal = status == ObjectStatus::Normal;
            // Object properties: no Prior Group or Object ID Gap, so no gap can pass its ID.
            let properties = if h.properties && normal { properties(r) } else { Properties::new() };
            let payload = if normal { r.bytes(24) } else { vec![] };
            SubgroupObject { object_id: id, properties, status, payload }
        })
        .collect();
    (h, objs)
}

fn fetch(r: &mut Rng, order: GroupOrder) -> Vec<FetchEntry> {
    let mut g: u64 = match order {
        GroupOrder::Ascending => r.small(),
        GroupOrder::Descending => 10_000 + r.small(),
    };
    let mut o: u64 = r.small();
    let mut first = true;
    let mut out = Vec::new();
    for _ in 0..r.below(8) {
        let move_group = first || r.chance(1, 3);
        if move_group && !first {
            let step = 1 + r.below(5);
            g = match order {
                GroupOrder::Ascending => g + step,
                GroupOrder::Descending => g - step,
            };
            o = r.small();
        } else if !first {
            o += 1 + r.below(3);
        }
        if move_group && r.chance(1, 4) {
            out.push(FetchEntry::EndOfRange { kind: [RangeEnd::NonExistent, RangeEnd::Unknown, RangeEnd::TimedOut][r.below(3) as usize], location: Location::new(g, o) });
        } else {
            out.push(FetchEntry::Object(FetchObject {
                group_id: g,
                object_id: o,
                subgroup_id: r.chance(4, 5).then(|| r.below(4)),
                publisher_priority: r.below(3) as u8,
                properties: r.chance(1, 3).then(|| properties(r)),
                payload: r.bytes(16),
            }));
        }
        first = false;
    }
    out
}

fn datagram(r: &mut Rng) -> ObjectDatagram {
    let status = r.chance(1, 3);
    let st = [ObjectStatus::Normal, ObjectStatus::EndOfGroup, ObjectStatus::EndOfTrack][r.below(3) as usize];
    let with_props = r.chance(1, 2) && (!status || st == ObjectStatus::Normal);
    let mut props = properties(r);
    if props.is_empty() {
        props.push(KeyValuePair::varint(0x3a40, r.wide()));
    }
    ObjectDatagram {
        track_alias: r.wide(),
        group_id: r.wide(),
        object_id: if r.chance(1, 3) { 0 } else { r.wide() },
        publisher_priority: r.chance(1, 2).then(|| r.next() as u8),
        properties: with_props.then_some(props),
        end_of_group: !status && r.chance(1, 2),
        content: if status { DatagramContent::Status(st) } else { DatagramContent::Payload(r.bytes(24)) },
    }
}

/// Every decoder, on bytes that may be anything: each must return.
fn feed_everything(b: &[u8], r: &mut Rng) {
    let _ = varint::get(b, "x");
    let _ = ControlMessage::decode_exact(b);
    let _ = decode_uni_stream(b, GroupOrder::Ascending);
    let _ = decode_uni_stream(b, GroupOrder::Descending);
    let _ = Datagram::decode(b);
    let _ = Properties::decode(b, "x").map(|p| (p.check_object(0, 0), p.check_track(), p.immutable()));
    let _ = Token::decode(b);
    let _ = WaiReceiptProperty::from_bytes(b).map(|p| p.check());
    let mut d = UniStreamDecoder::new(GroupOrder::Ascending, Limits::default());
    let mut at = 0;
    'feed: while at < b.len() {
        let n = (1 + r.below(9) as usize).min(b.len() - at);
        d.push(&b[at..at + n]);
        at += n;
        loop {
            match d.next() {
                Ok(Some(_)) => {}
                Ok(None) => break,
                Err(_) => break 'feed,
            }
        }
    }
    let _ = d.finish();
}

fn mutate(b: &[u8], r: &mut Rng) -> Vec<u8> {
    let mut m = b.to_vec();
    for _ in 0..1 + r.below(3) {
        if m.is_empty() {
            m.push(r.next() as u8);
            continue;
        }
        let i = r.below(m.len() as u64) as usize;
        match r.below(6) {
            0 => m[i] ^= 1 << r.below(8),
            1 => m.truncate(i),
            2 => m.insert(i, r.next() as u8),
            3 => m[i] = r.next() as u8,
            4 => {
                // A length or a count swapped for a huge one.
                let big = varint::to_vec(u64::MAX - r.below(3));
                m.splice(i..i + 1, big);
            }
            _ => m[i] = [0x00, 0x7f, 0x80, 0xbf, 0xc0, 0xfe, 0xff][r.below(7) as usize],
        }
    }
    m
}

#[test]
fn encode_decode_encode_is_the_identity() {
    let mut r = Rng(0x6d6f_7174_2d32_3201);
    for i in 0..4000u64 {
        // Control messages, each type in turn.
        let m = control(&mut r, i);
        let b = m.encode().unwrap_or_else(|e| panic!("{} did not encode: {e}\n{m:?}", m.name()));
        let back = ControlMessage::decode_exact(&b).unwrap_or_else(|e| panic!("{} did not decode: {e}\n{}", m.name(), json_hex(&b)));
        assert_eq!(back, m, "{}", json_hex(&b));
        assert_eq!(back.encode().unwrap(), b);

        // A subgroup stream.
        let (h, objs) = subgroup(&mut r);
        let b = encode_subgroup_stream(&h, &objs).unwrap();
        let ev = decode_uni_stream(&b, GroupOrder::Ascending).unwrap_or_else(|e| panic!("subgroup: {e}\n{}", json_hex(&b)));
        let mut want = vec![UniEvent::Header(StreamHeader::Subgroup(h))];
        want.extend(objs.iter().cloned().map(UniEvent::SubgroupObject));
        assert_eq!(ev, want);
        assert_eq!(encode_subgroup_stream(&h, &objs).unwrap(), b);

        // A fetch stream, in each order.
        for order in [GroupOrder::Ascending, GroupOrder::Descending] {
            let entries = fetch(&mut r, order);
            let fh = FetchHeader { request_id: r.wide() };
            let b = encode_fetch_stream(&fh, order, &entries).unwrap_or_else(|e| panic!("fetch did not encode: {e}\n{entries:?}"));
            let ev = decode_uni_stream(&b, order).unwrap_or_else(|e| panic!("fetch: {e}\n{}", json_hex(&b)));
            let mut want = vec![UniEvent::Header(StreamHeader::Fetch(fh))];
            want.extend(entries.iter().cloned().map(UniEvent::FetchEntry));
            assert_eq!(ev, want);
        }

        // A datagram.
        let d = datagram(&mut r);
        let b = d.encode().unwrap_or_else(|e| panic!("datagram did not encode: {e}\n{d:?}"));
        assert_eq!(Datagram::decode(&b).unwrap(), Datagram::Object(d.clone()));

        // A WAI_RECEIPT value.
        let mut ch = [0u8; 32];
        ch.iter_mut().for_each(|x| *x = r.next() as u8);
        let p = WaiReceiptProperty { content_hash: ch, joules_micro: r.wide(), acquisition: String::from_utf8_lossy(&r.bytes(20)).into_owned(), acquisition_evidence: r.bytes(20), group_receipt: ch };
        assert_eq!(WaiReceiptProperty::from_bytes(&p.to_bytes()).unwrap(), p);

        // Every varint length.
        let v = r.wide();
        let len = varint::encoded_len(v) + r.below((10 - varint::encoded_len(v)) as u64) as usize;
        let mut vb = Vec::new();
        varint::put_with_len(&mut vb, v, len.min(9)).unwrap();
        assert_eq!(varint::get(&vb, "v").unwrap(), (v, vb.len()));
    }
}

fn json_hex(b: &[u8]) -> String {
    joule_moq::draft22::json::hex(b)
}

#[test]
fn no_input_panics_any_decoder() {
    let mut r = Rng(0x686f_7374_696c_6501);
    let corpus: serde_json::Value =
        serde_json::from_str(&std::fs::read_to_string(concat!(env!("CARGO_MANIFEST_DIR"), "/tests/vectors/draft-22/vectors.json")).unwrap()).unwrap();
    let mut seeds: Vec<Vec<u8>> = corpus["vectors"]
        .as_array()
        .unwrap()
        .iter()
        .map(|v| {
            let h = v["hex"].as_str().unwrap();
            (0..h.len()).step_by(2).map(|i| u8::from_str_radix(&h[i..i + 2], 16).unwrap()).collect()
        })
        .collect();
    for i in 0..300u64 {
        seeds.push(control(&mut r, i).encode().unwrap());
        let (h, objs) = subgroup(&mut r);
        seeds.push(encode_subgroup_stream(&h, &objs).unwrap());
        seeds.push(datagram(&mut r).encode().unwrap());
    }
    assert!(!seeds.is_empty());
    let mut fed = 0u64;
    for round in 0..25 {
        for s in &seeds {
            let m = if round == 0 { s.clone() } else { mutate(s, &mut r) };
            let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                let mut r2 = Rng(m.len() as u64 + 1);
                feed_everything(&m, &mut r2);
            }));
            assert!(caught.is_ok(), "a decoder panicked on {}", json_hex(&m));
            fed += 1;
        }
        // Bytes with no structure at all.
        let noise = r.bytes(64);
        feed_everything(&noise, &mut r);
    }
    assert!(fed > 0);
}