mpeg2ts 0.6.0

MPEG2-TS decoding/encoding library
Documentation
use crate::Result;
use crate::ts::payload::{Bytes, Null, Pat, Pes, Pmt};
use crate::ts::{AdaptationField, Pid, TsHeader, TsPacket, TsPayload};
use std::collections::HashMap;
use std::io::Read;

/// The `ReadTsPacket` trait allows for reading TS packets from a source.
pub trait ReadTsPacket {
    /// Reads a TS packet.
    ///
    /// If the end of the stream is reached, it will return `Ok(None)`.
    fn read_ts_packet(&mut self) -> Result<Option<TsPacket>>;
}

/// TS packet reader.
#[derive(Debug)]
pub struct TsPacketReader<R> {
    stream: R,
    pids: HashMap<Pid, PidKind>,
}
impl<R: Read> TsPacketReader<R> {
    /// Makes a new `TsPacketReader` instance.
    pub fn new(stream: R) -> Self {
        TsPacketReader {
            stream,
            pids: HashMap::new(),
        }
    }

    /// Returns a reference to the underlaying byte stream.
    pub fn stream(&self) -> &R {
        &self.stream
    }

    /// Converts `TsPacketReader` into the underlaying byte stream `R`.
    pub fn into_stream(self) -> R {
        self.stream
    }
}
impl<R: Read> ReadTsPacket for TsPacketReader<R> {
    fn read_ts_packet(&mut self) -> Result<Option<TsPacket>> {
        let mut reader = self.stream.by_ref().take(TsPacket::SIZE as u64);
        let mut peek = [0; 1];
        let eos = reader.read(&mut peek)? == 0;
        if eos {
            return Ok(None);
        }

        let (header, adaptation_field_control, payload_unit_start_indicator) =
            TsHeader::read_from(peek.chain(&mut reader))?;

        let adaptation_field = if adaptation_field_control.has_adaptation_field() {
            AdaptationField::read_from(&mut reader)?
        } else {
            None
        };

        let payload = if adaptation_field_control.has_payload() {
            let payload = match header.pid.as_u16() {
                Pid::PAT => {
                    let pat = Pat::read_from(&mut reader)?;
                    for pa in &pat.table {
                        self.pids.insert(pa.program_map_pid, PidKind::Pmt);
                    }
                    TsPayload::Pat(pat)
                }
                Pid::NULL => {
                    let null = Null::read_from(&mut reader)?;
                    TsPayload::Null(null)
                }
                0x01..=0x1F | 0x1FFB => {
                    // Unknown (unsupported) packets
                    let bytes = Bytes::read_from(&mut reader)?;
                    TsPayload::Raw(bytes)
                }
                _ => {
                    let kind = self.pids.get(&header.pid).cloned().ok_or_else(|| {
                        crate::Error::invalid_input(format!("Unknown PID: header={:?}", header))
                    })?;
                    match kind {
                        PidKind::Pmt => {
                            let pmt = Pmt::read_from(&mut reader)?;
                            for es in &pmt.es_info {
                                self.pids.insert(es.elementary_pid, PidKind::Pes);
                            }
                            TsPayload::Pmt(pmt)
                        }
                        PidKind::Pes => {
                            if payload_unit_start_indicator {
                                let pes = Pes::read_from(&mut reader)?;
                                TsPayload::PesStart(pes)
                            } else {
                                let bytes = Bytes::read_from(&mut reader)?;
                                TsPayload::PesContinuation(bytes)
                            }
                        }
                    }
                }
            };
            Some(payload)
        } else {
            None
        };

        if reader.limit() != 0 {
            return Err(crate::Error::invalid_input(
                "Unexpected remaining data in TS packet",
            ));
        }
        Ok(Some(TsPacket {
            header,
            adaptation_field,
            payload,
        }))
    }
}

#[derive(Debug, Clone)]
enum PidKind {
    Pmt,
    Pes,
}