arcly-stream 0.8.2

//! Native SRT ingest handler (feature `srt`).
//!
//! Accepts SRT (Secure Reliable Transport) connections in **listener** mode,
//! parses the control/data packet stream, and extracts the encapsulated
//! MPEG-TS payload into elementary media frames that flow onto the engine bus
//! through the standard [`InboundProtocol`] seam.
//!
//! [`InboundProtocol`]: crate::inbound::InboundProtocol
//!
//! # Pipeline
//!
//! ```text
//!  UDP datagram ─▶ SrtPacket::parse ─▶ data payload ─▶ TsDemuxer ─▶ AccessUnit
//!                       │                                              │
//!                       └─ handshake (induction/conclusion)           ▼
//!                                                              PublishSession
//! ```
//!
//! The SRT [packet header][SrtPacket] (SRT RFC, draft-sharabayko-srt) and the
//! [`SrtHandshake`] induction/conclusion exchange are parsed here; the MPEG-TS
//! bytes ride in SRT *data* packets and are demuxed by [`TsDemuxer`].
//!
//! # Scope
//!
//! Loss recovery is modeled with a full **ARQ window**: the listener reorders by
//! sequence number, NAKs gaps, and periodically ACKs; the caller keeps a
//! retransmission buffer and replays NAK'd packets. With the `srt-encrypt`
//! feature, **AES-CTR** media encryption and the `KMREQ`/`KMRSP` key-material
//! exchange are supported via a pure-Rust crypto core — no dependency, still
//! `#![forbid(unsafe_code)]`. Without that feature an encrypted feed is rejected
//! at handshake.
//!
//! A caller's **Stream ID** (`streamid`) is honored: the HSv5 `SRT_CMD_SID`
//! extension a libsrt caller sends to select its stream is parsed (by the
//! internal `handshake::stream_id`) and used as the published stream id within
//! the configured app namespace (falling back to the handler's key when absent).
//!
//! Interop note: the handshake/ARQ/KM/`streamid` wire formats follow
//! draft-sharabayko-srt and interoperate with this crate's own caller/listener
//! and with libsrt callers that send a Stream ID; exhaustive option negotiation
//! with third-party SRT stacks is a follow-up.

mod arq;
mod conn;
#[cfg(feature = "srt-encrypt")]
mod crypto;
mod egress;
mod handshake;
#[cfg(feature = "srt-encrypt")]
mod keymaterial;
mod packet;

pub use egress::SrtCaller;
pub use handshake::{HandshakeType, SrtHandshake};
pub use packet::{ControlType, SrtPacket};
// The MPEG-TS demuxer now lives in the shared `protocol::tsdemux` module (also
// used by `udp` ingest); re-exported here so existing paths keep working.
pub use crate::protocol::tsdemux::{TsDemuxer, TsPayload, TsTrackKind};

use crate::bus::PlaybackRegistry;
use crate::inbound::{InboundProtocol, IngestContext};
use crate::{Result, StreamKey};
use async_trait::async_trait;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};

/// Out-of-order reorder window before a stuck gap is force-skipped.
const ARQ_WINDOW: usize = 1024;
/// Cadence of ACK/NAK control feedback and buffer-relief checks.
const CONTROL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(10);
/// Inbound-datagram queue depth per connection task.
const CONN_QUEUE: usize = 512;

/// SRT **server**: one UDP listener that demultiplexes many concurrent SRT
/// connections by peer address and routes each by its `streamid` mode —
/// `m=publish` callers are *ingested* (published onto the bus), `m=request`
/// callers are *served* (a stream muxed back to them as MPEG-TS). Each
/// connection runs as its own task with independent ARQ/encryption state.
#[derive(Clone)]
pub struct SrtHandler {
    bind: SocketAddr,
    /// Default stream key when a caller omits `streamid` (publish only).
    key: StreamKey,
    /// Playback side, required to serve `m=request` (egress) callers. `None`
    /// disables egress (publish-only listener).
    playback: Option<Arc<dyn PlaybackRegistry>>,
    /// Egress gate (per-app toggle + play token) for `m=request`. `None` = open.
    gate: Option<crate::auth::EgressGate>,
    /// Pre-shared passphrase enabling AES-CTR (feature `srt-encrypt`).
    #[cfg(feature = "srt-encrypt")]
    passphrase: Option<String>,
}

/// Resolve a caller's `streamid` to a [`StreamKey`]. libsrt convention prefixes
/// the resource with `#!::r=`; an embedded `/` selects `app/stream`, otherwise
/// the value is a stream within the listener's default app (`default.app`). This
/// lets a single listener route to any `app/stream` (e.g.
/// `streamid=#!::r=live/test` or `streamid=live/test`).
fn resolve_streamid(default: &StreamKey, sid: &str) -> StreamKey {
    // Strip a libsrt `#!::{...}` wrapper, taking the `r=` (resource) field.
    let resource = sid
        .strip_prefix("#!::")
        .and_then(|rest| {
            rest.split(',')
                .find_map(|kv| kv.strip_prefix("r=").or_else(|| kv.strip_prefix("s=")))
        })
        .unwrap_or(sid)
        .trim_matches('/');
    match resource.split_once('/') {
        Some((app, stream)) if !app.is_empty() && !stream.is_empty() => StreamKey::new(app, stream),
        _ => StreamKey::new(default.app.clone(), resource),
    }
}

impl SrtHandler {
    /// A listener bound to `bind`. Publish-only until [`with_playback`](Self::with_playback) is set.
    /// `key` is the fallback stream key when a `m=publish` caller omits `streamid`.
    pub fn new(bind: SocketAddr, key: StreamKey) -> Self {
        Self {
            bind,
            key,
            playback: None,
            gate: None,
            #[cfg(feature = "srt-encrypt")]
            passphrase: None,
        }
    }

    /// Enable serving `m=request` (egress) callers from `playback`.
    pub fn with_playback(mut self, playback: Arc<dyn PlaybackRegistry>) -> Self {
        self.playback = Some(playback);
        self
    }

    /// Gate `m=request` (egress) callers through `gate` (per-app toggle + token).
    pub fn with_gate(mut self, gate: crate::auth::EgressGate) -> Self {
        self.gate = Some(gate);
        self
    }

    /// Pre-shared passphrase for AES-CTR (both encrypted ingest and egress).
    /// Feature `srt-encrypt`.
    #[cfg(feature = "srt-encrypt")]
    pub fn with_passphrase(mut self, passphrase: impl Into<String>) -> Self {
        self.passphrase = Some(passphrase.into());
        self
    }
}

#[async_trait]
impl InboundProtocol for SrtHandler {
    fn name(&self) -> &'static str {
        "srt"
    }

    async fn serve(&self, ctx: IngestContext, shutdown: CancellationToken) -> Result<()> {
        use tokio::net::UdpSocket;
        use tokio::sync::mpsc;

        let socket = Arc::new(UdpSocket::bind(self.bind).await?);
        info!(bind = %self.bind, "srt server listening");

        // One inbound queue per peer; a connection task drains it and owns all
        // per-connection state (handshake, ARQ, encryption, publish/play).
        let mut conns: HashMap<SocketAddr, mpsc::Sender<Vec<u8>>> = HashMap::new();
        let mut buf = vec![0u8; 1500];

        loop {
            tokio::select! {
                _ = shutdown.cancelled() => break,
                r = socket.recv_from(&mut buf) => {
                    let (n, from) = match r {
                        Ok(v) => v,
                        Err(e) => { warn!(error = %e, "srt recv failed"); continue; }
                    };
                    let datagram = buf[..n].to_vec();

                    // Route to the peer's connection task if it exists.
                    if let Some(tx) = conns.get(&from) {
                        if tx.try_send(datagram).is_err() {
                            // Task gone or backlogged past the queue: drop the peer;
                            // a fresh handshake will re-establish it.
                            conns.remove(&from);
                        }
                        continue;
                    }

                    // New peer: only a handshake (induction) starts a connection.
                    if !conn::is_handshake(&datagram) {
                        continue;
                    }
                    let (tx, rx) = mpsc::channel(CONN_QUEUE);
                    let cfg = conn::ConnConfig {
                        socket: socket.clone(),
                        peer: from,
                        ctx: ctx.clone(),
                        playback: self.playback.clone(),
                        gate: self.gate.clone(),
                        default_key: self.key.clone(),
                        #[cfg(feature = "srt-encrypt")]
                        passphrase: self.passphrase.clone(),
                        shutdown: shutdown.clone(),
                    };
                    tokio::spawn(conn::run(cfg, rx));
                    let _ = tx.try_send(datagram);
                    conns.insert(from, tx);
                }
            }
        }
        Ok(())
    }
}

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

    #[test]
    fn resolve_streamid_parses_app_and_stream() {
        let def = StreamKey::new("live", "srt");
        // bare app/stream
        let k = resolve_streamid(&def, "show/cam1");
        assert_eq!(k.app.as_str(), "show");
        assert_eq!(k.stream_id.as_str(), "cam1");
        // libsrt resource wrapper
        let k = resolve_streamid(&def, "#!::r=show/cam1,m=publish");
        assert_eq!(k.app.as_str(), "show");
        assert_eq!(k.stream_id.as_str(), "cam1");
        // stream-only falls back to the listener's default app
        let k = resolve_streamid(&def, "cam9");
        assert_eq!(k.app.as_str(), "live");
        assert_eq!(k.stream_id.as_str(), "cam9");
    }

    #[test]
    fn handler_reports_name_and_key() {
        let h = SrtHandler::new(
            "127.0.0.1:9000".parse().unwrap(),
            StreamKey::new("live", "feed"),
        );
        assert_eq!(h.name(), "srt");
        assert_eq!(h.key.stream_id.as_str(), "feed");
    }

    use super::arq::Receiver;
    use super::packet::SrtPacket;
    use std::collections::HashSet;
    use std::time::Duration;
    use tokio::net::UdpSocket;
    use tokio::sync::mpsc;
    use tokio::time::timeout;

    /// End-to-end lossy-link recovery: a caller streams packets through a test
    /// listener that drops two of them on first sight and NAKs; the caller must
    /// retransmit so the receiver delivers every payload, in order.
    #[tokio::test]
    async fn arq_recovers_dropped_packets() {
        let listener = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let (tx, rx) = mpsc::channel(32);
        let shutdown = CancellationToken::new();
        let caller_task = tokio::spawn(SrtCaller::new(addr).run(rx, shutdown.clone()));

        // Handshake: induction then conclusion, both echoed by `respond`.
        let mut buf = [0u8; 2048];
        let (n, peer) = listener.recv_from(&mut buf).await.unwrap();
        let dest = SrtHandshake::parse(&buf[..n]).unwrap().socket_id;
        listener
            .send_to(&handshake::respond(&buf[..n]).unwrap(), peer)
            .await
            .unwrap();
        let (n, peer) = listener.recv_from(&mut buf).await.unwrap();
        listener
            .send_to(&handshake::respond(&buf[..n]).unwrap(), peer)
            .await
            .unwrap();

        // Feed 12 uniquely-tagged payloads.
        let count = 12usize;
        for i in 0..count {
            tx.send(bytes::Bytes::from(vec![i as u8; 100]))
                .await
                .unwrap();
        }

        let mut receiver = Receiver::new(64);
        let mut delivered: Vec<Vec<u8>> = Vec::new();
        let mut dropped_once: HashSet<u32> = HashSet::new();
        while delivered.len() < count {
            let (n, from) = timeout(Duration::from_secs(5), listener.recv_from(&mut buf))
                .await
                .expect("packet within timeout")
                .unwrap();
            let SrtPacket::Data {
                sequence,
                payload_offset,
                ..
            } = SrtPacket::parse(&buf[..n]).unwrap()
            else {
                continue;
            };
            // Drop seq 3 and 8 the first time they appear, and NAK them.
            if (sequence == 3 || sequence == 8) && dropped_once.insert(sequence) {
                let nak = arq::build_nak(&[(sequence, sequence)], 0, dest);
                listener.send_to(&nak, from).await.unwrap();
                continue;
            }
            let payload = buf[payload_offset..n].to_vec();
            delivered.extend(receiver.push(sequence, payload));
        }

        shutdown.cancel();
        let _ = caller_task.await;

        let expected: Vec<Vec<u8>> = (0..count).map(|i| vec![i as u8; 100]).collect();
        assert_eq!(delivered, expected, "every payload recovered, in order");
        assert_eq!(dropped_once.len(), 2, "both losses were actually injected");
    }

    /// End-to-end encryption: an encrypted caller completes the KMREQ exchange,
    /// and its on-wire payloads decrypt (only) with the recovered key material.
    #[cfg(feature = "srt-encrypt")]
    #[tokio::test]
    async fn encrypted_caller_payload_decrypts_with_recovered_km() {
        let pass = "swordfish-correct-horse";
        let listener = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let (tx, rx) = mpsc::channel(8);
        let shutdown = CancellationToken::new();
        let caller_task = tokio::spawn(
            SrtCaller::new(addr)
                .with_passphrase(pass)
                .run(rx, shutdown.clone()),
        );

        let mut buf = [0u8; 2048];
        // Induction (no KM yet).
        let (n, peer) = listener.recv_from(&mut buf).await.unwrap();
        let (reply, none) = handshake::respond_with_km(&buf[..n], pass.as_bytes()).unwrap();
        assert!(none.is_none(), "induction carries no key material");
        listener.send_to(&reply, peer).await.unwrap();
        // Conclusion carries KMREQ; recover the SEK with the right passphrase.
        let (n, peer) = listener.recv_from(&mut buf).await.unwrap();
        let (reply, km) = handshake::respond_with_km(&buf[..n], pass.as_bytes()).unwrap();
        let km = km.expect("conclusion KMREQ yields key material");
        // A wrong passphrase must fail to unwrap the very same handshake.
        assert!(handshake::respond_with_km(&buf[..n], b"wrong").is_none());
        listener.send_to(&reply, peer).await.unwrap();

        let plain = vec![0x47u8; TS_BYTES_PER_DATAGRAM_TEST];
        tx.send(bytes::Bytes::from(plain.clone())).await.unwrap();

        let (n, _) = timeout(Duration::from_secs(5), listener.recv_from(&mut buf))
            .await
            .expect("data packet")
            .unwrap();
        let SrtPacket::Data {
            sequence,
            key_flag,
            payload_offset,
            ..
        } = SrtPacket::parse(&buf[..n]).unwrap()
        else {
            panic!("expected a data packet");
        };
        assert_eq!(key_flag, 1, "payload flagged as even-key encrypted");
        let mut wire = buf[payload_offset..n].to_vec();
        assert_ne!(wire, plain, "payload is ciphertext on the wire");
        km.transform(sequence, &mut wire);
        assert_eq!(wire, plain, "recovered key material decrypts the payload");

        shutdown.cancel();
        let _ = caller_task.await;
    }

    /// One full SRT datagram of TS payload (7 × 188).
    #[cfg(feature = "srt-encrypt")]
    const TS_BYTES_PER_DATAGRAM_TEST: usize = 7 * 188;
}