use std::collections::VecDeque;
use std::pin::Pin;
use std::task::{Context, Poll};
use dvb_si::demux::{SectionEvent, SiDemux, SiDemuxBuilder};
use futures_core::Stream;
use tokio::io::AsyncRead;
use tokio::io::ReadBuf;
use crate::resync::{resync, TS_PACKET_SIZE, TS_SYNC_BYTE};
use crate::ResyncStats;
const READ_BUF_SIZE: usize = TS_PACKET_SIZE * 7;
pub struct SectionStream<R> {
reader: R,
demux: SiDemux,
queue: VecDeque<SectionEvent>,
buf: Vec<u8>,
filled: usize,
eof: bool,
synced: bool,
resync_stats: ResyncStats,
}
impl<R: AsyncRead + Unpin> SectionStream<R> {
#[must_use]
pub fn new(reader: R) -> Self {
Self::with_demux(reader, SiDemux::builder().build())
}
#[must_use]
pub fn with_builder(reader: R, builder: SiDemuxBuilder) -> Self {
Self::with_demux(reader, builder.build())
}
#[must_use]
pub fn with_demux(reader: R, demux: SiDemux) -> Self {
Self {
reader,
demux,
queue: VecDeque::new(),
buf: vec![0u8; READ_BUF_SIZE],
filled: 0,
eof: false,
synced: false,
resync_stats: ResyncStats::default(),
}
}
#[must_use]
pub fn stats(&self) -> dvb_si::demux::Stats {
self.demux.stats()
}
#[must_use]
pub fn resync_stats(&self) -> ResyncStats {
self.resync_stats
}
fn feed_buf(&mut self, data: &[u8]) {
let start = if self.synced {
0
} else {
match resync(data) {
Some(off) => {
self.synced = true;
self.resync_stats.resyncs += 1;
self.resync_stats.bytes_discarded += off as u64;
off
}
None => {
self.resync_stats.bytes_discarded += data.len() as u64;
return;
}
}
};
let aligned = &data[start..];
let n_packets = aligned.len() / TS_PACKET_SIZE;
for i in 0..n_packets {
let pkt_start = i * TS_PACKET_SIZE;
let pkt = &aligned[pkt_start..pkt_start + TS_PACKET_SIZE];
if pkt[0] != TS_SYNC_BYTE {
self.resync_stats.desyncs += 1;
let discarded = aligned.len() - pkt_start;
self.resync_stats.bytes_discarded += discarded as u64;
self.synced = false;
self.filled = 0;
return;
}
for event in self.demux.feed(pkt) {
self.queue.push_back(event);
}
}
let aligned_end = start + (data[start..].len() / TS_PACKET_SIZE) * TS_PACKET_SIZE;
let remainder = &data[aligned_end..];
if !remainder.is_empty() {
self.buf[..remainder.len()].copy_from_slice(remainder);
self.filled = remainder.len();
} else {
self.filled = 0;
}
}
}
impl<R: AsyncRead + Unpin> Stream for SectionStream<R> {
type Item = SectionEvent;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
if let Some(event) = this.queue.pop_front() {
return Poll::Ready(Some(event));
}
if this.eof {
return Poll::Ready(None);
}
let buf_len = this.buf.len();
let read_from = this.filled;
let mut read_buf = ReadBuf::new(&mut this.buf[read_from..buf_len]);
match Pin::new(&mut this.reader).poll_read(cx, &mut read_buf) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Err(_)) => {
this.eof = true;
return Poll::Ready(None);
}
Poll::Ready(Ok(())) => {
let n = read_buf.filled().len();
if n == 0 {
this.eof = true;
return Poll::Ready(None);
}
let total = read_from + n;
let data: Vec<u8> = this.buf[..total].to_vec();
this.feed_buf(&data);
}
}
}
}
}
#[cfg(feature = "udp")]
pub struct UdpReader {
pub(crate) socket: tokio::net::UdpSocket,
}
#[cfg(feature = "udp")]
impl AsyncRead for UdpReader {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
self.get_mut().socket.poll_recv(cx, buf)
}
}
#[cfg(feature = "udp")]
impl SectionStream<UdpReader> {
pub async fn bind_multicast(
bind_addr: std::net::SocketAddrV4,
multicast_addr: std::net::Ipv4Addr,
) -> std::io::Result<Self> {
use tokio::net::UdpSocket;
let socket = UdpSocket::bind(bind_addr).await?;
socket.join_multicast_v4(multicast_addr, *bind_addr.ip())?;
Ok(Self::new(UdpReader { socket }))
}
}