Skip to main content

rtc_rtp/codec/h264/
mod.rs

1//! H.264 RTP payload format ([RFC 6184]).
2//!
3//! A NAL unit that fits the MTU is sent as-is. Larger ones are split into FU-A fragments, and
4//! several small ones (typically SPS and PPS) may be combined into one STAP-A aggregate. The
5//! `*_NALU_TYPE` constants name the types this payloader recognises, and the `*_BITMASK`
6//! constants describe how the NAL header and FU header pack their fields.
7//!
8//! [RFC 6184]: https://datatracker.ietf.org/doc/html/rfc6184
9#[cfg(test)]
10mod h264_test;
11
12use crate::packetizer::{Depacketizer, Payloader};
13use shared::error::{Error, Result};
14
15use bytes::{BufMut, Bytes, BytesMut};
16
17/// H264Payloader payloads H264 packets
18#[derive(Default, Debug, Clone)]
19pub struct H264Payloader {
20    sps_nalu: Option<Bytes>,
21    pps_nalu: Option<Bytes>,
22}
23
24/// NAL type 24: STAP-A, which aggregates several small NAL units into one payload.
25pub const STAPA_NALU_TYPE: u8 = 24;
26/// NAL type 28: FU-A, which fragments one NAL unit across several payloads.
27pub const FUA_NALU_TYPE: u8 = 28;
28/// NAL type 29: FU-B, FU-A with a decoding-order number. Not used by WebRTC.
29pub const FUB_NALU_TYPE: u8 = 29;
30/// NAL type 7: sequence parameter set.
31pub const SPS_NALU_TYPE: u8 = 7;
32/// NAL type 8: picture parameter set.
33pub const PPS_NALU_TYPE: u8 = 8;
34/// NAL type 9: access unit delimiter.
35pub const AUD_NALU_TYPE: u8 = 9;
36/// NAL type 12: filler data, which is dropped rather than sent.
37pub const FILLER_NALU_TYPE: u8 = 12;
38
39/// Bytes of FU-A header prefixed to each fragment.
40pub const FUA_HEADER_SIZE: usize = 2;
41/// Bytes of STAP-A header prefixed to an aggregate.
42pub const STAPA_HEADER_SIZE: usize = 1;
43/// Bytes of length prefix before each NAL unit inside a STAP-A.
44pub const STAPA_NALU_LENGTH_SIZE: usize = 2;
45
46/// Mask selecting the NAL type from a NAL header byte.
47pub const NALU_TYPE_BITMASK: u8 = 0x1F;
48/// Mask selecting `nal_ref_idc` from a NAL header byte.
49pub const NALU_REF_IDC_BITMASK: u8 = 0x60;
50/// FU header start bit: this fragment begins a NAL unit.
51pub const FU_START_BITMASK: u8 = 0x80;
52/// FU header end bit: this fragment completes a NAL unit.
53pub const FU_END_BITMASK: u8 = 0x40;
54
55/// The STAP-A header byte this payloader emits.
56pub const OUTPUT_STAP_AHEADER: u8 = 0x78;
57
58/// The Annex B start code (`00 00 00 01`) that delimits NAL units in a byte stream.
59pub static ANNEXB_NALUSTART_CODE: Bytes = Bytes::from_static(&[0x00, 0x00, 0x00, 0x01]);
60
61impl H264Payloader {
62    fn next_ind(nalu: &Bytes, start: usize) -> (isize, isize) {
63        let mut zero_count = 0;
64
65        for (i, &b) in nalu[start..].iter().enumerate() {
66            if b == 0 {
67                zero_count += 1;
68                continue;
69            } else if b == 1 && zero_count >= 2 {
70                return ((start + i - zero_count) as isize, zero_count as isize + 1);
71            }
72            zero_count = 0
73        }
74        (-1, -1)
75    }
76
77    fn emit(&mut self, nalu: &Bytes, mtu: usize, payloads: &mut Vec<Bytes>) {
78        if nalu.is_empty() {
79            return;
80        }
81
82        let nalu_type = nalu[0] & NALU_TYPE_BITMASK;
83        let nalu_ref_idc = nalu[0] & NALU_REF_IDC_BITMASK;
84
85        if nalu_type == AUD_NALU_TYPE || nalu_type == FILLER_NALU_TYPE {
86            return;
87        } else if nalu_type == SPS_NALU_TYPE {
88            self.sps_nalu = Some(nalu.clone());
89            return;
90        } else if nalu_type == PPS_NALU_TYPE {
91            self.pps_nalu = Some(nalu.clone());
92            return;
93        } else if let (Some(sps_nalu), Some(pps_nalu)) = (&self.sps_nalu, &self.pps_nalu) {
94            // Pack current NALU with SPS and PPS as STAP-A
95            let sps_len = (sps_nalu.len() as u16).to_be_bytes();
96            let pps_len = (pps_nalu.len() as u16).to_be_bytes();
97
98            let mut stap_a_nalu = Vec::with_capacity(1 + 2 + sps_nalu.len() + 2 + pps_nalu.len());
99            stap_a_nalu.push(OUTPUT_STAP_AHEADER);
100            stap_a_nalu.extend(sps_len);
101            stap_a_nalu.extend_from_slice(sps_nalu);
102            stap_a_nalu.extend(pps_len);
103            stap_a_nalu.extend_from_slice(pps_nalu);
104            if stap_a_nalu.len() <= mtu {
105                payloads.push(Bytes::from(stap_a_nalu));
106            }
107        }
108
109        if self.sps_nalu.is_some() && self.pps_nalu.is_some() {
110            self.sps_nalu = None;
111            self.pps_nalu = None;
112        }
113
114        // Single NALU
115        if nalu.len() <= mtu {
116            payloads.push(nalu.clone());
117            return;
118        }
119
120        // FU-A
121        let max_fragment_size = mtu as isize - FUA_HEADER_SIZE as isize;
122
123        // The FU payload consists of fragments of the payload of the fragmented
124        // NAL unit so that if the fragmentation unit payloads of consecutive
125        // FUs are sequentially concatenated, the payload of the fragmented NAL
126        // unit can be reconstructed.  The NAL unit type octet of the fragmented
127        // NAL unit is not included as such in the fragmentation unit payload,
128        // 	but rather the information of the NAL unit type octet of the
129        // fragmented NAL unit is conveyed in the F and NRI fields of the FU
130        // indicator octet of the fragmentation unit and in the type field of
131        // the FU header.  An FU payload MAY have any number of octets and MAY
132        // be empty.
133
134        let nalu_data = nalu;
135        // According to the RFC, the first octet is skipped due to redundant information
136        let mut nalu_data_index = 1;
137        let nalu_data_length = nalu.len() as isize - nalu_data_index;
138        let mut nalu_data_remaining = nalu_data_length;
139
140        if std::cmp::min(max_fragment_size, nalu_data_remaining) <= 0 {
141            return;
142        }
143
144        while nalu_data_remaining > 0 {
145            let current_fragment_size = std::cmp::min(max_fragment_size, nalu_data_remaining);
146            //out: = make([]byte, fuaHeaderSize + currentFragmentSize)
147            let mut out = BytesMut::with_capacity(FUA_HEADER_SIZE + current_fragment_size as usize);
148            // +---------------+
149            // |0|1|2|3|4|5|6|7|
150            // +-+-+-+-+-+-+-+-+
151            // |F|NRI|  Type   |
152            // +---------------+
153            let b0 = FUA_NALU_TYPE | nalu_ref_idc;
154            out.put_u8(b0);
155
156            // +---------------+
157            //|0|1|2|3|4|5|6|7|
158            //+-+-+-+-+-+-+-+-+
159            //|S|E|R|  Type   |
160            //+---------------+
161
162            let mut b1 = nalu_type;
163            if nalu_data_remaining == nalu_data_length {
164                // Set start bit
165                b1 |= 1 << 7;
166            } else if nalu_data_remaining - current_fragment_size == 0 {
167                // Set end bit
168                b1 |= 1 << 6;
169            }
170            out.put_u8(b1);
171
172            out.put(
173                &nalu_data
174                    [nalu_data_index as usize..(nalu_data_index + current_fragment_size) as usize],
175            );
176            payloads.push(out.freeze());
177
178            nalu_data_remaining -= current_fragment_size;
179            nalu_data_index += current_fragment_size;
180        }
181    }
182}
183
184impl Payloader for H264Payloader {
185    /// Payload fragments a H264 packet across one or more byte arrays
186    fn payload(&mut self, mtu: usize, payload: &Bytes) -> Result<Vec<Bytes>> {
187        if payload.is_empty() || mtu == 0 {
188            return Ok(vec![]);
189        }
190
191        let mut payloads = vec![];
192
193        let (mut next_ind_start, mut next_ind_len) = H264Payloader::next_ind(payload, 0);
194        if next_ind_start == -1 {
195            self.emit(payload, mtu, &mut payloads);
196        } else {
197            while next_ind_start != -1 {
198                let prev_start = (next_ind_start + next_ind_len) as usize;
199                let (next_ind_start2, next_ind_len2) = H264Payloader::next_ind(payload, prev_start);
200                next_ind_start = next_ind_start2;
201                next_ind_len = next_ind_len2;
202                if next_ind_start != -1 {
203                    self.emit(
204                        &payload.slice(prev_start..next_ind_start as usize),
205                        mtu,
206                        &mut payloads,
207                    );
208                } else {
209                    // Emit until end of stream, no end indicator found
210                    self.emit(&payload.slice(prev_start..), mtu, &mut payloads);
211                }
212            }
213        }
214
215        Ok(payloads)
216    }
217
218    fn clone_to(&self) -> Box<dyn Payloader> {
219        Box::new(self.clone())
220    }
221}
222
223/// H264Packet represents the H264 header that is stored in the payload of an RTP Packet
224#[derive(PartialEq, Eq, Debug, Default, Clone)]
225pub struct H264Packet {
226    /// Whether to emit AVCC length-prefixed output instead of Annex B start codes.
227    pub is_avc: bool,
228    fua_buffer: Option<BytesMut>,
229}
230
231impl Depacketizer for H264Packet {
232    /// depacketize parses the passed byte slice and stores the result in the H264Packet this method is called upon
233    fn depacketize(&mut self, packet: &Bytes) -> Result<Bytes> {
234        if packet.len() <= 1 {
235            return Err(Error::ErrShortPacket);
236        }
237
238        // NALU Types
239        // https://tools.ietf.org/html/rfc6184#section-5.4
240        let b0 = packet[0];
241        let nalu_type = b0 & NALU_TYPE_BITMASK;
242
243        // The AUD NALU can be size 2 (1 byte header, 1 byte payload)
244        if packet.len() <= 2 && nalu_type != AUD_NALU_TYPE {
245            return Err(Error::ErrShortPacket);
246        }
247
248        let mut payload = BytesMut::new();
249
250        match nalu_type {
251            1..=23 => {
252                if self.is_avc {
253                    payload.put_u32(packet.len() as u32);
254                } else {
255                    payload.put(&*ANNEXB_NALUSTART_CODE);
256                }
257                payload.put(&*packet.clone());
258                Ok(payload.freeze())
259            }
260            STAPA_NALU_TYPE => {
261                let mut curr_offset = STAPA_HEADER_SIZE;
262                while curr_offset < packet.len() {
263                    let nalu_size =
264                        ((packet[curr_offset] as usize) << 8) | packet[curr_offset + 1] as usize;
265                    curr_offset += STAPA_NALU_LENGTH_SIZE;
266
267                    if packet.len() < curr_offset + nalu_size {
268                        return Err(Error::StapASizeLargerThanBuffer(
269                            nalu_size,
270                            packet.len() - curr_offset,
271                        ));
272                    }
273
274                    if self.is_avc {
275                        payload.put_u32(nalu_size as u32);
276                    } else {
277                        payload.put(&*ANNEXB_NALUSTART_CODE);
278                    }
279                    payload.put(&*packet.slice(curr_offset..curr_offset + nalu_size));
280                    curr_offset += nalu_size;
281                }
282
283                Ok(payload.freeze())
284            }
285            FUA_NALU_TYPE => {
286                if packet.len() < FUA_HEADER_SIZE {
287                    return Err(Error::ErrShortPacket);
288                }
289
290                if self.fua_buffer.is_none() {
291                    self.fua_buffer = Some(BytesMut::new());
292                }
293
294                if let Some(fua_buffer) = &mut self.fua_buffer {
295                    fua_buffer.put(&*packet.slice(FUA_HEADER_SIZE..));
296                }
297
298                let b1 = packet[1];
299                if b1 & FU_END_BITMASK != 0 {
300                    let nalu_ref_idc = b0 & NALU_REF_IDC_BITMASK;
301                    let fragmented_nalu_type = b1 & NALU_TYPE_BITMASK;
302
303                    if let Some(fua_buffer) = self.fua_buffer.take() {
304                        if self.is_avc {
305                            payload.put_u32((fua_buffer.len() + 1) as u32);
306                        } else {
307                            payload.put(&*ANNEXB_NALUSTART_CODE);
308                        }
309                        payload.put_u8(nalu_ref_idc | fragmented_nalu_type);
310                        payload.put(fua_buffer);
311                    }
312
313                    Ok(payload.freeze())
314                } else {
315                    Ok(Bytes::new())
316                }
317            }
318            _ => Err(Error::NaluTypeIsNotHandled(nalu_type)),
319        }
320    }
321
322    /// is_partition_head checks if this is the head of a packetized nalu stream.
323    fn is_partition_head(&self, payload: &Bytes) -> bool {
324        if payload.len() < 2 {
325            return false;
326        }
327
328        if payload[0] & NALU_TYPE_BITMASK == FUA_NALU_TYPE
329            || payload[0] & NALU_TYPE_BITMASK == FUB_NALU_TYPE
330        {
331            (payload[1] & FU_START_BITMASK) != 0
332        } else {
333            true
334        }
335    }
336
337    fn is_partition_tail(&self, marker: bool, _payload: &Bytes) -> bool {
338        marker
339    }
340}