open-standards/joule-code/joule-code-rs/crates/joule-moq-ingest/tests/shutdown_ports.rs

//! A shutdown returns with every port the gateway listened on closed, so a gateway can start
//! again at once on the same ports, in the same process, for a new room: HTTP, RTMP, SRT, and
//! a RIST session's ports (with the `rist` feature). By then every connection the gateway held
//! is closed, and every session that was running has signed its record. Each test runs in a
//! task on a runtime of one worker, as a restart from a spawned task does: there a task that a
//! shutdown only told to go (`JoinHandle::abort`) cannot have run to its end, listener and
//! all, until the shutdown's caller yields, so a bind straight after such a shutdown fails
//! every time. Each test shuts the gateway down and starts it again on its ports many times.
//!
//! Two tests hold the shutdown to what it ends past its bound (`TASKS_END`, 5 s), with the
//! `rist` feature: a RIST session whose relay does not answer its dial (a dial lasts up to
//! 10 s) ends at the stop and signs its record, and a RIST listener held by a task that is
//! still running at the bound is closed however that task was ended.

mod common;

use common::*;
use joule_contrib::srt::{self, CallConfig};
use joule_moq_ingest::{GatewayConfig, GatewayServer};
use std::net::SocketAddr;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::time::timeout;

const ROUNDS: usize = 30;
const SESSION_ROUNDS: usize = 6;
const RIST: bool = cfg!(feature = "rist");

fn any_port() -> SocketAddr {
    "127.0.0.1:0".parse().unwrap()
}

/// The gateway's configuration with the ports it bound in place of the ones it was given:
/// what a restart on fixed ports starts with.
fn fixed(s: &GatewayServer) -> GatewayConfig {
    let mut c = s.gateway.cfg.clone();
    c.http.listen = s.addr;
    c.rtmp.listen = s.rtmp_addr;
    c.contrib.srt.listen = s.srt_addr;
    c
}

/// Bind each address at once, as a gateway starting again on it does first.
fn free(tcp: &[SocketAddr], udp: &[SocketAddr], round: usize) {
    for &a in tcp {
        if let Err(e) = std::net::TcpListener::bind(a) {
            panic!("round {round}: TCP {a} is still bound when the shutdown returns: {e}");
        }
    }
    for &a in udp {
        if let Err(e) = std::net::UdpSocket::bind(a) {
            panic!("round {round}: UDP {a} is still bound when the shutdown returns: {e}");
        }
    }
}

/// An HTTP connection the gateway answered once and holds, kept alive and idle.
async fn idle_http(at: SocketAddr) -> TcpStream {
    let mut s = TcpStream::connect(at).await.unwrap();
    s.write_all(b"GET /whip HTTP/1.1\r\nhost: gateway\r\n\r\n").await.unwrap();
    let mut got = Vec::new();
    let mut b = [0u8; 512];
    while !got.windows(4).any(|w| w == b"\r\n\r\n") {
        let n = timeout(STEP, s.read(&mut b)).await.expect("an answer in time").unwrap();
        assert!(n > 0, "the gateway answers before it closes");
        got.extend_from_slice(&b[..n]);
    }
    assert!(got.starts_with(b"HTTP/1.1 204"), "{}", String::from_utf8_lossy(&got));
    s
}

/// Whether the gateway has closed its end of `s`: a read returns at once, with nothing.
async fn closed_by_gateway(s: &mut TcpStream) -> bool {
    let mut b = [0u8; 64];
    matches!(timeout(Duration::from_secs(2), s.read(&mut b)).await, Ok(Ok(0) | Err(_)))
}

/// A one-time session at `/<kind>` (`srt` or `rist`) for `who`'s grant, and its JSON.
async fn open(w: &World, who: &ed25519_dalek::SigningKey, kind: &str, body: &str) -> serde_json::Value {
    let auth = w.auth(who, 0);
    let r = auth.send("POST", &format!("{}/{kind}", w.base()), &[("content-type", "application/json")], body.as_bytes()).await;
    assert_eq!(r.status, 201, "POST /{kind}: {}", r.text());
    serde_json::from_slice(&r.body).unwrap()
}

/// Shut `w`'s gateway down, check `tcp` and `udp` are free when that returns, and start it
/// again on its ports, with `adjust` applied to its configuration.
async fn restart(w: World, tcp: &[SocketAddr], udp: &[SocketAddr], round: usize, adjust: impl FnOnce(&mut GatewayConfig)) -> World {
    let World { relay, gateway, authority, descriptor, room, publisher, gateway_key } = w;
    let mut cfg = fixed(&gateway);
    adjust(&mut cfg);
    gateway.shutdown().await;
    free(tcp, udp, round);
    let gateway = GatewayServer::start(cfg).await.unwrap_or_else(|e| panic!("round {round}: the gateway starts again on the same ports: {e}"));
    World { relay, gateway, authority, descriptor, room, publisher, gateway_key }
}

/// A RIST Simple Profile flow to the session at `at` under `cname`, held open.
#[cfg(feature = "rist")]
async fn rist_flow(at: SocketAddr, cname: &str) -> joule_contrib::link::Link {
    let cfg = joule_contrib::rist::SendConfig { cname: cname.into(), buffer_ms: 300, ..Default::default() };
    joule_contrib::rist::send(at, cfg).await.unwrap()
}

#[cfg(not(feature = "rist"))]
async fn rist_flow(_: SocketAddr, _: &str) -> joule_contrib::link::Link {
    unreachable!("a build without RIST opens no RIST session")
}

fn ports(c: &mut GatewayConfig) {
    c.rtmp.listen = Some(any_port());
    c.contrib.srt.listen = Some(any_port());
    c.contrib.rist.enabled = RIST;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn a_shutdown_from_a_task_frees_every_port_for_a_start_at_once() {
    let t = tokio::spawn(async {
        let mut w = World::start(ports).await;
        let mut rist_port = None;
        for round in 0..ROUNDS {
            let g = w.gateway.gateway.clone();
            let (http, rtmp, srt) = (w.gateway.addr, w.gateway.rtmp_addr.unwrap(), w.gateway.srt_addr.unwrap());
            let mut udp = vec![srt];
            // An HTTP connection kept alive after its request, and an RTMP connection that has
            // not begun its handshake: each a task of the gateway's.
            let mut h = idle_http(http).await;
            let mut r = TcpStream::connect(rtmp).await.unwrap();
            assert!(until(STEP, || g.rtmp.lock().unwrap().len() == 1).await, "the gateway holds the RTMP connection");
            if RIST {
                // A RIST session waiting for its flow, on the port the gateway picked in the
                // first round and is given from then on.
                let s = open(&w, &w.publisher, "rist", r#"{"profile":"simple"}"#).await;
                let port = s["port"].as_u64().unwrap() as u16;
                assert_eq!(*rist_port.get_or_insert(port), port, "the session has the configured port");
                udp.extend([SocketAddr::new(http.ip(), port), SocketAddr::new(http.ip(), port + 1)]);
            }
            w = restart(w, &[http, rtmp], &udp, round, |c| {
                if let Some(p) = rist_port {
                    (c.contrib.rist.port_min, c.contrib.rist.port_max) = (p, p);
                }
            })
            .await;
            assert!(closed_by_gateway(&mut h).await, "round {round}: the HTTP connection is closed when the shutdown returns");
            assert!(closed_by_gateway(&mut r).await, "round {round}: the RTMP connection is closed when the shutdown returns");
        }
        w.stop().await;
    });
    timeout(Duration::from_secs(300), t).await.expect("in time").unwrap();
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn a_shutdown_ends_the_sessions_that_run_and_frees_their_ports() {
    let t = tokio::spawn(async {
        let mut w = World::start(ports).await;
        // Two publishers: one principal holds one contribution session at a time.
        let second = ed25519_dalek::SigningKey::from_bytes(&[0x54; 32]);
        let mut rist_port = None;
        for round in 0..SESSION_ROUNDS {
            let g = w.gateway.gateway.clone();
            let (http, rtmp, srt_at) = (w.gateway.addr, w.gateway.rtmp_addr.unwrap(), w.gateway.srt_addr.unwrap());
            let mut udp = vec![srt_at];
            // An SRT session publishing: its connection's task holds the listener's socket.
            let s = open(&w, &w.publisher, "srt", "").await;
            let cfg = CallConfig {
                stream_id: s["stream_id"].as_str().map(str::to_string),
                passphrase: s["passphrase"].as_str().map(str::to_string),
                keylen: s["pbkeylen"].as_u64().unwrap() as usize,
                latency_ms: 120,
                ..CallConfig::default()
            };
            let srt_link = timeout(STEP, srt::call(srt_at, cfg)).await.unwrap().expect("the call is admitted");
            assert!(contrib_live(&w, s["session"].as_str().unwrap()).await, "the SRT session publishes its tracks");
            let mut rist_link = None;
            if RIST {
                // A RIST session running: its listener holds its ports for the session's life.
                let s = open(&w, &second, "rist", r#"{"profile":"simple"}"#).await;
                let port = s["port"].as_u64().unwrap() as u16;
                assert_eq!(*rist_port.get_or_insert(port), port, "the session has the configured port");
                udp.extend([SocketAddr::new(http.ip(), port), SocketAddr::new(http.ip(), port + 1)]);
                rist_link = Some(rist_flow(SocketAddr::new(http.ip(), port), s["cname"].as_str().unwrap()).await);
                assert!(contrib_live(&w, s["session"].as_str().unwrap()).await, "the RIST session publishes its tracks");
            }
            w = restart(w, &[http, rtmp], &udp, round, |c| {
                if let Some(p) = rist_port {
                    (c.contrib.rist.port_min, c.contrib.rist.port_max) = (p, p);
                }
            })
            .await;
            // Each session the shutdown ended has signed its record by then, as ended by it.
            let recs = g.records();
            let ended = |protocol: &str| recs.iter().filter(|r| r.record.protocol.starts_with(protocol)).map(|r| r.record.ended.clone()).collect::<Vec<_>>();
            assert_eq!(ended("srt"), ["deleted"], "round {round}: the SRT session's record");
            if RIST {
                assert_eq!(ended("rist"), ["deleted"], "round {round}: the RIST session's record");
            }
            srt_link.close();
            if let Some(l) = rist_link {
                l.close();
            }
        }
        w.stop().await;
    });
    timeout(Duration::from_secs(300), t).await.expect("in time").unwrap();
}

/// A RIST session whose flow is taken while the relay does not answer: the session dials it,
/// and a dial lasts up to 10 s (`relay::STEP`), past the shutdown's wait for the gateway's
/// tasks. The shutdown ends the dial; the session ends at once, as one that starts once the
/// shutdown has begun does, and signs its record; its ports are free when the shutdown returns.
#[cfg(feature = "rist")]
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn a_session_still_dialing_the_relay_at_a_stop_signs_its_record_and_frees_its_port() {
    let t = tokio::spawn(async {
        // Takes every datagram and answers none.
        let silent = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
        for round in 0..2 {
            let (at, port) = (silent.local_addr().unwrap(), low_pair());
            let w = World::start(|c| {
                c.contrib.rist.enabled = true;
                (c.contrib.rist.port_min, c.contrib.rist.port_max) = (port, port);
                c.relay.addr = at;
            })
            .await;
            let g = w.gateway.gateway.clone();
            let ip = w.gateway.addr.ip();
            let s = open(&w, &w.publisher, "rist", r#"{"profile":"simple"}"#).await;
            assert_eq!(s["port"].as_u64(), Some(port as u64), "round {round}: the session has the configured port");
            let flow = rist_flow(SocketAddr::new(ip, port), s["cname"].as_str().unwrap()).await;
            // The flow is taken: the session runs, dialing the relay.
            assert!(until(STEP, || g.contrib.active.load(std::sync::atomic::Ordering::SeqCst) == 1).await, "round {round}: the session takes its flow");
            assert!(g.contrib.pending.lock().unwrap().is_empty(), "round {round}: the session is no longer pending");
            let World { relay, gateway, .. } = w;
            let began = std::time::Instant::now();
            gateway.shutdown().await;
            let took = began.elapsed();
            free(&[], &[SocketAddr::new(ip, port), SocketAddr::new(ip, port + 1)], round);
            let recs = g.records();
            let rist: Vec<_> = recs.iter().filter(|r| r.record.protocol.starts_with("rist")).map(|r| r.record.ended.clone()).collect();
            assert_eq!(rist, ["gateway_stopped"], "round {round}: the session's record is signed, as ended by the stop (the shutdown took {took:?})");
            assert_eq!(g.contrib.active.load(std::sync::atomic::Ordering::SeqCst), 0, "round {round}: the session has ended");
            flow.close();
            relay.shutdown().await;
        }
    });
    timeout(Duration::from_secs(300), t).await.expect("in time").unwrap();
}

/// A RIST listener held by a task of the gateway's that does not end by itself: the shutdown
/// ends that task past its bound, and the listener, dropped with it, would only tell its
/// receiver to go. Held by the gateway as well, it is closed, and its ports are free when the
/// shutdown returns. Only the listener's ports are checked: the gateway's own are let go 5 s
/// before the shutdown returns, and those it was handed for port 0 can be any socket's by then.
#[cfg(feature = "rist")]
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn a_rist_listener_whose_task_is_ended_past_the_bound_is_closed_by_the_shutdown() {
    use joule_contrib::rist::{recv::RecvConfig, Profile, RistListener, RistListenerConfig};
    let t = tokio::spawn(async {
        for round in 0..3 {
            let World { relay, gateway, .. } = World::start(|c| c.contrib.rist.enabled = true).await;
            let g = gateway.gateway.clone();
            let ip = gateway.addr.ip();
            let cfg = |port| RistListenerConfig {
                profile: Profile::Simple,
                bind_ip: ip,
                port,
                recv: RecvConfig { buffer_us: 200_000, ..RecvConfig::default() },
                secret: None,
                expect_cname: Some("held".into()),
                tag: 7,
            };
            let mut l = None;
            for _ in 0..8 {
                if let Ok(x) = RistListener::bind(cfg(low_pair())).await {
                    l = Some(x);
                    break;
                }
            }
            let l = l.unwrap_or_else(|| panic!("round {round}: a listener binds a free pair"));
            let port = l.addr.port();
            g.hold_rist(&l);
            g.spawn(async move {
                let _held = l;
                std::future::pending::<()>().await
            });
            gateway.shutdown().await;
            free(&[], &[SocketAddr::new(ip, port), SocketAddr::new(ip, port + 1)], round);
            relay.shutdown().await;
        }
    });
    timeout(Duration::from_secs(300), t).await.expect("in time").unwrap();
}