webrtc_media/io/h264_writer/
mod.rs

1#[cfg(test)]
2mod h264_writer_test;
3
4use std::io::{Seek, Write};
5
6use rtp::codecs::h264::H264Packet;
7use rtp::packetizer::Depacketizer;
8
9use crate::error::Result;
10use crate::io::Writer;
11
12const NALU_TTYPE_STAP_A: u32 = 24;
13const NALU_TTYPE_SPS: u32 = 7;
14const NALU_TYPE_BITMASK: u32 = 0x1F;
15
16fn is_key_frame(data: &[u8]) -> bool {
17    if data.len() < 4 {
18        false
19    } else {
20        let word = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
21        let nalu_type = (word >> 24) & NALU_TYPE_BITMASK;
22        (nalu_type == NALU_TTYPE_STAP_A && (word & NALU_TYPE_BITMASK) == NALU_TTYPE_SPS)
23            || (nalu_type == NALU_TTYPE_SPS)
24    }
25}
26
27/// H264Writer is used to take RTP packets, parse them and
28/// write the data to an io.Writer.
29/// Currently it only supports non-interleaved mode
30/// Therefore, only 1-23, 24 (STAP-A), 28 (FU-A) NAL types are allowed.
31/// <https://tools.ietf.org/html/rfc6184#section-5.2>
32pub struct H264Writer<W: Write + Seek> {
33    writer: W,
34    has_key_frame: bool,
35    cached_packet: Option<H264Packet>,
36}
37
38impl<W: Write + Seek> H264Writer<W> {
39    // new initializes a new H264 writer with an io.Writer output
40    pub fn new(writer: W) -> Self {
41        H264Writer {
42            writer,
43            has_key_frame: false,
44            cached_packet: None,
45        }
46    }
47}
48
49impl<W: Write + Seek> Writer for H264Writer<W> {
50    /// write_rtp adds a new packet and writes the appropriate headers for it
51    fn write_rtp(&mut self, packet: &rtp::packet::Packet) -> Result<()> {
52        if packet.payload.is_empty() {
53            return Ok(());
54        }
55
56        if !self.has_key_frame {
57            self.has_key_frame = is_key_frame(&packet.payload);
58            if !self.has_key_frame {
59                // key frame not defined yet. discarding packet
60                return Ok(());
61            }
62        }
63
64        if self.cached_packet.is_none() {
65            self.cached_packet = Some(H264Packet::default());
66        }
67
68        if let Some(cached_packet) = &mut self.cached_packet {
69            let payload = cached_packet.depacketize(&packet.payload)?;
70
71            self.writer.write_all(&payload)?;
72        }
73
74        Ok(())
75    }
76
77    /// close closes the underlying writer
78    fn close(&mut self) -> Result<()> {
79        self.cached_packet = None;
80        self.writer.flush()?;
81        Ok(())
82    }
83}
webrtc_media/io/h264_writer/mod.rs

webrtc_media/io/h264_writer/
mod.rs
