arcly-stream 0.7.2

//! SRT **caller egress** — dial an SRT listener and push an MPEG-TS feed.
//!
//! The egress counterpart to the [listener ingest](super::SrtHandler):
//! [`SrtCaller`] connects to a remote SRT listener (`srt://host:port` in caller
//! mode), runs the caller handshake (`INDUCTION` → `CONCLUSION`), then wraps an
//! MPEG-TS byte stream into SRT data packets and sends them.
//!
//! The transport stays **MPEG-TS-agnostic of the engine**: the caller consumes
//! already-muxed TS bytes from an [`mpsc`] channel, so the host pairs it with a
//! [`MpegTsMuxer`](crate::packager::MpegTsMuxer) to get a full frame → SRT path
//! without coupling this module to the packager.
//!
//! Scope: unencrypted, and the minimal HSv5 handshake (no `HSREQ`/`KMREQ`
//! extension blocks) — it interoperates with this crate's own listener; full
//! extension interop with third-party SRT stacks is a follow-up.

use std::net::SocketAddr;
use std::time::{Duration, Instant};

use tokio::net::UdpSocket;
use tokio::sync::mpsc;
use tokio::time::timeout;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info};

use super::arq::{self, SendBuffer};
use super::handshake::{caller_conclusion, caller_induction, SrtHandshake};
use super::packet::{build_data_packet, ControlType};
use crate::{Result, StreamError};

/// 7 × 188-byte TS packets = the classic 1316-byte SRT payload.
const TS_BYTES_PER_DATAGRAM: usize = 7 * 188;
/// Bound the handshake so a dead listener can't wedge the caller.
const SETUP_TIMEOUT: Duration = Duration::from_secs(5);
/// Retransmission window: packets kept available for NAK-driven replay.
const SEND_WINDOW: usize = 2048;

/// Pushes an MPEG-TS feed to a remote SRT listener in caller mode.
pub struct SrtCaller {
    addr: SocketAddr,
    socket_id: u32,
    /// Passphrase enabling AES-CTR encryption (feature `srt-encrypt`).
    #[cfg(feature = "srt-encrypt")]
    passphrase: Option<String>,
    /// SEK length in bytes (16/24/32) when encrypting (feature `srt-encrypt`).
    #[cfg(feature = "srt-encrypt")]
    key_len: usize,
}

impl SrtCaller {
    /// A caller targeting the listener at `addr`.
    pub fn new(addr: SocketAddr) -> Self {
        // A non-zero, stable-ish local socket id (real callers randomize).
        let socket_id = (0x4152_434C ^ (addr.port() as u32).rotate_left(16)) | 1;
        Self {
            addr,
            socket_id,
            #[cfg(feature = "srt-encrypt")]
            passphrase: None,
            #[cfg(feature = "srt-encrypt")]
            key_len: 16,
        }
    }

    /// Encrypt the feed with AES-CTR under `passphrase` (default AES-128; pair
    /// with [`with_key_len`](Self::with_key_len)). Feature `srt-encrypt`.
    #[cfg(feature = "srt-encrypt")]
    pub fn with_passphrase(mut self, passphrase: impl Into<String>) -> Self {
        self.passphrase = Some(passphrase.into());
        self
    }

    /// Set the AES key length in bytes (16, 24, or 32). Feature `srt-encrypt`.
    #[cfg(feature = "srt-encrypt")]
    pub fn with_key_len(mut self, key_len: usize) -> Self {
        self.key_len = key_len;
        self
    }

    /// Connect, handshake, and forward MPEG-TS chunks from `ts` until the channel
    /// closes or `shutdown` fires. Each inbound chunk is split into
    /// 1316-byte SRT datagrams.
    pub async fn run(
        self,
        mut ts: mpsc::Receiver<bytes::Bytes>,
        shutdown: CancellationToken,
    ) -> Result<()> {
        let sock = UdpSocket::bind(("0.0.0.0", 0)).await?;
        sock.connect(self.addr).await?;
        let mut buf = [0u8; 1500];

        // INDUCTION → cookie.
        sock.send(&caller_induction(self.socket_id, 0)).await?;
        let n = timeout(SETUP_TIMEOUT, sock.recv(&mut buf))
            .await
            .map_err(|_| StreamError::protocol("srt induction timed out"))??;
        let induction = SrtHandshake::parse(&buf[..n])
            .ok_or_else(|| StreamError::protocol("malformed srt induction response"))?;

        // Optional encryption: generate key material and send it in the
        // CONCLUSION's KMREQ extension; the listener unwraps it with the same
        // passphrase. Each payload is then AES-CTR'd before framing.
        #[cfg(feature = "srt-encrypt")]
        let key_material = self
            .passphrase
            .as_ref()
            .map(|_| super::keymaterial::KeyMaterial::generate(self.key_len));

        // CONCLUSION (echo the cookie, carry KMREQ if encrypting) → agreement.
        #[cfg(feature = "srt-encrypt")]
        let conclusion = match (&self.passphrase, &key_material) {
            (Some(pass), Some(km)) => super::handshake::caller_conclusion_encrypted(
                self.socket_id,
                0,
                induction.cookie,
                self.key_len,
                &km.to_message(pass.as_bytes()),
            ),
            _ => caller_conclusion(self.socket_id, 0, induction.cookie),
        };
        #[cfg(not(feature = "srt-encrypt"))]
        let conclusion = caller_conclusion(self.socket_id, 0, induction.cookie);

        sock.send(&conclusion).await?;
        let n = timeout(SETUP_TIMEOUT, sock.recv(&mut buf))
            .await
            .map_err(|_| StreamError::protocol("srt conclusion timed out"))??;
        let agreement = SrtHandshake::parse(&buf[..n])
            .ok_or_else(|| StreamError::protocol("malformed srt conclusion response"))?;
        // The data packets' destination is the socket id the listener returned.
        let dest = agreement.socket_id;
        info!(addr = %self.addr, dest, "srt caller connected");

        // The key flag (`KK`) stamped on every data packet: 1 (even key) when
        // encrypting, 0 otherwise.
        #[cfg(feature = "srt-encrypt")]
        let key_flag: u8 = if key_material.is_some() { 1 } else { 0 };
        #[cfg(not(feature = "srt-encrypt"))]
        let key_flag: u8 = 0;

        let start = Instant::now();
        let mut seq = 0u32;
        let mut msg = 0u32;
        let mut send_buf = SendBuffer::new(SEND_WINDOW);
        loop {
            tokio::select! {
                _ = shutdown.cancelled() => break,

                // Control feedback from the listener: retransmit on NAK, release
                // the send buffer on ACK.
                r = sock.recv(&mut buf) => {
                    let Ok(n) = r else { continue };
                    let dg = &buf[..n];
                    match arq::control_type(dg) {
                        Some(ControlType::Nak) => {
                            for pkt in send_buf.retransmit(&arq::parse_nak(dg)) {
                                let _ = sock.send(&pkt).await;
                            }
                        }
                        Some(ControlType::Ack) => {
                            if let Some(ack) = arq::parse_ack(dg) {
                                send_buf.acknowledge(ack);
                            }
                            // ACKACK back so the listener can retire the ACK.
                            if let Some(ack_no) = arq::ack_seqno(dg) {
                                let ts_us = start.elapsed().as_micros() as u32;
                                let _ = sock.send(&arq::build_ackack(ack_no, ts_us, dest)).await;
                            }
                        }
                        _ => {}
                    }
                }

                chunk = ts.recv() => match chunk {
                    Some(bytes) => {
                        for piece in bytes.chunks(TS_BYTES_PER_DATAGRAM) {
                            let ts_us = start.elapsed().as_micros() as u32;
                            let payload = piece.to_vec();
                            #[cfg(feature = "srt-encrypt")]
                            let mut payload = payload;
                            #[cfg(feature = "srt-encrypt")]
                            if let Some(km) = &key_material {
                                km.transform(seq, &mut payload);
                            }
                            let pkt =
                                build_data_packet(seq, msg, ts_us, dest, key_flag, false, &payload);
                            if sock.send(&pkt).await.is_err() {
                                return Ok(()); // listener gone
                            }
                            send_buf.record(seq, pkt);
                            seq = seq.wrapping_add(1) & 0x7FFF_FFFF;
                            msg = msg.wrapping_add(1) & 0x03FF_FFFF;
                        }
                    }
                    None => break, // source closed
                }
            }
        }
        debug!(addr = %self.addr, "srt caller finished");
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::srt::{handshake::respond, SrtPacket};

    #[tokio::test]
    async fn caller_handshakes_and_sends_data_over_loopback() {
        // A fake listener: this crate's own `respond` for the handshake, then it
        // captures the first data packet — a real end-to-end UDP loopback of the
        // caller flow (handshake + data framing).
        let listener = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let (tx, rx) = mpsc::channel(4);
        let shutdown = CancellationToken::new();
        let sh = shutdown.clone();
        let caller = SrtCaller::new(addr);
        let caller_task = tokio::spawn(async move { caller.run(rx, sh).await });

        let mut buf = [0u8; 1500];
        // Induction.
        let (n, peer) = listener.recv_from(&mut buf).await.unwrap();
        let reply = respond(&buf[..n]).unwrap();
        listener.send_to(&reply, peer).await.unwrap();
        // Conclusion.
        let (n, peer) = listener.recv_from(&mut buf).await.unwrap();
        let reply = respond(&buf[..n]).unwrap();
        listener.send_to(&reply, peer).await.unwrap();

        // Feed one TS datagram; expect an SRT data packet on the wire.
        tx.send(bytes::Bytes::from(vec![0x47u8; TS_BYTES_PER_DATAGRAM]))
            .await
            .unwrap();
        let (n, _) = timeout(Duration::from_secs(5), listener.recv_from(&mut buf))
            .await
            .expect("data packet arrived")
            .unwrap();
        assert!(
            matches!(SrtPacket::parse(&buf[..n]), Some(SrtPacket::Data { .. })),
            "caller sent an SRT data packet after the handshake"
        );

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