netpix_common/
mpegts.rs

1#![allow(unused_imports)]
2pub mod adaptation_field;
3pub mod aggregator;
4pub mod constants;
5pub mod descriptors;
6pub mod header;
7pub mod payload;
8pub mod pes;
9pub mod psi;
10#[cfg(test)]
11mod tests;
12
13use crate::mpegts::adaptation_field::AdaptationField;
14use crate::mpegts::header::Header;
15use crate::mpegts::header::{AdaptationFieldControl, PIDTable, TransportScramblingControl};
16use crate::mpegts::payload::RawPayload;
17use crate::utils::bits::BitReader;
18use bincode::{Decode, Encode};
19use constants::*;
20
21#[derive(Decode, Encode, Debug, Clone, Eq, PartialEq)]
22pub struct MpegtsPacket {
23    pub number_of_fragments: usize,
24    pub fragments: Vec<MpegtsFragment>,
25}
26
27#[derive(Decode, Encode, Debug, Clone, Eq, PartialEq)]
28pub struct MpegtsFragment {
29    pub header: Header,
30    pub adaptation_field: Option<AdaptationField>,
31    pub payload: Option<RawPayload>,
32    pub size: usize,
33}
34
35#[cfg(not(target_arch = "wasm32"))]
36impl MpegtsPacket {
37    pub fn build(packet: &super::Packet) -> Option<Self> {
38        packet
39            .payload
40            .as_ref()
41            .and_then(|payload| Self::unmarshall(payload))
42    }
43
44    fn unmarshall(buffer: &[u8]) -> Option<Self> {
45        if buffer.len() % FRAGMENT_SIZE != 0 || buffer.len() > FRAGMENT_SIZE * MAX_FRAGMENTS {
46            return None;
47        }
48
49        let fragments: Vec<_> = (0..buffer.len())
50            .step_by(FRAGMENT_SIZE)
51            .filter_map(|start_index| {
52                ((buffer[start_index] & SYNC_BYTE_MASK) == SYNC_BYTE)
53                    .then(|| Self::get_fragment(buffer, start_index, start_index / FRAGMENT_SIZE))
54                    .flatten()
55            })
56            .collect();
57
58        (!fragments.is_empty()).then_some(Self {
59            number_of_fragments: fragments.len(),
60            fragments,
61        })
62    }
63
64    fn get_fragment(
65        buffer: &[u8],
66        start_index: usize,
67        fragment_number: usize,
68    ) -> Option<MpegtsFragment> {
69        let header = Self::get_header(buffer, start_index)?;
70        let current_index = start_index + HEADER_SIZE;
71
72        let (adaptation_field, payload_start) =
73            Self::process_adaptation_field(&header, buffer, current_index)?;
74
75        let payload = Self::process_payload(&header, buffer, payload_start, fragment_number);
76
77        Some(MpegtsFragment {
78            header,
79            adaptation_field: adaptation_field.clone(),
80            payload: payload.clone(),
81            size: Self::calculate_fragment_size(&adaptation_field, &payload),
82        })
83    }
84
85    fn process_adaptation_field(
86        header: &Header,
87        buffer: &[u8],
88        start_index: usize,
89    ) -> Option<(Option<AdaptationField>, usize)> {
90        match header.adaptation_field_control {
91            AdaptationFieldControl::AdaptationFieldOnly
92            | AdaptationFieldControl::AdaptationFieldAndPayload => {
93                let field = AdaptationField::unmarshall(&buffer[start_index..])?;
94                let next_index = start_index + field.adaptation_field_length as usize + 1;
95                Some((Some(field), next_index))
96            }
97            _ => Some((None, start_index)),
98        }
99    }
100
101    fn process_payload(
102        header: &Header,
103        buffer: &[u8],
104        start_index: usize,
105        fragment_number: usize,
106    ) -> Option<RawPayload> {
107        match header.adaptation_field_control {
108            AdaptationFieldControl::PayloadOnly
109            | AdaptationFieldControl::AdaptationFieldAndPayload => {
110                Self::get_payload(buffer, start_index, fragment_number)
111            }
112            _ => None,
113        }
114    }
115
116    fn get_header(buffer: &[u8], start_index: usize) -> Option<Header> {
117        let reader = BitReader::at_position(buffer, start_index);
118        Some(Header {
119            transport_error_indicator: reader.get_bit(1, 7)?,
120            payload_unit_start_indicator: reader.get_bit(1, 6)?,
121            transport_priority: reader.get_bit(1, 5)?,
122            pid: PIDTable::from(reader.get_bits_u16(1, PID_MASK_UPPER, 0xFF)?),
123            transport_scrambling_control: match reader.get_bits(3, TSC_MASK, 6)? {
124                0 => TransportScramblingControl::NotScrambled,
125                val => TransportScramblingControl::UserDefined(val),
126            },
127            adaptation_field_control: match reader.get_bits(3, AFC_MASK, 4)? {
128                1 => AdaptationFieldControl::PayloadOnly,
129                2 => AdaptationFieldControl::AdaptationFieldOnly,
130                3 => AdaptationFieldControl::AdaptationFieldAndPayload,
131                _ => return None,
132            },
133            continuity_counter: reader.get_bits(3, CC_MASK, 0)?,
134        })
135    }
136
137    fn calculate_fragment_size(
138        adaptation_field: &Option<AdaptationField>,
139        payload: &Option<RawPayload>,
140    ) -> usize {
141        HEADER_SIZE
142            + adaptation_field
143                .as_ref()
144                .map_or(0, |af| af.adaptation_field_length as usize + 1)
145            + payload.as_ref().map_or(0, |p| p.data.len())
146    }
147
148    fn get_payload(
149        buffer: &[u8],
150        start_index: usize,
151        fragment_number: usize,
152    ) -> Option<RawPayload> {
153        let end_index = (fragment_number + 1) * FRAGMENT_SIZE;
154        let length = end_index.saturating_sub(start_index);
155
156        if length == 0 {
157            return None;
158        }
159
160        let data = buffer[start_index..end_index.min(buffer.len())].to_vec();
161        Some(RawPayload { data })
162    }
163}