arcly-stream 0.8.2

//! RTCP packet builders and a tolerant compound parser (RFC 3550 / 4585 / 5104).
//!
//! **Build:** the keyframe-request messages an ingest sends back to a WHIP
//! publisher — PLI (Picture Loss Indication) and FIR (Full Intra Request) — plus
//! a Sender Report ([`build_sr`]) for A/V sync on transports that own RTCP
//! generation themselves.
//!
//! **Parse:** [`parse_compound`] decodes a (decrypted) RTCP compound packet into
//! [`RtcpFeedback`] messages — PLI/FIR, generic NACK, Receiver Reports, and REMB
//! — so the WHEP egress can react to viewer feedback. TWCC is left to the
//! transport's own bandwidth estimator (see
//! [`DtlsSrtpTransport::estimated_bitrate`](super::DtlsSrtpTransport::estimated_bitrate)).

/// RTCP payload type for Payload-Specific Feedback (PSFB), RFC 4585.
const PT_PSFB: u8 = 206;

/// Build a PLI (Picture Loss Indication) RTCP packet.
///
/// PLI is a PSFB message with feedback message type (FMT) 1 and no parameters
/// (RFC 4585 §6.3.1). `sender_ssrc` is this server's SSRC; `media_ssrc`
/// identifies the publisher's video stream to refresh.
pub fn build_pli(sender_ssrc: u32, media_ssrc: u32) -> Vec<u8> {
    let mut p = Vec::with_capacity(12);
    // V=2, P=0, FMT=1 (PLI).
    p.push(0x80 | 1);
    p.push(PT_PSFB);
    // Length in 32-bit words minus one: header(1) + 2 SSRC words = 3 → 2.
    p.extend_from_slice(&2u16.to_be_bytes());
    p.extend_from_slice(&sender_ssrc.to_be_bytes());
    p.extend_from_slice(&media_ssrc.to_be_bytes());
    p
}

/// Build a FIR (Full Intra Request) RTCP packet, RFC 5104 §4.3.1.
///
/// FIR is a PSFB message with FMT 4, carrying one FCI entry: the target SSRC, a
/// monotonically increasing `seq_nr` (so retransmitted FIRs are deduplicated),
/// and three reserved bytes.
pub fn build_fir(sender_ssrc: u32, media_ssrc: u32, seq_nr: u8) -> Vec<u8> {
    let mut p = Vec::with_capacity(20);
    // V=2, P=0, FMT=4 (FIR).
    p.push(0x80 | 4);
    p.push(PT_PSFB);
    // header(1) + sender(1) + media(1) + FCI(2) = 5 words → length 4.
    p.extend_from_slice(&4u16.to_be_bytes());
    p.extend_from_slice(&sender_ssrc.to_be_bytes());
    p.extend_from_slice(&0u32.to_be_bytes()); // media SSRC field is 0 for FIR
    p.extend_from_slice(&media_ssrc.to_be_bytes()); // FCI: target SSRC
    p.push(seq_nr);
    p.extend_from_slice(&[0, 0, 0]); // reserved
    p
}

/// Build a Sender Report (SR, RFC 3550 §6.4.1) with no report blocks.
///
/// An SR ties an `rtp_ts` (the stream's 90 kHz/48 kHz clock) to a wall-clock NTP
/// timestamp, which is how a receiver synchronizes audio against video. Pass the
/// 64-bit NTP time from [`ntp_now`] (or a fixed value in tests) and the running
/// `packet_count`/`octet_count` for the SSRC.
///
/// Transports that own RTCP generation themselves (e.g. the str0m backend, which
/// emits SRs for its send streams natively) do not need this; it is for the
/// in-crate / custom `DtlsSrtpTransport` egress path.
pub fn build_sr(
    sender_ssrc: u32,
    ntp: u64,
    rtp_ts: u32,
    packet_count: u32,
    octet_count: u32,
) -> Vec<u8> {
    let mut p = Vec::with_capacity(28);
    p.push(0x80); // V=2, P=0, RC=0
    p.push(PT_SR);
    p.extend_from_slice(&6u16.to_be_bytes()); // 7 words → length 6
    p.extend_from_slice(&sender_ssrc.to_be_bytes());
    p.extend_from_slice(&ntp.to_be_bytes()); // NTP MSW+LSW (64-bit)
    p.extend_from_slice(&rtp_ts.to_be_bytes());
    p.extend_from_slice(&packet_count.to_be_bytes());
    p.extend_from_slice(&octet_count.to_be_bytes());
    p
}

/// The current wall clock as a 64-bit NTP timestamp (seconds since 1900 in the
/// high 32 bits, fractional seconds in the low 32) for [`build_sr`].
pub fn ntp_now() -> u64 {
    /// Seconds between the NTP epoch (1900) and the Unix epoch (1970).
    const NTP_UNIX_OFFSET: u64 = 2_208_988_800;
    let d = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default();
    let secs = d.as_secs() + NTP_UNIX_OFFSET;
    let frac = ((d.subsec_nanos() as u64) << 32) / 1_000_000_000;
    (secs << 32) | frac
}

/// RTCP payload type for Transport-layer Feedback (RTPFB), RFC 4585.
const PT_RTPFB: u8 = 205;
/// RTCP payload type for a Receiver Report, RFC 3550.
const PT_RR: u8 = 201;
/// RTCP payload type for a Sender Report, RFC 3550.
const PT_SR: u8 = 200;

/// A single feedback message decoded from an RTCP compound packet.
///
/// RTCP packets arrive *compounded* — several stacked in one datagram — so
/// [`parse_compound`] walks the whole buffer and yields every message it
/// recognizes. Unknown payload types are skipped (their length field lets the
/// walker step over them) rather than aborting the parse.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RtcpFeedback {
    /// Picture Loss Indication (PSFB FMT 1) — the receiver wants a fresh IDR for
    /// `media_ssrc`. The egress path turns this into an upstream keyframe request.
    Pli {
        /// SSRC of the report sender (the viewer).
        sender_ssrc: u32,
        /// SSRC of the media stream the keyframe is requested for.
        media_ssrc: u32,
    },
    /// Full Intra Request (PSFB FMT 4) — like PLI but reliable/seq-numbered.
    Fir {
        /// SSRC of the report sender (the viewer).
        sender_ssrc: u32,
        /// SSRC of the media stream the IDR is requested for.
        media_ssrc: u32,
    },
    /// Generic NACK (RTPFB FMT 1, RFC 4585 §6.2.1) — the receiver is missing the
    /// listed RTP sequence numbers and wants them retransmitted.
    Nack {
        /// SSRC of the report sender (the viewer).
        sender_ssrc: u32,
        /// SSRC of the media stream the losses belong to.
        media_ssrc: u32,
        /// The lost sequence numbers, expanded from each (PID, bitmask) FCI entry.
        lost: Vec<u16>,
    },
    /// REMB — Receiver Estimated Maximum Bitrate (PSFB FMT 15, `goog-remb`). The
    /// receiver's estimate of the bitrate the path can sustain, in **bits/sec**.
    /// A direct bandwidth-estimate input for egress rate adaptation.
    Remb {
        /// SSRC of the report sender (the viewer).
        sender_ssrc: u32,
        /// Estimated maximum bitrate the path can sustain, in bits per second.
        bitrate_bps: u64,
    },
    /// A Receiver Report block (RFC 3550 §6.4.2) — loss/jitter stats for one
    /// source, used for QoS metrics and bandwidth estimation.
    ReceiverReport {
        /// SSRC of the media source this report block describes.
        ssrc: u32,
        /// Fraction of packets lost since the previous report (8.8 fixed-point).
        fraction_lost: u8,
        /// Cumulative number of packets lost over the session (24-bit).
        cumulative_lost: u32,
        /// Interarrival jitter estimate, in RTP timestamp units.
        jitter: u32,
    },
}

/// Parse a (decrypted) RTCP compound packet into its feedback messages.
///
/// Tolerant by construction: a truncated or unknown sub-packet ends the walk
/// instead of panicking, and unrecognized payload types are skipped via their
/// declared length. Returns every message understood, in wire order.
pub fn parse_compound(mut buf: &[u8]) -> Vec<RtcpFeedback> {
    let mut out = Vec::new();
    while buf.len() >= 4 {
        // Common header: V(2)P(1)RC/FMT(5) | PT(8) | length(16, words-1).
        let version = buf[0] >> 6;
        if version != 2 {
            break;
        }
        let fmt = buf[0] & 0x1f;
        let pt = buf[1];
        let len_words = u16::from_be_bytes([buf[2], buf[3]]) as usize;
        let pkt_len = (len_words + 1) * 4;
        if pkt_len == 0 || pkt_len > buf.len() {
            break;
        }
        let pkt = &buf[..pkt_len];
        match (pt, fmt) {
            (PT_PSFB, 1) if pkt_len >= 12 => out.push(RtcpFeedback::Pli {
                sender_ssrc: be32(&pkt[4..]),
                media_ssrc: be32(&pkt[8..]),
            }),
            (PT_PSFB, 4) if pkt_len >= 16 => out.push(RtcpFeedback::Fir {
                sender_ssrc: be32(&pkt[4..]),
                media_ssrc: be32(&pkt[12..]),
            }),
            // REMB rides on PSFB FMT 15 with a "REMB" application identifier in
            // the FCI; the other FMT-15 user (TWCC) is RTPFB and is handled by the
            // transport's own BWE, not parsed here.
            (PT_PSFB, 15) if pkt_len >= 20 && &pkt[12..16] == b"REMB" => out.push(parse_remb(pkt)),
            (PT_RTPFB, 1) if pkt_len >= 12 => out.push(parse_nack(pkt)),
            (PT_RR, _) if pkt_len >= 8 => parse_report_blocks(&pkt[8..], fmt, &mut out),
            // SR's header carries a 20-byte sender-info block before its report
            // blocks; skip past it to the per-source RR blocks.
            (PT_SR, _) if pkt_len >= 28 => parse_report_blocks(&pkt[28..], fmt, &mut out),
            _ => {}
        }
        buf = &buf[pkt_len..];
    }
    out
}

fn parse_nack(pkt: &[u8]) -> RtcpFeedback {
    let sender_ssrc = be32(&pkt[4..]);
    let media_ssrc = be32(&pkt[8..]);
    let mut lost = Vec::new();
    // FCI entries start at byte 12: each is PID(16) + bitmask(16). The bitmask's
    // bit i (LSB first) flags loss of PID+i+1.
    let mut off = 12;
    while off + 4 <= pkt.len() {
        let pid = u16::from_be_bytes([pkt[off], pkt[off + 1]]);
        let blp = u16::from_be_bytes([pkt[off + 2], pkt[off + 3]]);
        lost.push(pid);
        for i in 0..16 {
            if blp & (1 << i) != 0 {
                lost.push(pid.wrapping_add(i + 1));
            }
        }
        off += 4;
    }
    RtcpFeedback::Nack {
        sender_ssrc,
        media_ssrc,
        lost,
    }
}

/// Decode a REMB FCI: `Num SSRC` then a 6-bit exponent + 18-bit mantissa, with
/// `bitrate = mantissa << exp` (bits/sec). The trailing feedback SSRCs are not
/// needed for a rate estimate and are ignored.
fn parse_remb(pkt: &[u8]) -> RtcpFeedback {
    let sender_ssrc = be32(&pkt[4..]);
    // FCI: [12..16]="REMB", [16]=num ssrc, [17]=EEEEEEMM, [18],[19]=mantissa.
    let exp = (pkt[17] >> 2) as u32;
    let mantissa = (((pkt[17] & 0x03) as u64) << 16) | ((pkt[18] as u64) << 8) | pkt[19] as u64;
    // Saturate rather than overflow on an absurd exponent from a malformed peer.
    let bitrate_bps = mantissa.checked_shl(exp).unwrap_or(u64::MAX);
    RtcpFeedback::Remb {
        sender_ssrc,
        bitrate_bps,
    }
}

/// Decode the `count` (RC) 24-byte report blocks following an SR/RR header.
fn parse_report_blocks(blocks: &[u8], count: u8, out: &mut Vec<RtcpFeedback>) {
    for i in 0..count as usize {
        let off = i * 24;
        if off + 24 > blocks.len() {
            break;
        }
        let b = &blocks[off..off + 24];
        out.push(RtcpFeedback::ReceiverReport {
            ssrc: be32(b),
            fraction_lost: b[4],
            cumulative_lost: u32::from_be_bytes([0, b[5], b[6], b[7]]),
            // Layout: [8..12] ext highest seq, [12..16] interarrival jitter.
            jitter: be32(&b[12..]),
        });
    }
}

#[inline]
fn be32(b: &[u8]) -> u32 {
    u32::from_be_bytes([b[0], b[1], b[2], b[3]])
}

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

    #[test]
    fn pli_has_correct_header_and_ssrcs() {
        let p = build_pli(0x1111_1111, 0x2222_2222);
        assert_eq!(p.len(), 12);
        assert_eq!(p[0], 0x81); // V=2, FMT=1
        assert_eq!(p[1], 206); // PSFB
        assert_eq!(u16::from_be_bytes([p[2], p[3]]), 2); // length words-1
        assert_eq!(&p[4..8], &0x1111_1111u32.to_be_bytes());
        assert_eq!(&p[8..12], &0x2222_2222u32.to_be_bytes());
    }

    #[test]
    fn fir_carries_seq_and_target_ssrc() {
        let p = build_fir(1, 0xDEAD_BEEF, 7);
        assert_eq!(p.len(), 20);
        assert_eq!(p[0], 0x84); // FMT=4
        assert_eq!(&p[12..16], &0xDEAD_BEEFu32.to_be_bytes());
        assert_eq!(p[16], 7); // seq_nr
    }

    #[test]
    fn parses_our_own_pli_round_trip() {
        let p = build_pli(0xAAAA_AAAA, 0xBBBB_BBBB);
        assert_eq!(
            parse_compound(&p),
            vec![RtcpFeedback::Pli {
                sender_ssrc: 0xAAAA_AAAA,
                media_ssrc: 0xBBBB_BBBB,
            }]
        );
    }

    #[test]
    fn parses_fir_round_trip() {
        let p = build_fir(1, 0xDEAD_BEEF, 3);
        assert_eq!(
            parse_compound(&p),
            vec![RtcpFeedback::Fir {
                sender_ssrc: 1,
                media_ssrc: 0xDEAD_BEEF,
            }]
        );
    }

    #[test]
    fn parses_generic_nack_with_bitmask() {
        // RTPFB FMT 1, sender 1, media 2, one FCI: PID=100, BLP bit0 + bit2 set
        // → lost 100, 101, 103.
        let mut p = Vec::new();
        p.push(0x80 | 1);
        p.push(PT_RTPFB);
        p.extend_from_slice(&3u16.to_be_bytes()); // 4 words → len 3
        p.extend_from_slice(&1u32.to_be_bytes());
        p.extend_from_slice(&2u32.to_be_bytes());
        p.extend_from_slice(&100u16.to_be_bytes());
        p.extend_from_slice(&0b0000_0000_0000_0101u16.to_be_bytes());
        assert_eq!(
            parse_compound(&p),
            vec![RtcpFeedback::Nack {
                sender_ssrc: 1,
                media_ssrc: 2,
                lost: vec![100, 101, 103],
            }]
        );
    }

    #[test]
    fn parses_receiver_report_block() {
        // RR, RC=1, reporter SSRC, one 24-byte block.
        let mut p = Vec::new();
        p.push(0x80 | 1); // RC=1
        p.push(PT_RR);
        p.extend_from_slice(&7u16.to_be_bytes()); // header(2w) + block(6w)=8w → 7
        p.extend_from_slice(&0xCAFEu32.to_be_bytes()); // reporter ssrc
        p.extend_from_slice(&0x1234_5678u32.to_be_bytes()); // source ssrc
        p.push(64); // fraction lost
        p.extend_from_slice(&[0, 0, 10]); // cumulative lost = 10
        p.extend_from_slice(&55u32.to_be_bytes()); // ext highest seq (ignored)
        p.extend_from_slice(&99u32.to_be_bytes()); // jitter
        p.extend_from_slice(&0u32.to_be_bytes()); // LSR
        p.extend_from_slice(&0u32.to_be_bytes()); // DLSR
        assert_eq!(
            parse_compound(&p),
            vec![RtcpFeedback::ReceiverReport {
                ssrc: 0x1234_5678,
                fraction_lost: 64,
                cumulative_lost: 10,
                jitter: 99,
            }]
        );
    }

    #[test]
    fn sr_has_sender_info_layout() {
        let p = build_sr(0xABCD_1234, 0x1122_3344_5566_7788, 90_000, 5, 1000);
        assert_eq!(p.len(), 28);
        assert_eq!(p[0], 0x80); // V=2, RC=0
        assert_eq!(p[1], 200); // SR
        assert_eq!(u16::from_be_bytes([p[2], p[3]]), 6);
        assert_eq!(&p[4..8], &0xABCD_1234u32.to_be_bytes());
        assert_eq!(&p[8..16], &0x1122_3344_5566_7788u64.to_be_bytes());
        assert_eq!(&p[16..20], &90_000u32.to_be_bytes());
        assert_eq!(&p[20..24], &5u32.to_be_bytes());
        assert_eq!(&p[24..28], &1000u32.to_be_bytes());
    }

    #[test]
    fn ntp_now_is_after_2020() {
        // High 32 bits are seconds since 1900; 2020-01-01 is 3_786_825_600.
        assert!((ntp_now() >> 32) > 3_786_825_600);
    }

    #[test]
    fn parses_remb_bitrate() {
        // PSFB FMT 15, sender 9, media 0, FCI "REMB" num=1, exp=3 mantissa=1000
        // → bitrate = 1000 << 3 = 8000 bps.
        let mut p = Vec::new();
        p.push(0x80 | 15);
        p.push(PT_PSFB);
        p.extend_from_slice(&5u16.to_be_bytes()); // 6 words → len 5
        p.extend_from_slice(&9u32.to_be_bytes());
        p.extend_from_slice(&0u32.to_be_bytes());
        p.extend_from_slice(b"REMB");
        p.push(1); // num ssrc
        let mantissa: u32 = 1000;
        let exp: u32 = 3;
        p.push(((exp << 2) as u8) | ((mantissa >> 16) as u8 & 0x03));
        p.push((mantissa >> 8) as u8);
        p.push(mantissa as u8);
        p.extend_from_slice(&0x1234u32.to_be_bytes()); // one feedback ssrc
        assert_eq!(
            parse_compound(&p),
            vec![RtcpFeedback::Remb {
                sender_ssrc: 9,
                bitrate_bps: 8000,
            }]
        );
    }

    #[test]
    fn walks_compound_and_skips_unknown() {
        // Unknown PT (e.g. SDES=202) sandwiched between two PLIs; the walker steps
        // over the unknown via its length and still yields both PLIs.
        let mut p = build_pli(1, 2);
        p.push(0x80);
        p.push(202); // SDES
        p.extend_from_slice(&1u16.to_be_bytes()); // 2 words
        p.extend_from_slice(&0u32.to_be_bytes());
        p.extend_from_slice(&build_pli(3, 4));
        assert_eq!(
            parse_compound(&p),
            vec![
                RtcpFeedback::Pli {
                    sender_ssrc: 1,
                    media_ssrc: 2
                },
                RtcpFeedback::Pli {
                    sender_ssrc: 3,
                    media_ssrc: 4
                },
            ]
        );
    }

    #[test]
    fn tolerates_truncated_tail() {
        let mut p = build_pli(1, 2);
        p.extend_from_slice(&[0x80, 206]); // dangling 2 bytes
        assert_eq!(parse_compound(&p).len(), 1);
    }
}