rtp_parse/rtcp/
rtcp_fb_nack.rs

1use std::collections::BTreeSet;
2
3use crate::{
4    rtcp::{rtcp_fb_header::write_rtcp_fb_header, rtcp_header::write_rtcp_header},
5    PacketBuffer, PacketBufferMut,
6};
7use anyhow::{anyhow, bail, Context, Result};
8use bit_cursor::{
9    bit_read_exts::BitReadExts, bit_write_exts::BitWriteExts, byte_order::NetworkOrder,
10    nsw_types::u5,
11};
12
13use super::{rtcp_fb_header::RtcpFbHeader, rtcp_header::RtcpHeader};
14
15/// https://datatracker.ietf.org/doc/html/rfc4585#section-6.2.1
16///  0                   1                   2                   3
17///  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
18/// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
19/// |            PID                |             BLP               |
20/// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
21#[derive(Debug)]
22pub struct RtcpFbNackPacket {
23    pub header: RtcpHeader,
24    pub fb_header: RtcpFbHeader,
25    pub missing_seq_nums: BTreeSet<u16>,
26}
27
28// TODO: somewhere in here we need to enforce some kind of maximum number of sequence numbers that
29// can be included in a nack packet (or, a way for a nack to fail to serialize due to having more
30// sequence numbers that couldn't be added to the buffer)
31
32impl RtcpFbNackPacket {
33    pub const FMT: u5 = u5::new(1);
34}
35
36pub fn read_rtcp_fb_nack<B: PacketBuffer>(
37    buf: &mut B,
38    header: RtcpHeader,
39    fb_header: RtcpFbHeader,
40) -> Result<RtcpFbNackPacket> {
41    let mut missing_seq_nums = BTreeSet::new();
42    let mut nack_block_num = 1;
43    while buf.bytes_remaining() >= NackBlock::SIZE_BYTES {
44        let mut nack_block =
45            read_nack_block(buf).with_context(|| format!("nack block {nack_block_num}"))?;
46        missing_seq_nums.append(&mut nack_block.missing_seq_nums);
47        nack_block_num += 1;
48    }
49
50    Ok(RtcpFbNackPacket {
51        header,
52        fb_header,
53        missing_seq_nums,
54    })
55}
56
57pub fn write_rtcp_fb_nack<B: PacketBufferMut>(
58    buf: &mut B,
59    fb_nack: &RtcpFbNackPacket,
60) -> Result<()> {
61    write_rtcp_header(buf, &fb_nack.header).context("rtcp header")?;
62    write_rtcp_fb_header(buf, &fb_nack.fb_header).context("fb header")?;
63
64    for (i, missing_packet_chunk) in fb_nack
65        .missing_seq_nums
66        .chunk_by_max_difference(16)
67        .into_iter()
68        .enumerate()
69    {
70        let nack_block = NackBlock {
71            missing_seq_nums: missing_packet_chunk,
72        };
73        if buf.bytes_remaining() < NackBlock::SIZE_BYTES {
74            bail!("Not enough room in buffer for nack block {i}");
75        }
76        write_nack_block(buf, &nack_block).with_context(|| format!("nack block {i}"))?;
77    }
78
79    Ok(())
80}
81
82pub struct NackBlock {
83    // We don't use the BTreeSet here because when this type is used we're either parsing from a
84    // packet, where know the order will be right (since we parse it that way), or about to
85    // serialize, where we also know the order is already correct.
86    missing_seq_nums: BTreeSet<u16>,
87}
88
89impl NackBlock {
90    pub const SIZE_BYTES: usize = 4;
91}
92
93pub fn read_nack_block<B: PacketBuffer>(buf: &mut B) -> Result<NackBlock> {
94    let packet_id = buf.read_u16::<NetworkOrder>().context("packet id")?;
95    let blp = buf.read_u16::<NetworkOrder>().context("blp")?;
96
97    let mut missing_seq_nums = BTreeSet::new();
98    missing_seq_nums.insert(packet_id);
99    for shift_amount in 0..16 {
100        if (blp >> shift_amount) & 0x1 == 1 {
101            missing_seq_nums.insert(packet_id + shift_amount + 1);
102        }
103    }
104
105    Ok(NackBlock { missing_seq_nums })
106}
107
108pub fn write_nack_block<B: PacketBufferMut>(buf: &mut B, nack_block: &NackBlock) -> Result<()> {
109    let packet_id = nack_block.missing_seq_nums.first().ok_or(anyhow!(
110        "NackBlock must contain at least one sequence number"
111    ))?;
112    buf.write_u16::<NetworkOrder>(*packet_id)
113        .context("packet id")?;
114    let mut blp = 0u16;
115    // Skip past the first one, since that was used for the packet id
116    for missing_seq_num in nack_block.missing_seq_nums.iter().skip(1) {
117        let delta = missing_seq_num - packet_id;
118        if delta > 16 {
119            bail!("NACK missing sequence numbers can not span more than 16 sequence numbers");
120        }
121        let mask = 1u16 << (delta - 1);
122        blp |= mask;
123    }
124    buf.write_u16::<NetworkOrder>(blp).context("blp")?;
125
126    Ok(())
127}
128
129trait ChunkByMaxDifference<T> {
130    fn chunk_by_max_difference(&self, max_diff: T) -> Vec<BTreeSet<T>>;
131}
132
133impl ChunkByMaxDifference<u16> for BTreeSet<u16> {
134    fn chunk_by_max_difference(&self, max_diff: u16) -> Vec<BTreeSet<u16>> {
135        let mut all_chunks: Vec<BTreeSet<u16>> = Vec::new();
136        let Some(first) = self.first() else {
137            return all_chunks;
138        };
139        let mut curr_chunk: BTreeSet<u16> = BTreeSet::from([*first]);
140        for value in self.iter().skip(1) {
141            if value - curr_chunk.first().unwrap() > max_diff {
142                all_chunks.push(curr_chunk);
143                curr_chunk = BTreeSet::from([*value]);
144            } else {
145                curr_chunk.insert(*value);
146            }
147        }
148
149        all_chunks
150    }
151}
152
153#[cfg(test)]
154mod test {
155    use std::collections::BTreeSet;
156
157    use bit_cursor::bit_cursor::BitCursor;
158    use bitvec::{order::Msb0, vec::BitVec};
159
160    use super::{read_nack_block, write_nack_block};
161
162    #[test]
163    fn test_read_nack_block() {
164        // Missing seq nums 10, 11, 16, 18, 22, 24, 26
165        let data_buf: [u8; 4] = [0x00, 0x0A, 0xA8, 0xA1];
166        let mut cursor = BitCursor::new(BitVec::<u8, Msb0>::from_slice(&data_buf));
167        let nack_block = read_nack_block(&mut cursor).unwrap();
168        assert_eq!(
169            nack_block.missing_seq_nums,
170            BTreeSet::from([10, 11, 16, 18, 22, 24, 26]),
171        );
172    }
173
174    #[test]
175    fn test_write_nack_block() {
176        // Missing seq nums 10, 11, 16, 18, 22, 24, 26
177        let data_buf: [u8; 4] = [0x00, 0x0A, 0xA8, 0xA1];
178        let mut cursor = BitCursor::new(BitVec::<u8, Msb0>::from_slice(&data_buf));
179        let nack_block = read_nack_block(&mut cursor).unwrap();
180
181        let write_data_buf = [0u8; 4];
182        let mut write_cursor = BitCursor::new(BitVec::<u8, Msb0>::from_slice(&write_data_buf));
183        write_nack_block(&mut write_cursor, &nack_block).unwrap();
184        let write_data = write_cursor.into_inner().into_vec();
185        assert_eq!(&data_buf, &write_data[..]);
186    }
187}