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

//! The relay's own limits against the considerations of `draft-englishm-moq-relay-dos-02`
//! (`src/dos.rs`; the README's *Denial of service* table), each tripped with its stated answer:
//!
//! - **a flood of SETUPs to refuse** (the draft's §3.2: what is cheap to open in quantity): a
//!   session refused at SETUP is held open for its client to read the refusal, outside the
//!   session count. Ten from one address against a bound of three an address: the seven oldest
//!   are closed at once, UNAUTHORIZED, the three newest stay held, and none counts as a
//!   session. Against a bound of four in all, with the per-address bound out of the way, the
//!   six oldest go. Planted: the bounds removed, and all ten are held, which is what the test's
//!   own assertions would then refuse;
//! - **what one session holds** (§4.3, §6.1, §6.3): a third subscription, a third published
//!   track and a second namespace are each refused `EXCESSIVE_LOAD` (`session_limit`), nothing
//!   the session holds is ended, and a slot comes back when its request ends. Planted: the
//!   limit removed (0), and the same request is served;
//! - **rapid request cycles** (§4.2): the request past a session's rate is refused
//!   `EXCESSIVE_LOAD` (`request_rate`) with the Retry Interval after which one is served.

mod common;

use common::*;
use ed25519_dalek::SigningKey;
use joule_moq::draft22::{ControlMessage, NamespaceRequest, Parameters, Publish, RequestErrorCode, SessionError, StreamResetCode, TrackRequest};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::config::Grant;
use joule_moq_relay::transport::TransportError;
use joule_moq_relay::RelayServer;
use std::time::Duration;
use tokio::time::timeout;

fn grant_record(principal: &SigningKey) -> Vec<u8> {
    let mut g = b"WRG1\x01".to_vec();
    g.extend_from_slice(&[9; 32]);
    g.extend_from_slice(&principal.verifying_key().to_bytes());
    g.extend_from_slice(&[0x5a; 160]);
    g.extend_from_slice(&[0x77; 64]);
    g
}

/// A relay that refuses, at SETUP, a grant presented with no proof, and holds each refusal for
/// a minute unless a bound sheds it.
async fn refusing(overall: usize, per_address: usize) -> (RelayServer, ClientOptions) {
    let host = SigningKey::from_bytes(&[0x41; 32]);
    let record = grant_record(&host);
    let hash = hex::encode(joule_moq_relay::auth::sha256(&record));
    let server = 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 }];
        c.limits.setup_refusal_wait_ms = 60_000;
        c.limits.max_held_refusals = overall;
        c.limits.max_held_refusals_per_address = per_address;
    })
    .await;
    (server, ClientOptions { grants: vec![(record, host)], ignore_refusal: true, ..ClientOptions::default() })
}

/// Ten SETUPs the relay refuses, one after another, each client staying; then which of them
/// the relay has closed (UNAUTHORIZED) within a second, oldest first.
async fn flood(server: &RelayServer, opts: &ClientOptions) -> Vec<bool> {
    let mut clients = Vec::new();
    for _ in 0..10 {
        let (conn, ep) = dial(server, Via::Quic, None).await;
        let c = Client::connect(conn, opts.clone()).await.expect("the relay's SETUP, refusing");
        clients.push((c, ep));
        // In order: the relay holds each before the next arrives.
        let n = clients.len() as u64;
        assert!(until(STEP, || { let s = server.relay().stats(); s.refusals_held + s.refusals_shed >= n }).await, "refusal {n} is held or shed");
    }
    let mut closed = Vec::new();
    for (c, _) in &clients {
        closed.push(match timeout(Duration::from_secs(1), c.conn.closed()).await {
            Ok(TransportError::Closed { code, .. }) => {
                assert_eq!(code, SessionError::UNAUTHORIZED.0);
                true
            }
            Ok(other) => panic!("{other:?}"),
            Err(_) => false,
        });
    }
    for (c, _) in &clients {
        c.close(0, "done");
    }
    closed
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_flood_of_refused_setups_holds_a_bounded_number() {
    timeout(Duration::from_secs(180), async {
        // By address: three held, the seven oldest shed.
        let (server, opts) = refusing(256, 3).await;
        let closed = flood(&server, &opts).await;
        assert_eq!(closed, [true, true, true, true, true, true, true, false, false, false], "the oldest go first");
        let s = server.relay().stats();
        assert_eq!((s.refusals_shed, s.sessions_open), (7, 0), "shed, and none of them a session");
        assert!(until(STEP, || server.relay().stats().refusals_held == 0).await, "a client that leaves frees its place");
        server.shutdown().await;

        // In all: four held.
        let (server, opts) = refusing(4, 1000).await;
        let closed = flood(&server, &opts).await;
        assert_eq!(closed.iter().filter(|c| **c).count(), 6);
        assert!(closed[..6].iter().all(|c| *c) && closed[6..].iter().all(|c| !*c));
        server.shutdown().await;

        // Planted: the bounds removed. All ten are held for the whole wait, none shed: the
        // state the two cases above refuse.
        let (server, opts) = refusing(usize::MAX, usize::MAX).await;
        let closed = flood(&server, &opts).await;
        assert!(closed.iter().all(|c| !*c), "with no bound every refused SETUP holds its connection");
        assert_eq!(server.relay().stats().refusals_shed, 0);
        server.shutdown().await;
    })
    .await
    .expect("the test finishes in time");
}

async fn answer(c: &Client, m: ControlMessage) -> (joule_moq_relay::wire::ReqTx, joule_moq_relay::wire::ReqRx, ControlMessage) {
    let (tx, mut rx) = c.request(&m).await.unwrap();
    let r = Client::response(&mut rx).await.unwrap();
    (tx, rx, r)
}

fn refused_for(r: &ControlMessage, why: &str) -> u64 {
    let ControlMessage::RequestError(e) = r else { panic!("{} where a refusal was due", r.name()) };
    assert_eq!(e.code, RequestErrorCode::EXCESSIVE_LOAD);
    assert_eq!(e.reason, why.as_bytes());
    e.retry_interval
}

async fn held_to(max: u64) {
    let server = relay(|c| {
        c.limits.max_subscriptions_per_session = max;
        c.limits.max_tracks_per_session = max;
        c.limits.max_namespaces_per_session = max / 2;
        c.limits.max_requests_per_s = 0;
    })
    .await;
    let limited = max != 0;
    let (p, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let tracks: Vec<_> = (0..3u64).map(|i| (name(&format!("t{i}")), 10 + i)).collect();
    let space = ns(&["live", "r"]);
    let publish = |i: usize| ControlMessage::Publish(Publish { request_id: p.next_request_id(), namespace: space.clone(), name: tracks[i].0.clone(), track_alias: tracks[i].1, params: Parameters::new(), track_properties: track_props(128, None) });
    let (_p0t, _p0r, r0) = answer(&p, publish(0)).await;
    let (_p1t, _p1r, r1) = answer(&p, publish(1)).await;
    assert!(matches!((&r0, &r1), (ControlMessage::RequestOk(_), ControlMessage::RequestOk(_))));
    let (_p2t, _p2r, r2) = answer(&p, publish(2)).await;
    if limited {
        refused_for(&r2, joule_moq_relay::dos::SESSION_LIMIT);
    } else {
        assert!(matches!(r2, ControlMessage::RequestOk(_)), "planted: with the limit removed a third track is published");
    }
    let nsreq = |c: &Client, f: &str| ControlMessage::PublishNamespace(NamespaceRequest { request_id: c.next_request_id(), prefix: ns(&["live", f]), params: Parameters::new() });
    let (_n0t, _n0r, n0) = answer(&p, nsreq(&p, "a")).await;
    assert!(matches!(n0, ControlMessage::RequestOk(_)));
    let (_n1t, _n1r, n1) = answer(&p, nsreq(&p, "b")).await;
    if limited {
        refused_for(&n1, joule_moq_relay::dos::SESSION_LIMIT);
    } else {
        assert!(matches!(n1, ControlMessage::RequestOk(_)));
    }

    let (c, _) = connect(&server, Via::WebSocket, ClientOptions::default()).await;
    let sub = |i: usize| ControlMessage::Subscribe(TrackRequest { request_id: c.next_request_id(), namespace: space.clone(), name: tracks[i].0.clone(), params: Parameters::new() });
    let (s0t, _s0r, s0) = answer(&c, sub(0)).await;
    let (_s1t, _s1r, s1) = answer(&c, sub(1)).await;
    assert!(matches!((&s0, &s1), (ControlMessage::SubscribeOk(_), ControlMessage::SubscribeOk(_))));
    // A third: the same track again is a subscription all the same.
    let (_s2t, _s2r, s2) = answer(&c, sub(0)).await;
    if limited {
        refused_for(&s2, joule_moq_relay::dos::SESSION_LIMIT);
        assert_eq!(server.relay().stats().load_refused, 3);
        // What the session holds stands, and the slot of a request that ends comes back.
        s0t.reset(StreamResetCode::CANCELLED.0).await;
        let mut served = false;
        for _ in 0..50 {
            let (_t, _r, again) = answer(&c, sub(1)).await;
            if matches!(again, ControlMessage::SubscribeOk(_)) {
                served = true;
                break;
            }
            refused_for(&again, joule_moq_relay::dos::SESSION_LIMIT);
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
        assert!(served, "a subscription is served once one of the two ended");
    } else {
        assert!(matches!(s2, ControlMessage::SubscribeOk(_)), "planted: with the limit removed a third subscription is served");
        assert_eq!(server.relay().stats().load_refused, 0);
    }
    c.close(0, "done");
    p.close(0, "done");
    server.shutdown().await;
}

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

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn requests_past_a_sessions_rate_are_refused_with_when_to_retry() {
    timeout(Duration::from_secs(120), async {
        let server = relay(|c| c.limits.max_requests_per_s = 4).await;
        let (c, _) = connect(&server, Via::WebSocket, ClientOptions::default()).await;
        let status = |c: &Client| ControlMessage::TrackStatus(TrackRequest { request_id: c.next_request_id(), namespace: ns(&["x"]), name: name("y"), params: Parameters::new() });
        let mut retry = None;
        // A second's worth at once is served (each DOES_NOT_EXIST: no such track); within the
        // same instant the next few are not.
        for i in 0..12 {
            let (_t, _r, a) = answer(&c, status(&c)).await;
            match &a {
                ControlMessage::RequestError(e) if e.code == RequestErrorCode::EXCESSIVE_LOAD => {
                    assert!(i >= 4, "the first four are served");
                    retry = Some(refused_for(&a, joule_moq_relay::dos::REQUEST_RATE));
                    break;
                }
                _ => {}
            }
        }
        let retry = retry.expect("a request past the rate is refused");
        assert!((2..=252).contains(&retry), "the Retry Interval is milliseconds plus one, within a request's worth: {retry}");
        tokio::time::sleep(Duration::from_millis(retry + 20)).await;
        let (_t, _r, a) = answer(&c, status(&c)).await;
        assert!(!matches!(&a, ControlMessage::RequestError(e) if e.code == RequestErrorCode::EXCESSIVE_LOAD), "served after the interval");
        // The session stands throughout.
        assert_eq!(server.relay().stats().sessions_open, 1);
        c.close(0, "done");
        server.shutdown().await;
    })
    .await
    .expect("the test finishes in time");
}