rtp/codecs/h264/
mod.rs

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

rtp/codecs/h264/
mod.rs
