arcly-stream 0.4.0

//! Native RTSP ingest handler (feature `rtsp`).
//!
//! Pulls live media from RTSP sources — IP cameras, hardware encoders, and
//! restreamers — and bridges it onto the engine bus through the standard
//! [`InboundProtocol`] seam. The handler acts as an **RTSP client**: for each
//! configured [`RtspSource`] it drives the session state machine
//! (`OPTIONS → DESCRIBE → SETUP → PLAY → TEARDOWN`), then depacketizes the RTP
//! media into Annex-B access units and publishes them.
//!
//! # Transport
//!
//! TCP-interleaved transport (RTP-over-RTSP, RFC 2326 §10.12) is the default and
//! the most camera-compatible: media and control share the one TCP connection,
//! so it traverses NAT and firewalls that block the classic UDP transport. The
//! [interleaved framing][InterleavedFrame] (`$ channel length …`) is parsed
//! here; the RTP payloads feed the shared [`H264Depacketizer`].
//!
//! [`InboundProtocol`]: crate::inbound::InboundProtocol
//! [`H264Depacketizer`]: crate::protocol::rtp::H264Depacketizer
//!
//! # Async behavior & teardown
//!
//! [`serve`](crate::inbound::InboundProtocol::serve) spawns one pull task per
//! source and supervises them until `shutdown` fires, at which point each task
//! issues `TEARDOWN` and releases its [`PublishSession`](crate::inbound::PublishSession). A source that drops is
//! retried with backoff so a flaky camera link self-heals.
//!
//! # Scope
//!
//! The message, SDP, and interleaved-framing parsers are complete and unit
//! tested. **Digest and Basic authentication** are handled natively: embed
//! credentials in the URL (`rtsp://user:pass@host/path`) and the client answers a
//! `401` challenge. ONVIF device discovery remains out of scope.

mod auth;
mod egress;
mod message;
mod sdp;

pub use egress::RtspServer;
pub use message::{InterleavedFrame, RtspMethod, RtspRequest, RtspResponse};
pub use sdp::{MediaDescription, Sdp};

use crate::inbound::{InboundProtocol, IngestContext};
use crate::protocol::rtp::{AacDepacketizer, H264Depacketizer, RtpHeader};
use crate::{CodecId, MediaFrame, Result, StreamKey};
use async_trait::async_trait;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
use tracing::{debug, warn};

/// One RTSP source to pull and the stream key it publishes to.
#[derive(Debug, Clone)]
pub struct RtspSource {
    /// Absolute `rtsp://` URL (credentials, if any, embedded by the host).
    pub url: String,
    /// Engine stream key the pulled media is published under.
    pub key: StreamKey,
}

impl RtspSource {
    /// A source pulling `url` and publishing it as `key`.
    pub fn new(url: impl Into<String>, key: StreamKey) -> Self {
        Self {
            url: url.into(),
            key,
        }
    }
}

/// RTSP ingest worker — pulls every configured [`RtspSource`] concurrently.
#[derive(Debug)]
pub struct RtspHandler {
    sources: Vec<RtspSource>,
    retry_backoff: Duration,
}

impl Default for RtspHandler {
    fn default() -> Self {
        Self::new()
    }
}

impl RtspHandler {
    /// A handler with no sources. Add them with [`source`](Self::source).
    pub fn new() -> Self {
        Self {
            sources: Vec::new(),
            retry_backoff: Duration::from_secs(3),
        }
    }

    /// Register a source to pull.
    pub fn source(mut self, source: RtspSource) -> Self {
        self.sources.push(source);
        self
    }

    /// Override the reconnect backoff applied after a source drops (default 3s).
    pub fn retry_backoff(mut self, backoff: Duration) -> Self {
        self.retry_backoff = backoff;
        self
    }

    /// Pull one source until `shutdown`, reconnecting on failure. Owned
    /// arguments so each source runs on its own spawned task.
    async fn run_source(
        source: RtspSource,
        ctx: IngestContext,
        shutdown: CancellationToken,
        backoff: Duration,
    ) {
        loop {
            if shutdown.is_cancelled() {
                return;
            }
            // Race the whole pull session against shutdown. `pull_once` performs
            // blocking network I/O (TCP connect, request/response, the RTP read
            // loop); a peer that silently drops packets would otherwise wedge the
            // connect/read for the OS timeout and make graceful shutdown hang.
            // Cancelling drops the future, closing the socket immediately.
            tokio::select! {
                _ = shutdown.cancelled() => return,
                res = Self::pull_once(&source, &ctx, &shutdown) => {
                    if let Err(e) = res {
                        warn!(url = %source.url, error = %e, "rtsp source dropped; will retry");
                    }
                }
            }
            tokio::select! {
                _ = shutdown.cancelled() => return,
                _ = tokio::time::sleep(backoff) => {}
            }
        }
    }

    /// One full pull session for `source`. Connects, negotiates, and streams
    /// interleaved RTP until the link drops or `shutdown` fires.
    async fn pull_once(
        source: &RtspSource,
        ctx: &IngestContext,
        shutdown: &CancellationToken,
    ) -> Result<()> {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpStream;

        // The host may embed credentials in the URL (rtsp://user:pass@host/…);
        // strip them off the wire and answer any 401 Digest/Basic challenge.
        let (url, creds) = auth::split_userinfo(&source.url);
        let mut challenge: Option<auth::Challenge> = None;

        let (host, port) = message::host_port(&url)
            .ok_or_else(|| crate::StreamError::protocol("malformed rtsp url"))?;
        let mut stream = TcpStream::connect((host.as_str(), port)).await?;
        let mut cseq = 1u32;

        // OPTIONS → DESCRIBE → SETUP → PLAY (each retried once with auth on a 401).
        Self::send_authed(
            &mut stream,
            RtspMethod::Options,
            &url,
            &mut cseq,
            &[],
            &creds,
            &mut challenge,
        )
        .await?;

        let describe = Self::send_authed(
            &mut stream,
            RtspMethod::Describe,
            &url,
            &mut cseq,
            &[("Accept", "application/sdp")],
            &creds,
            &mut challenge,
        )
        .await?;
        let sdp = Sdp::parse(&describe.body);
        debug!(url = %url, media = sdp.media.len(), "rtsp DESCRIBE parsed");

        // SETUP the first video track over interleaved channels 0/1.
        let setup_url = sdp.first_video_control(&url).unwrap_or_else(|| url.clone());
        let setup = Self::send_authed(
            &mut stream,
            RtspMethod::Setup,
            &setup_url,
            &mut cseq,
            &[("Transport", "RTP/AVP/TCP;unicast;interleaved=0-1")],
            &creds,
            &mut challenge,
        )
        .await?;
        let session_id = message::session_id(&setup).unwrap_or_default();

        // If the session offers an AAC audio track, SETUP it on channels 2/3.
        let mut audio_clock = None;
        if sdp.has_aac_audio() {
            if let Some(audio_url) = sdp.first_audio_control(&url) {
                Self::send_authed(
                    &mut stream,
                    RtspMethod::Setup,
                    &audio_url,
                    &mut cseq,
                    &[
                        ("Transport", "RTP/AVP/TCP;unicast;interleaved=2-3"),
                        ("Session", &session_id),
                    ],
                    &creds,
                    &mut challenge,
                )
                .await?;
                audio_clock = sdp
                    .media
                    .iter()
                    .find(|m| m.media == "audio")
                    .and_then(|m| m.clock_rate)
                    .or(Some(48_000));
                debug!(url = %url, "rtsp AAC audio track set up on ch 2/3");
            }
        }

        Self::send_authed(
            &mut stream,
            RtspMethod::Play,
            &url,
            &mut cseq,
            &[("Session", &session_id)],
            &creds,
            &mut challenge,
        )
        .await?;

        // Stream interleaved RTP → depacketize → publish.
        let session = ctx.open_publish(source.key.clone()).await?;
        let mut depack = H264Depacketizer::new();
        let (size_len, index_len) = sdp.audio_aac_lengths();
        let aac = AacDepacketizer::with_lengths(size_len, index_len);
        let mut buf = Vec::with_capacity(64 * 1024);
        let mut read = [0u8; 16 * 1024];

        loop {
            tokio::select! {
                _ = shutdown.cancelled() => break,
                n = stream.read(&mut read) => {
                    let n = n?;
                    if n == 0 { break; }
                    buf.extend_from_slice(&read[..n]);
                    Self::drain_interleaved(&mut buf, &mut depack, &aac, audio_clock, &session)?;
                }
            }
        }

        // Best-effort TEARDOWN, then release the publish slot.
        let _ = Self::send_authed(
            &mut stream,
            RtspMethod::Teardown,
            &url,
            &mut cseq,
            &[("Session", &session_id)],
            &creds,
            &mut challenge,
        )
        .await;
        let _ = stream.shutdown().await;
        session.finish().await
    }

    /// Send one RTSP request and return its response, transparently answering a
    /// `401` with a Digest/Basic `Authorization` and retrying once. `cseq` is
    /// advanced for every request actually sent.
    #[allow(clippy::too_many_arguments)]
    async fn send_authed<S>(
        stream: &mut S,
        method: RtspMethod,
        uri: &str,
        cseq: &mut u32,
        headers: &[(&str, &str)],
        creds: &Option<(String, String)>,
        challenge: &mut Option<auth::Challenge>,
    ) -> Result<RtspResponse>
    where
        S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
    {
        // First attempt — include auth if we already learned a challenge.
        let auth1 = match (challenge.as_ref(), creds.as_ref()) {
            (Some(ch), Some((u, p))) => Some(auth::authorization(ch, u, p, method.as_str(), uri)),
            _ => None,
        };
        let mut hdrs: Vec<(&str, &str)> = headers.to_vec();
        if let Some(a) = &auth1 {
            hdrs.push(("Authorization", a));
        }
        message::write_request(stream, method, uri, *cseq, &hdrs).await?;
        *cseq += 1;
        let resp = message::read_response(stream).await?;

        // Not a challenge (or we have no credentials) → return as-is.
        if resp.status != 401 || creds.is_none() {
            return Ok(resp);
        }
        let Some(ch) = resp
            .header("WWW-Authenticate")
            .and_then(auth::parse_challenge)
        else {
            return Ok(resp);
        };
        *challenge = Some(ch);

        // Retry once with the freshly-parsed challenge.
        let (u, p) = creds.as_ref().unwrap();
        let auth2 = auth::authorization(challenge.as_ref().unwrap(), u, p, method.as_str(), uri);
        let mut hdrs: Vec<(&str, &str)> = headers.to_vec();
        hdrs.push(("Authorization", &auth2));
        message::write_request(stream, method, uri, *cseq, &hdrs).await?;
        *cseq += 1;
        message::read_response(stream).await
    }

    /// Consume whole interleaved frames from `buf`, depacketizing channel-0 video
    /// RTP and (when `audio_clock` is set) channel-2 AAC RTP into frames and
    /// publishing them. Leaves any partial frame in `buf`.
    fn drain_interleaved(
        buf: &mut Vec<u8>,
        depack: &mut H264Depacketizer,
        aac: &AacDepacketizer,
        audio_clock: Option<u32>,
        session: &crate::inbound::PublishSession,
    ) -> Result<()> {
        let mut consumed = 0;
        while let Some((frame, len)) = InterleavedFrame::parse(&buf[consumed..]) {
            consumed += len;
            let Some(header) = RtpHeader::parse(frame.payload) else {
                continue;
            };
            let payload = &frame.payload[header.payload_offset..];
            // Channel 0 = video RTP, channel 2 = audio RTP; odd channels are
            // RTCP, ignored on ingest.
            match frame.channel {
                0 => {
                    match depack.push(payload, header.marker, header.timestamp, header.sequence) {
                        Ok(Some(au)) => {
                            let pts = (au.timestamp / 90) as i64; // 90 kHz → ms
                            let mf = MediaFrame::new_video(
                                pts,
                                pts,
                                au.data,
                                CodecId::H264,
                                au.keyframe,
                            );
                            let _ = session.publish_frame(mf)?;
                        }
                        Ok(None) => {}
                        Err(e) => debug!(?e, "rtp depacketize skip"),
                    }
                }
                2 => {
                    if let Some(clock) = audio_clock {
                        match aac.push(payload) {
                            Ok(units) => {
                                for au in units {
                                    let pts =
                                        (header.timestamp as i64 * 1000) / clock.max(1) as i64;
                                    let mf = MediaFrame::new_audio(pts, au, CodecId::AAC);
                                    let _ = session.publish_frame(mf)?;
                                }
                            }
                            Err(e) => debug!(?e, "aac depacketize skip"),
                        }
                    }
                }
                _ => {}
            }
        }
        buf.drain(..consumed);
        Ok(())
    }
}

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

    async fn serve(&self, ctx: IngestContext, shutdown: CancellationToken) -> Result<()> {
        // Pull every source concurrently on its own task; await them all so the
        // worker only returns once every source has drained on shutdown.
        let mut tasks = tokio::task::JoinSet::new();
        for source in &self.sources {
            tasks.spawn(Self::run_source(
                source.clone(),
                ctx.clone(),
                shutdown.clone(),
                self.retry_backoff,
            ));
        }
        while tasks.join_next().await.is_some() {}
        Ok(())
    }
}

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

    #[test]
    fn source_builder_sets_url_and_key() {
        let s = RtspSource::new("rtsp://cam/stream", StreamKey::new("live", "cam1"));
        assert_eq!(s.url, "rtsp://cam/stream");
        assert_eq!(s.key.stream_id.as_str(), "cam1");
    }

    #[test]
    fn handler_collects_sources() {
        let h = RtspHandler::new()
            .source(RtspSource::new("rtsp://a/1", StreamKey::new("live", "a")))
            .source(RtspSource::new("rtsp://b/2", StreamKey::new("live", "b")));
        assert_eq!(h.sources.len(), 2);
        assert_eq!(h.name(), "rtsp");
    }
}