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

//! A session over WebTransport is closed by closing the WebTransport session
//! (`draft-ietf-moq-transport-22` §6.6): the MOQT session termination code and the reason
//! ride a WT_CLOSE_SESSION capsule on the CONNECT stream (`draft-ietf-webtrans-http3-16`
//! §6), and the other end reads both exactly, as it does over the WebSocket binding. A QUIC
//! connection close alone carries an HTTP/3 error code, which no end reads as MOQT's.
//!
//! The relay closes a session four ways here: refusing a room grant presented without its
//! proof, after SETUP; refusing a stream that breaks a rule before SETUP, which it stops only
//! after the session's close; refusing a session past its session count, before any MOQT
//! message; and ending an established session at its time limit. And a client closes one,
//! whose code the relay's signed summary records as the peer's.

mod common;

use common::*;
use ed25519_dalek::SigningKey;
use joule_moq::draft22::{SessionError, Token};
use joule_moq_relay::client::ClientOptions;
use joule_moq_relay::config::Grant;
use joule_moq_relay::transport::{webtransport, Connection, TransportError};
use joule_moq_relay::RelayServer;
use std::time::Duration;
use tokio::time::timeout;

/// A record shaped as a room grant (`WRG1`, version 1, a room, the principal), which is
/// all the relay reads of it: its rule names the record by its SHA-256.
fn grant_record(room: u8, principal: &SigningKey) -> Vec<u8> {
    let mut g = b"WRG1\x01".to_vec();
    g.extend_from_slice(&[room; 32]);
    g.extend_from_slice(&principal.verifying_key().to_bytes());
    g.extend_from_slice(&[0x5a; 160]);
    g.extend_from_slice(&[0x77; 64]);
    g
}

async fn closed(conn: &dyn Connection, what: &str) -> TransportError {
    timeout(STEP, conn.closed()).await.unwrap_or_else(|_| panic!("{what}: not closed in time"))
}

fn by_relay(code: SessionError, reason: &str) -> TransportError {
    TransportError::Closed { code: code.0, reason: reason.into(), by_peer: true }
}

/// A WebTransport session to `server` with no MOQT on it yet.
async fn session(server: &RelayServer) -> webtransport::WtConnection {
    let url = format!("https://127.0.0.1:{}{}", server.quic_addr.unwrap().port(), server.relay().config().path);
    timeout(STEP, webtransport::connect(&url, Some(server.identity.leaf_sha256), loopback()))
        .await
        .expect("in time")
        .expect("a WebTransport session")
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_session_over_webtransport_closes_with_its_code_and_reason() {
    timeout(Duration::from_secs(120), run()).await.expect("the test finishes in time");
}

async fn run() {
    // A room grant presented as a bare token: the relay refuses the session after SETUP,
    // UNAUTHORIZED (WAI live-rooms §10.5.2).
    let host = SigningKey::from_bytes(&[0x51; 32]);
    let record = grant_record(3, &host);
    let hash = hex::encode(joule_moq_relay::auth::sha256(&record));
    let rooms = relay(|c| {
        c.auth.required = true;
        c.auth.grants = vec![Grant { prefix: vec!["room".into()], token_sha256: hash, publish: true, subscribe: true, proof: true, unlimited: false, seek: false }];
    })
    .await;
    let bare = ClientOptions { tokens: vec![Token::UseValue { token_type: 0, value: record }], ..ClientOptions::default() };
    let (c, _) = connect(&rooms, Via::WebTransport, bare).await;
    assert_eq!(closed(c.conn.as_ref(), "a grant with no proof").await, by_relay(SessionError::UNAUTHORIZED, "live-rooms §10.5: pop_missing"));

    // A stream that breaks a rule before SETUP: refused, PROTOCOL_VIOLATION. The relay ends
    // that stream only after the session's close, so the client reads the close and never
    // a STOP_SENDING ahead of it.
    let raw = session(&rooms).await;
    let mut s = timeout(STEP, raw.open_uni()).await.expect("in time").expect("a stream");
    s.write_all(&[0x3f; 4]).await.expect("written");
    match closed(&raw, "a stream that breaks a rule").await {
        TransportError::Closed { code, reason, by_peer: true } if code == SessionError::PROTOCOL_VIOLATION.0 => assert!(reason.contains("invalid flags"), "{reason}"),
        other => panic!("a stream that breaks a rule: {other:?}"),
    }
    match timeout(STEP, s.delivered()).await.expect("in time") {
        Err(TransportError::Closed { .. }) => {}
        other => panic!("the relay stopped the stream ahead of the session's close: {other:?}"),
    }

    // A client closes its session with a code of its own: the relay reads it as the peer's,
    // and its summary records it; the client's own end reads its close as local.
    let (c, _) = connect(&rooms, Via::WebTransport, ClientOptions::default()).await;
    // The client has its handshake once the relay's SETUP arrives; the relay reads the
    // client's SETUP after sending its own, and a session closed before then was never served
    // and gets no summary. So the close waits until the relay holds the session.
    assert!(until(STEP, || rooms.relay().session_diagnostics().len() == 1).await, "the relay holds the session");
    c.close(SessionError::PROTOCOL_VIOLATION.0, "a client's own close");
    assert_eq!(
        closed(c.conn.as_ref(), "the client's close").await,
        TransportError::Closed { code: 0, reason: "closed by this endpoint".into(), by_peer: false }
    );
    let recorded = until(STEP, || {
        rooms.relay().summaries().iter().any(|s| s.summary.transport == "webtransport" && s.summary.close_code == Some(SessionError::PROTOCOL_VIOLATION.0))
    })
    .await;
    let codes: Vec<_> = rooms.relay().summaries().iter().map(|s| (s.summary.transport.clone(), s.summary.close_code)).collect();
    assert!(recorded, "the relay records the client's code: {codes:?}");
    rooms.shutdown().await;

    // A relay that holds one session at a time, each for three seconds at most.
    let limited = relay(|c| {
        c.limits.max_sessions = 1;
        c.limits.max_session_ms = 3_000;
    })
    .await;
    let (held, _) = connect(&limited, Via::WebTransport, ClientOptions::default()).await;
    assert!(until(STEP, || limited.relay().session_diagnostics().len() == 1).await, "the first session is held");
    // Past its session count: refused before any MOQT message, INTERNAL_ERROR.
    let refused = session(&limited).await;
    assert_eq!(closed(&refused, "a session past the count").await, by_relay(SessionError::INTERNAL_ERROR, "the relay is at its session limit"));
    // At its time limit, an established session is ended with NO_ERROR and the reason, which
    // a FIN on the CONNECT stream alone (code 0, no reason) would not carry.
    assert_eq!(closed(held.conn.as_ref(), "a session at its time limit").await, by_relay(SessionError::NO_ERROR, "the relay's session limit of 3 s"));
    limited.shutdown().await;
}