open-standards/joule-code/joule-code-rs/crates/joule-moq-relay/tests/common/mod.rs

//! What the relay's integration tests share: a relay on loopback, clients over each
//! transport, publishers that sign their groups, and subscribers that read and verify.
//!
//! Every network address is loopback, and every wait here and in the tests has a hard
//! bound, so no loop can outlive its test.

#![allow(dead_code)]

use ed25519_dalek::SigningKey;
use joule_moq::carriage::{object_properties, ObjectExtras, ReceiptBundle};
use joule_moq::draft22::data::{ObjectStatus as WireStatus, SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::kvp::{prop, KeyValuePair, Properties};
use joule_moq::draft22::{ControlMessage, Parameters, StreamHeader, SubscribeOk, UniEvent};
use joule_moq::messages::ObjectHeader;
use joule_moq::receipt::{ObjectReceipt, ReceiptBuilder, ReceiptVerifier, VerifiedGroup};
use joule_moq::{GroupId, ObjectId, ObjectStatus, TrackAlias, TrackName, TrackNamespace};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::transport::mux::MuxConfig;
use joule_moq_relay::transport::{quic, webtransport, websocket, TransportError};
use joule_moq_relay::wire::{ReqRx, ReqTx, StreamError, UniReader};
use joule_moq_relay::{RelayConfig, RelayServer};
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use std::time::Duration;
use tokio::time::{timeout, Instant};

/// The bound on any single wait in a test.
pub const STEP: Duration = Duration::from_secs(20);

pub fn loopback() -> std::net::SocketAddr {
    "127.0.0.1:0".parse().unwrap()
}

/// A relay on loopback ports, the configuration adjusted by `f`.
pub async fn relay(f: impl FnOnce(&mut RelayConfig)) -> RelayServer {
    let mut c = RelayConfig::loopback();
    f(&mut c);
    RelayServer::start(c).await.expect("the relay starts")
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Via {
    Quic,
    WebTransport,
    WebSocket,
    /// Native QUIC offering only this draft's ALPN; the session bridged to -22.
    QuicDraft(joule_moq::draft::Draft),
    /// WebTransport offering only this draft; the session bridged to -22.
    WtDraft(joule_moq::draft::Draft),
}

/// A client of `server` over `via`. A QUIC client keeps its endpoint, returned with it.
pub async fn connect(server: &RelayServer, via: Via, opts: ClientOptions) -> (Client, Option<quinn::Endpoint>) {
    connect_with(server, via, opts, None).await
}

pub async fn connect_with(server: &RelayServer, via: Via, opts: ClientOptions, transport: Option<quinn::TransportConfig>) -> (Client, Option<quinn::Endpoint>) {
    let quic_addr = server.quic_addr.expect("a UDP endpoint");
    let path = server.relay().config().path.clone();
    match via {
        Via::Quic => {
            let ep = quic::client_endpoint(loopback(), vec![server.identity.chain[0].clone()], vec![joule_moq::draft22::ALPN.to_vec()], transport).unwrap();
            let conn = timeout(STEP, quic::connect(&ep, quic_addr, "localhost")).await.expect("in time").expect("a QUIC session");
            let c = Client::connect(Arc::new(conn), opts).await.expect("SETUP over QUIC");
            (c, Some(ep))
        }
        Via::WebTransport => {
            let url = format!("https://127.0.0.1:{}{}", quic_addr.port(), path);
            let conn = timeout(STEP, webtransport::connect(&url, Some(server.identity.leaf_sha256), loopback())).await.expect("in time").expect("a WebTransport session");
            let c = Client::connect(Arc::new(conn), opts).await.expect("SETUP over WebTransport");
            (c, None)
        }
        Via::QuicDraft(d) => {
            let ep = quic::client_endpoint(loopback(), vec![server.identity.chain[0].clone()], vec![d.alpn().to_vec()], transport).unwrap();
            let (conn, got) = timeout(STEP, quic::connect_any(&ep, quic_addr, "localhost")).await.expect("in time").expect("a QUIC session");
            assert_eq!(got, d, "the relay chose the one draft offered");
            let conn = joule_moq_relay::transport::bridge::wrap(Arc::new(conn), d, joule_moq::draft22::Perspective::Client);
            let c = Client::connect(conn, opts).await.expect("SETUP over QUIC");
            (c, Some(ep))
        }
        Via::WtDraft(d) => {
            let url = format!("https://127.0.0.1:{}{}", quic_addr.port(), path);
            let conn = timeout(STEP, webtransport::connect_draft(&url, Some(server.identity.leaf_sha256), loopback(), d)).await.expect("in time").expect("a WebTransport session");
            let conn = joule_moq_relay::transport::bridge::wrap(Arc::new(conn), d, joule_moq::draft22::Perspective::Client);
            let c = Client::connect(conn, opts).await.expect("SETUP over WebTransport");
            (c, None)
        }
        Via::WebSocket => {
            let url = format!("ws://{}{}", server.websocket_addr.expect("a WebSocket listener"), path);
            let conn = timeout(STEP, websocket::connect(&url, MuxConfig::default())).await.expect("in time").expect("a WebSocket session");
            let c = Client::connect(Arc::new(conn), opts).await.expect("SETUP over the WebSocket binding");
            (c, None)
        }
    }
}

pub fn ns(fields: &[&str]) -> TrackNamespace {
    TrackNamespace::from_strs(fields.iter().copied())
}

pub fn name(s: &str) -> TrackName {
    TrackName::from_str(s)
}

/// Track properties: a default publisher priority, and an OBJECT_DELIVERY_TIMEOUT.
pub fn track_props(priority: u8, object_timeout_ms: Option<u64>) -> Properties {
    let mut p = Properties::new();
    p.push(KeyValuePair::varint(prop::DEFAULT_PUBLISHER_PRIORITY, priority as u64));
    if let Some(t) = object_timeout_ms {
        p.push(KeyValuePair::varint(prop::OBJECT_DELIVERY_TIMEOUT, t));
    }
    p
}

/// One group of a media track, one object, and its receipts as group `g` of the receipts
/// track carries them (WAI moq-streaming-format §4.1).
pub fn signed_group(key: &SigningKey, receipt_alias: u64, g: u64, payload: &[u8]) -> (ObjectReceipt, ReceiptBundle) {
    let hash = ObjectReceipt::hash_payload(payload);
    let mut b = ReceiptBuilder::new(TrackAlias(receipt_alias), GroupId(g));
    b.add_labelled_object(hash, 0, "test-publisher", "2026-10-17T00:00:00Z", vec![0]).unwrap();
    let receipt = ObjectReceipt {
        track_alias: TrackAlias(receipt_alias),
        group_id: GroupId(g),
        object_id: ObjectId(0),
        content_hash: hash,
        joules_micro: 0,
        origin: "test-publisher".into(),
        merkle_proof: b.proof_for(0).unwrap(),
        energy_class: Some(vec![0]),
    };
    let group = b.finalize_labelled(key, "test-publisher", "2026-10-17T00:00:01Z", None).unwrap();
    (receipt.clone(), ReceiptBundle { group, objects: vec![receipt] })
}

/// A media track and its receipts track, as one publisher sends them.
#[derive(Clone)]
pub struct Tracks {
    pub ns: TrackNamespace,
    pub media: TrackName,
    pub receipts: TrackName,
    pub media_alias: u64,
    pub receipts_alias: u64,
    pub receipt_alias: u64,
    pub priority: u8,
}

impl Tracks {
    pub fn new(ns: TrackNamespace, media: &str, media_alias: u64, receipt_alias: u64, priority: u8) -> Self {
        Tracks { ns, media: name(media), receipts: name(&format!("{media}_receipts")), media_alias, receipts_alias: media_alias + 1, receipt_alias, priority }
    }
}

/// Publish group `g` of `t`: the receipts first (no later than the group's first object,
/// moq-streaming-format §4.1), then the object, which carries its receipt as a property.
pub async fn publish_group(c: &Client, key: &SigningKey, t: &Tracks, g: u64, payload: &[u8]) -> ObjectReceipt {
    let (receipt, bundle) = signed_group(key, t.receipt_alias, g, payload);
    let rh = SubgroupHeader { track_alias: t.receipts_alias, group_id: g, subgroup_id: SubgroupId::Zero, publisher_priority: Some(t.priority), properties: false, end_of_group: true, first_object: true };
    c.send_subgroup(rh, &[SubgroupObject::normal(0, bundle.to_payload().unwrap())]).await.expect("the receipts");
    let header = ObjectHeader {
        track_alias: TrackAlias(t.media_alias),
        group_id: GroupId(g),
        object_id: ObjectId(0),
        publisher_priority: t.priority,
        status: ObjectStatus::Normal,
        payload_len: payload.len() as u64,
        mrl_layer: 0,
        mrl_total_layers: 0,
        wai_envelope: false,
    };
    let props = object_properties(&header, &ObjectExtras { receipt: Some(receipt.clone()), wai_receipt: None, delivery_timeout_ms: None }).unwrap();
    let mh = SubgroupHeader { track_alias: t.media_alias, group_id: g, subgroup_id: SubgroupId::Zero, publisher_priority: Some(t.priority), properties: true, end_of_group: true, first_object: true };
    c.send_subgroup(mh, &[SubgroupObject { object_id: 0, properties: props, status: WireStatus::Normal, payload: payload.to_vec() }]).await.expect("the object");
    receipt
}

/// PUBLISH a media track and its receipts track (§9.8).
pub async fn publish_tracks(c: &Client, t: &Tracks, object_timeout_ms: Option<u64>) -> Vec<(ReqTx, ReqRx)> {
    let a = c.publish(t.ns.clone(), t.media.clone(), t.media_alias, Parameters::new(), track_props(t.priority, object_timeout_ms)).await.expect("PUBLISH the media track");
    let b = c.publish(t.ns.clone(), t.receipts.clone(), t.receipts_alias, Parameters::new(), track_props(t.priority, None)).await.expect("PUBLISH the receipts track");
    vec![a, b]
}

/// Answer the relay's SUBSCRIBEs for `t`'s tracks (the namespace route, §7.6) until both
/// are Established; the streams are returned to keep the subscriptions open.
pub async fn answer_subscribes(c: &Client, t: &Tracks) -> Vec<(ReqTx, ReqRx)> {
    let mut held = Vec::new();
    let mut seen = (false, false);
    let deadline = Instant::now() + STEP;
    while !(seen.0 && seen.1) {
        let left = deadline.saturating_duration_since(Instant::now());
        let (tx, rx, first) = c.next_request(left).await.expect("the relay's SUBSCRIBE in time");
        let ControlMessage::Subscribe(r) = first else { panic!("the relay sent {}", first.name()) };
        let (alias, props) = if r.name == t.media {
            seen.0 = true;
            (t.media_alias, track_props(t.priority, None))
        } else {
            assert_eq!(r.name, t.receipts, "the relay subscribes only to the tracks asked for");
            seen.1 = true;
            (t.receipts_alias, track_props(t.priority, None))
        };
        tx.send(&ControlMessage::SubscribeOk(SubscribeOk { track_alias: alias, params: Parameters::new(), track_properties: props })).await.unwrap();
        held.push((tx, rx));
    }
    held
}

/// What a subscriber read from one data stream.
#[derive(Debug)]
pub struct Got {
    pub header: SubgroupHeader,
    pub objects: Vec<SubgroupObject>,
    /// The stream was reset, with this code.
    pub reset: Option<u64>,
    pub arrived: Instant,
}

/// Read one subgroup stream to its end, each read bounded.
pub async fn read_stream(mut r: UniReader) -> Option<Got> {
    let arrived = Instant::now();
    let StreamHeader::Subgroup(header) = r.header.clone()? else { return None };
    let mut objects = Vec::new();
    loop {
        match timeout(STEP, r.next()).await {
            Ok(Ok(Some(UniEvent::SubgroupObject(o)))) => objects.push(o),
            Ok(Ok(Some(_))) => return None,
            Ok(Ok(None)) => return Some(Got { header, objects, reset: None, arrived }),
            Ok(Err(StreamError::Transport(TransportError::Reset(code)))) => return Some(Got { header, objects, reset: Some(code), arrived }),
            Ok(Err(e)) => panic!("a data stream: {e}"),
            Err(_) => panic!("a data stream stalled"),
        }
    }
}

/// Read data streams until `done` holds of what was read, or `limit` passes.
pub async fn collect(c: &Client, limit: Duration, mut done: impl FnMut(&[Got]) -> bool) -> Vec<Got> {
    let deadline = Instant::now() + limit;
    let mut out = Vec::new();
    while !done(&out) {
        let left = deadline.saturating_duration_since(Instant::now());
        if left.is_zero() {
            break;
        }
        match c.next_data(left).await {
            Some(r) => {
                if let Some(g) = read_stream(r).await {
                    out.push(g);
                }
            }
            None => break,
        }
    }
    out
}

/// Read every data stream at once, each in a task of its own, as a subscriber does that
/// reads its streams concurrently: until `done` holds of what was read and `quiet` has
/// passed with no new stream or object, or `limit` passes. Each read is bounded.
///
/// Reading one stream at a time in the order they arrive can stall: QUIC's
/// connection-level flow control is shared by a session's streams, so the bytes of the
/// stream being waited on can wait for credit that unread later streams hold.
pub async fn collect_concurrent(c: &Client, limit: Duration, quiet: Duration, mut done: impl FnMut(&[Got]) -> bool) -> Vec<Got> {
    let deadline = Instant::now() + limit;
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Got>();
    let mut out = Vec::new();
    let mut readers = Vec::new();
    let mut last = Instant::now();
    loop {
        let now = Instant::now();
        if now >= deadline || (done(&out) && now.duration_since(last) >= quiet) {
            break;
        }
        let tick = deadline.min(now + quiet).saturating_duration_since(now);
        tokio::select! {
            s = c.next_stream(tick) => {
                if let Some(mut r) = s {
                    last = Instant::now();
                    let tx = tx.clone();
                    readers.push(tokio::spawn(async move {
                        if let Ok(Ok(_)) = timeout(STEP, r.start()).await
                            && let Some(g) = read_stream(r).await
                        {
                            let _ = tx.send(g);
                        }
                    }));
                }
            }
            g = rx.recv() => {
                if let Some(g) = g {
                    last = Instant::now();
                    out.push(g);
                }
            }
        }
    }
    for r in readers {
        r.abort();
    }
    while let Ok(g) = rx.try_recv() {
        out.push(g);
    }
    out
}

/// A subscriber's view of one publisher's tracks: the downstream aliases the relay gave.
pub struct Subscribed {
    pub media_alias: u64,
    pub receipts_alias: u64,
    pub held: Vec<(ReqTx, ReqRx)>,
}

pub async fn subscribe_both(c: &Client, t: &Tracks, media_params: Parameters) -> Subscribed {
    subscribe_both_with(c, t, media_params, Parameters::new()).await
}

pub async fn subscribe_both_with(c: &Client, t: &Tracks, media_params: Parameters, receipts_params: Parameters) -> Subscribed {
    let (m, mtx, mrx) = c.subscribe(t.ns.clone(), t.media.clone(), media_params).await.expect("SUBSCRIBE the media track");
    let (r, rtx, rrx) = c.subscribe(t.ns.clone(), t.receipts.clone(), receipts_params).await.expect("SUBSCRIBE the receipts track");
    Subscribed { media_alias: m.track_alias, receipts_alias: r.track_alias, held: vec![(mtx, mrx), (rtx, rrx)] }
}

/// Every media object of `sub` verified against the bundles of its receipts track, under
/// the publisher's key alone. Returns (objects verified, bundles verified, the verified
/// groups by Group ID, the object receipts by Group ID).
pub fn verify(got: &[Got], sub: &Subscribed, publisher: &SigningKey) -> (usize, usize, BTreeMap<u64, VerifiedGroup>, HashMap<u64, ObjectReceipt>) {
    let verifier = ReceiptVerifier::with_anchors(vec![publisher.verifying_key()]);
    let mut groups = BTreeMap::new();
    let mut receipts = HashMap::new();
    for g in got.iter().filter(|g| g.header.track_alias == sub.receipts_alias && g.reset.is_none()) {
        let b = ReceiptBundle::from_payload(&g.objects[0].payload).expect("a bundle");
        let vg = verifier.verify_group(&b.group).expect("the group receipt verifies under the publisher's key");
        assert_eq!(b.group.group_id, GroupId(g.header.group_id), "the receipt names its MOQT group");
        groups.insert(g.header.group_id, vg);
        receipts.insert(g.header.group_id, b.objects[0].clone());
    }
    let mut objects = 0;
    for g in got.iter().filter(|g| g.header.track_alias == sub.media_alias && g.reset.is_none()) {
        for o in &g.objects {
            let vg = groups.get(&g.header.group_id).expect("the receipts of every delivered group arrived");
            let r = &receipts[&g.header.group_id];
            verifier.verify_payload(&o.payload, r, vg).expect("the payload verifies against its signed root");
            objects += 1;
        }
    }
    (objects, groups.len(), groups, receipts)
}

pub fn payload(tag: &str, g: u64, len: usize) -> Vec<u8> {
    let mut v = format!("{tag}:{g}:").into_bytes();
    v.resize(len.max(v.len()), (g % 251) as u8);
    v
}

/// Wait, bounded, until `f` holds.
pub async fn until(limit: Duration, mut f: impl FnMut() -> bool) -> bool {
    let deadline = Instant::now() + limit;
    while !f() {
        if Instant::now() >= deadline {
            return false;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    true
}

/// The scenario every transport runs: two publishers, one reached by the namespace route
/// (PUBLISH_NAMESPACE, the relay's SUBSCRIBE) and one by PUBLISH, and three subscribers
/// of both, all over `via`. Every object reaches every subscriber and verifies against its
/// publisher's signed receipts, which the relay also checks at the door; one subscriber
/// discovers the first namespace by SUBSCRIBE_NAMESPACE; every session's summary is signed.
pub async fn two_publishers_three_subscribers(via: Via) {
    two_publishers_three_subscribers_on([via; 5]).await
}

/// [`two_publishers_three_subscribers`] with each session on a transport (and draft) of
/// its own: publisher A, publisher B, then the three subscribers.
pub async fn two_publishers_three_subscribers_on(vias: [Via; 5]) {
    let key_a = SigningKey::from_bytes(&[0xa1; 32]);
    let key_b = SigningKey::from_bytes(&[0xb2; 32]);
    let (ka, kb) = (hex::encode(key_a.verifying_key().to_bytes()), hex::encode(key_b.verifying_key().to_bytes()));
    let server = relay(|c| {
        c.receipts.door = vec![joule_moq_relay::config::DoorRule { prefix: vec!["live".into()], trust: vec![ka, kb], trust_field: None, hold_ms: 5_000 }];
    })
    .await;
    const GROUPS: u64 = 12;
    let ta = Tracks::new(ns(&["live", "a"]), "cam", 10, 8, 64);
    let tb = Tracks::new(ns(&["live", "b"]), "cam", 20, 9, 96);

    let (pa, _ea) = connect(&server, vias[0], ClientOptions::default()).await;
    let (pb, _eb) = connect(&server, vias[1], ClientOptions::default()).await;
    let _ns_a = pa.publish_namespace(ta.ns.clone(), Parameters::new()).await.expect("PUBLISH_NAMESPACE");
    let _pub_b = publish_tracks(&pb, &tb, None).await;

    // One subscriber discovers live/a by prefix (§4.2).
    let (s0, _e0) = connect(&server, vias[2], ClientOptions::default()).await;
    let (_nstx, mut nsrx) = s0.subscribe_namespace(ns(&["live"])).await.expect("SUBSCRIBE_NAMESPACE");
    match timeout(STEP, nsrx.next()).await.expect("NAMESPACE in time").expect("a message") {
        Some(ControlMessage::Namespace(suffix)) => assert_eq!(suffix, ns(&["a"]), "NAMESPACE carries the suffix after the prefix"),
        other => panic!("expected NAMESPACE, got {other:?}"),
    }

    let mut subs = vec![s0];
    let mut eps = Vec::new();
    for via in &vias[3..] {
        let (s, e) = connect(&server, *via, ClientOptions::default()).await;
        subs.push(s);
        eps.push(e);
    }
    // The first SUBSCRIBE to live/a makes the relay subscribe upstream; publisher A answers.
    let answering = {
        let ta = ta.clone();
        let pa_ref = &pa;
        async move { answer_subscribes(pa_ref, &ta).await }
    };
    let s_first = &subs[0];
    let (held_a, first_sub) = tokio::join!(answering, subscribe_both(s_first, &ta, Parameters::new()));
    let mut views: Vec<(Subscribed, Subscribed)> = vec![(first_sub, subscribe_both(&subs[0], &tb, Parameters::new()).await)];
    for s in &subs[1..] {
        views.push((subscribe_both(s, &ta, Parameters::new()).await, subscribe_both(s, &tb, Parameters::new()).await));
    }

    // Both publish, concurrently.
    let pub_a = async {
        for g in 0..GROUPS {
            publish_group(&pa, &key_a, &ta, g, &payload("a", g, 900)).await;
        }
    };
    let pub_b = async {
        for g in 0..GROUPS {
            publish_group(&pb, &key_b, &tb, g, &payload("b", g, 1400)).await;
        }
    };
    let reads = futures::future::join_all(subs.iter().map(|s| collect(s, STEP * 2, |got| got.len() as u64 >= 4 * GROUPS)));
    let ((), (), got) = tokio::join!(pub_a, pub_b, reads);

    for (i, (g, (va, vb))) in got.iter().zip(&views).enumerate() {
        let (oa, ba, _, _) = verify(g, va, &key_a);
        let (ob, bb, _, _) = verify(g, vb, &key_b);
        assert_eq!((oa, ba, ob, bb), (GROUPS as usize, GROUPS as usize, GROUPS as usize, GROUPS as usize), "subscriber {i}: every object and every receipt of both publishers");
        for x in g.iter().filter(|x| x.header.track_alias == va.media_alias) {
            assert_eq!(x.objects[0].payload, payload("a", x.header.group_id, 900));
        }
    }
    assert_eq!(server.relay().stats().refused_at_door, 0);
    drop(held_a);

    // Every session ends with a signed summary.
    for s in &subs {
        s.close(0, "done");
    }
    pa.close(0, "done");
    pb.close(0, "done");
    assert!(until(STEP, || server.relay().summaries().len() >= 5).await, "five summaries");
    for sum in server.relay().summaries() {
        sum.verify().expect("the summary verifies under the relay's key");
        assert_eq!(sum.summary.relay_energy.acquisition, "unmetered");
        assert_eq!(sum.summary.relay_energy.joules_micro, 0, "0 is the unmetered marker");
        let transports: Vec<&str> = vias.iter().map(|v| match v { Via::Quic | Via::QuicDraft(_) => "quic", Via::WebTransport | Via::WtDraft(_) => "webtransport", Via::WebSocket => "websocket" }).collect();
        assert!(transports.contains(&sum.summary.transport.as_str()), "{}", sum.summary.transport);
    }
    let as_subscriber: Vec<_> = server.relay().summaries().into_iter().filter(|s| s.summary.tally.subscriber.objects_forwarded > 0).collect();
    assert_eq!(as_subscriber.len(), 3);
    for s in &as_subscriber {
        assert_eq!(s.summary.tally.subscriber.objects_forwarded, 4 * GROUPS, "media and receipts of both publishers");
        assert_eq!(s.summary.tally.subscriber.dropped_delivery_timeout, 0);
        let signed = s.summary.tally.publisher_signed.get("unmetered").expect("the door verified the media objects' receipts");
        assert_eq!(signed.objects, 2 * GROUPS, "the publishers' signed figures, unmetered, for the media objects");
    }
    server.shutdown().await;
}