open-standards/joule-code/joule-code-rs/crates/joule-contrib/src/cli.rs

//! What the command-line tools share: their URLs, their sources of transport stream, and
//! links opened from a URL.
//!
//! ```text
//! srt://HOST:PORT[?streamid=S&passphrase=P&pbkeylen=16|24|32&latency=MS&mode=caller|listener]
//! rist://HOST:PORT[?profile=simple|main&cname=C&secret=P&aes-type=128|256&buffer=MS]
//! ```
//!
//! An SRT URL calls by default; `mode=listener` (or an empty host, `srt://:PORT`) listens.

use crate::link::Link;
use crate::srt::{self, Admitted, CallConfig, ListenerConfig, Role};
use crate::ts::{Mux, Pacer, ST_H264};
use std::collections::BTreeMap;
use std::net::{SocketAddr, ToSocketAddrs};
use std::sync::Arc;

#[derive(Clone, Debug)]
pub struct Url {
    pub scheme: String,
    pub host: String,
    pub port: u16,
    pub q: BTreeMap<String, String>,
}

impl Url {
    pub fn parse(s: &str) -> Result<Url, String> {
        let (scheme, rest) = s.split_once("://").ok_or("a URL is scheme://host:port")?;
        let (hp, q) = rest.split_once('?').unwrap_or((rest, ""));
        let (host, port) = hp.rsplit_once(':').ok_or("a URL names a port")?;
        let port: u16 = port.trim_end_matches('/').parse().map_err(|_| "the port is a number")?;
        let host = host.trim_start_matches('[').trim_end_matches(']').to_string();
        let mut m = BTreeMap::new();
        for kv in q.split('&').filter(|x| !x.is_empty()) {
            let (k, v) = kv.split_once('=').unwrap_or((kv, ""));
            m.insert(k.to_string(), percent_decode(v));
        }
        Ok(Url { scheme: scheme.to_ascii_lowercase(), host, port, q: m })
    }

    pub fn addr(&self) -> Result<SocketAddr, String> {
        let h = if self.host.is_empty() { "0.0.0.0" } else { &self.host };
        (h, self.port).to_socket_addrs().map_err(|e| e.to_string())?.next().ok_or_else(|| "the host does not resolve".into())
    }

    pub fn get(&self, k: &str) -> Option<&str> {
        self.q.get(k).map(String::as_str)
    }

    pub fn num<T: std::str::FromStr>(&self, k: &str, d: T) -> T {
        self.get(k).and_then(|v| v.parse().ok()).unwrap_or(d)
    }

    pub fn listens(&self) -> bool {
        self.get("mode") == Some("listener") || self.host.is_empty()
    }
}

fn percent_decode(s: &str) -> String {
    let b = s.as_bytes();
    let mut o = Vec::with_capacity(b.len());
    let mut i = 0;
    while i < b.len() {
        if b[i] == b'%'
            && i + 2 < b.len()
            && let Ok(v) = u8::from_str_radix(std::str::from_utf8(&b[i + 1..i + 3]).unwrap_or(""), 16)
        {
            o.push(v);
            i += 3;
        } else {
            o.push(b[i]);
            i += 1;
        }
    }
    String::from_utf8_lossy(&o).into_owned()
}

/// Open an SRT link from `u` in role `role`: call, or listen and take the first caller.
pub async fn srt_link(u: &Url, role: Role) -> Result<Link, String> {
    let latency_ms: u16 = u.num("latency", 120);
    let passphrase = u.get("passphrase").map(str::to_string);
    let keylen: usize = u.num("pbkeylen", 16);
    if u.listens() {
        let p2 = passphrase.clone();
        let adm = move |_: &srt::ConnRequest| Ok(Admitted { passphrase: p2.clone(), role, latency_ms: None, tag: 0 });
        let cfg = ListenerConfig { bind: u.addr()?, latency_ms, advertise_keylen: if passphrase.is_some() { keylen } else { 0 }, ..ListenerConfig::default() };
        let mut l = srt::Listener::bind(cfg, Arc::new(adm)).await.map_err(|e| e.to_string())?;
        eprintln!("listening on srt://{}", l.addr);
        let link = l.accept().await.ok_or("the listener stopped")?;
        // The listener keeps serving the link's packets until it ends.
        std::mem::forget(l);
        Ok(link)
    } else {
        let cfg = CallConfig { latency_ms, passphrase, keylen, stream_id: u.get("streamid").map(str::to_string), role, ..CallConfig::default() };
        srt::call(u.addr()?, cfg).await.map_err(|e| e.to_string())
    }
}

/// A synthetic program for `seconds`: one H.264-shaped access unit every 1/30 s (a
/// keyframe each second) and two ADTS frames of AAC every 1/25 s, multiplexed and timed.
/// The bytes are not decodable media; they have the shape of it.
pub fn synthetic(seconds: u64, kbps: u64) -> Vec<(u64, Vec<u8>)> {
    let mut m = Mux::new(ST_H264);
    let mut ts = Vec::new();
    let frames = seconds * 30;
    let per_frame = (kbps * 1000 / 8 / 30).max(200) as usize;
    let mut a = 0u64;
    for i in 0..frames {
        let key = i % 30 == 0;
        let mut au = vec![0, 0, 0, 1, 0x09, 0xf0];
        if key {
            au.extend_from_slice(&[0, 0, 0, 1, 0x67, 0x64, 0x00, 0x28, 0xac, 0, 0, 0, 1, 0x68, 0xee, 0x3c, 0x80, 0, 0, 0, 1, 0x65]);
        } else {
            au.extend_from_slice(&[0, 0, 0, 1, 0x41]);
        }
        let n = if key { per_frame * 3 } else { per_frame };
        au.extend((0..n).map(|j| ((i as usize * 131 + j * 7) % 251 + 1) as u8));
        let pts = 126_000 + i * 3000;
        ts.extend(m.pes(true, pts, &au, key));
        // Audio up to this frame's time: two frames of 1024 samples at 48 kHz a PES, 3840
        // ticks of 90 kHz.
        while a * 3840 <= i * 3000 {
            let mut f = Vec::new();
            for k in 0..2 {
                let body: Vec<u8> = (0..160).map(|j| ((a * 2 + k + j) % 200 + 9) as u8).collect();
                let len = body.len() + 7;
                f.extend_from_slice(&[0xff, 0xf1, (1 << 6) | (3 << 2), (2 << 6) | ((len >> 11) as u8 & 0x3), (len >> 3) as u8, ((len as u8 & 0x7) << 5) | 0x1f, 0xfc]);
                f.extend(body);
            }
            ts.extend(m.pes(false, 126_000 + a * 3840, &f, false));
            a += 1;
        }
    }
    Pacer::new(&ts, 7).chunks().to_vec()
}

/// A transport stream file's chunks of seven packets, timed by its PCR.
pub fn from_file(path: &str) -> Result<Vec<(u64, Vec<u8>)>, String> {
    let b = std::fs::read(path).map_err(|e| format!("{path}: {e}"))?;
    Ok(Pacer::new(&b, 7).chunks().to_vec())
}

/// Send `chunks` on `link`, each at its time from now; then wait `linger_ms` and close.
pub async fn pace(link: &Link, chunks: &[(u64, Vec<u8>)], linger_ms: u64) -> Result<(), String> {
    let tx = link.send.clone().ok_or("the link does not send")?;
    let start = tokio::time::Instant::now();
    for (at, c) in chunks {
        tokio::time::sleep_until(start + std::time::Duration::from_micros(*at)).await;
        if tx.send(c.clone()).await.is_err() {
            return Err(link.closed_reason().unwrap_or_else(|| "the link ended".into()));
        }
    }
    tokio::time::sleep(std::time::Duration::from_millis(linger_ms)).await;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn urls_read() {
        let u = Url::parse("srt://127.0.0.1:9000?streamid=%23!%3A%3Ar%3Dx,m%3Dpublish&passphrase=0123456789&latency=200").unwrap();
        assert_eq!((u.scheme.as_str(), u.port), ("srt", 9000));
        assert_eq!(u.get("streamid"), Some("#!::r=x,m=publish"));
        assert_eq!(u.num::<u16>("latency", 120), 200);
        assert!(!u.listens());
        assert!(Url::parse("srt://:9000").unwrap().listens());
        assert!(Url::parse("rist://[::1]:5000?profile=main").unwrap().addr().unwrap().is_ipv6());
    }

    #[test]
    fn a_synthetic_stream_is_timed() {
        let c = synthetic(2, 2000);
        assert!(c.len() > 100);
        assert!(c.last().unwrap().0 > 1_800_000);
    }
}