1use std::{
2 io::{self, Read, Write},
3 sync::Arc,
4};
5
6use bitflags::bitflags;
7use media_core::{buffer::Buffer, error::Error, invalid_param_error, Result};
8use num_rational::Rational64;
9
10bitflags! {
11 #[repr(transparent)]
12 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
13 pub struct PacketFlags: u32 {
14 const Key = 1;
15 const Corrupt = 2;
16 }
17}
18
19#[derive(Clone)]
20enum PacketData<'a> {
21 Borrowed(&'a [u8]),
22 Owned(Vec<u8>),
23 Buffer(Arc<Buffer>),
24}
25
26impl PacketData<'_> {
27 fn from_slice(slice: &[u8]) -> PacketData<'_> {
28 PacketData::Borrowed(slice)
29 }
30
31 fn from_vec(vec: Vec<u8>) -> PacketData<'static> {
32 PacketData::Owned(vec)
33 }
34
35 fn from_buffer(buffer: Arc<Buffer>) -> PacketData<'static> {
36 PacketData::Buffer(buffer)
37 }
38
39 fn as_slice(&self) -> &[u8] {
40 match self {
41 PacketData::Borrowed(slice) => slice,
42 PacketData::Owned(vec) => vec.as_slice(),
43 PacketData::Buffer(buffer) => buffer.data(),
44 }
45 }
46
47 fn as_mut_slice(&mut self) -> Option<&mut [u8]> {
48 match self {
49 PacketData::Borrowed(_) => None,
50 PacketData::Owned(ref mut vec) => Some(vec.as_mut_slice()),
51 PacketData::Buffer(buffer) => {
52 let mut_buffer = Arc::get_mut(buffer)?;
53 Some(mut_buffer.data_mut())
54 }
55 }
56 }
57
58 fn len(&self) -> usize {
59 match self {
60 PacketData::Borrowed(slice) => slice.len(),
61 PacketData::Owned(vec) => vec.len(),
62 PacketData::Buffer(buffer) => buffer.len(),
63 }
64 }
65
66 fn is_empty(&self) -> bool {
67 self.len() == 0
68 }
69
70 fn capacity(&self) -> usize {
71 match &self {
72 PacketData::Borrowed(slice) => slice.len(),
73 PacketData::Owned(vec) => vec.capacity(),
74 PacketData::Buffer(buffer) => buffer.capacity(),
75 }
76 }
77
78 fn into_owned(self) -> PacketData<'static> {
79 match self {
80 PacketData::Borrowed(slice) => PacketData::Owned(slice.to_vec()),
81 PacketData::Owned(vec) => PacketData::Owned(vec),
82 PacketData::Buffer(buffer) => PacketData::Buffer(buffer),
83 }
84 }
85}
86
87#[derive(Clone)]
88pub struct Packet<'a> {
89 pub pts: Option<i64>,
90 pub dts: Option<i64>,
91 pub duration: Option<i64>,
92 pub time_base: Option<Rational64>,
93 pub flags: PacketFlags,
94 pub pos: Option<usize>,
95 pub stream_index: Option<usize>,
96 data: PacketData<'a>,
97}
98
99impl<'a> Packet<'a> {
100 fn from_data(data: PacketData<'a>) -> Self {
101 Self {
102 pts: None,
103 dts: None,
104 duration: None,
105 time_base: None,
106 flags: PacketFlags::empty(),
107 pos: None,
108 stream_index: None,
109 data,
110 }
111 }
112
113 pub fn new(size: usize) -> Self {
114 Self::from_data(PacketData::from_vec(vec![0; size]))
115 }
116
117 pub fn with_capacity(capacity: usize) -> Self {
118 Self::from_data(PacketData::from_vec(Vec::with_capacity(capacity)))
119 }
120
121 pub fn from_slice(data: &'a [u8]) -> Self {
122 Self::from_data(PacketData::from_slice(data))
123 }
124
125 pub fn from_buffer(buffer: Arc<Buffer>) -> Self {
126 Self::from_data(PacketData::from_buffer(buffer))
127 }
128
129 pub fn into_owned(self) -> Packet<'static> {
130 Packet {
131 pts: self.pts,
132 dts: self.dts,
133 duration: self.duration,
134 time_base: self.time_base,
135 flags: self.flags,
136 pos: self.pos,
137 stream_index: self.stream_index,
138 data: self.data.into_owned(),
139 }
140 }
141
142 pub fn len(&self) -> usize {
143 self.data.len()
144 }
145
146 pub fn is_empty(&self) -> bool {
147 self.data.is_empty()
148 }
149
150 pub fn capacity(&self) -> usize {
151 self.data.capacity()
152 }
153
154 pub fn data(&self) -> &[u8] {
155 self.data.as_slice()
156 }
157
158 pub fn data_mut(&mut self) -> Option<&mut [u8]> {
159 self.data.as_mut_slice()
160 }
161
162 pub fn truncate(&mut self, len: usize) -> Result<()> {
163 let current_len = self.data.len();
164 if len > current_len {
165 return Err(invalid_param_error!(len));
166 }
167
168 match &mut self.data {
169 PacketData::Borrowed(slice) => {
170 self.data = PacketData::Borrowed(&slice[..len]);
171 }
172 PacketData::Owned(vec) => {
173 vec.truncate(len);
174 }
175 PacketData::Buffer(buffer) => {
176 let buffer = Arc::get_mut(buffer).ok_or_else(|| Error::Invalid("buffer is shared".to_string()))?;
177 buffer.resize(len);
178 }
179 }
180
181 Ok(())
182 }
183}
184
185pub trait ReadPacket: Read {
186 fn read_packet(&mut self, size: usize) -> io::Result<Packet<'_>> {
187 let mut packet = Packet::new(size);
188
189 if let Some(data_mut) = packet.data_mut() {
190 self.read_exact(data_mut)?;
191 } else {
192 unreachable!()
194 }
195
196 Ok(packet)
197 }
198}
199
200pub trait WritePacket: Write {
201 fn write_packet(&mut self, packet: &Packet) -> io::Result<()> {
202 self.write_all(packet.data())
203 }
204}
205
206impl<T: Read + ?Sized> ReadPacket for T {}
207impl<T: Write + ?Sized> WritePacket for T {}