av1_obu_parser/
lib.rs

1pub mod buffer;
2pub mod obu;
3
4/// IVF container parsing helpers for AV1 bitstreams.
5///
6/// IVF is a simple byte-oriented container:
7/// - file header: 32 bytes, beginning with the `DKIF` signature
8/// - repeated frame records: 12-byte frame header + frame payload
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct IvfHeader {
11    pub version: u16,
12    pub header_size: u16,
13    pub codec: [u8; 4],
14    pub width: u16,
15    pub height: u16,
16    pub timebase_denominator: u32,
17    pub timebase_numerator: u32,
18    pub num_frames: u32,
19    pub unused: u32,
20}
21
22impl IvfHeader {
23    pub fn codec_string(&self) -> &str {
24        std::str::from_utf8(&self.codec).unwrap_or("unknown")
25    }
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct IvfFrame<'a> {
30    pub index: usize,
31    pub pts: u64,
32    pub data: &'a [u8],
33}
34
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum IvfError {
37    InvalidSignature,
38    TruncatedHeader,
39    InvalidHeaderSize(u16),
40    TruncatedFrameHeader {
41        frame_index: usize,
42    },
43    TruncatedFrameData {
44        frame_index: usize,
45        frame_size: usize,
46        remaining_bytes: usize,
47    },
48}
49
50impl std::fmt::Display for IvfError {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        match self {
53            IvfError::InvalidSignature => {
54                write!(f, "file does not start with IVF signature 'DKIF'")
55            }
56            IvfError::TruncatedHeader => write!(f, "IVF file must be at least 32 bytes"),
57            IvfError::InvalidHeaderSize(size) => write!(f, "invalid IVF header size: {size}"),
58            IvfError::TruncatedFrameHeader { frame_index } => {
59                write!(f, "truncated IVF frame header at frame {frame_index}")
60            }
61            IvfError::TruncatedFrameData {
62                frame_index,
63                frame_size,
64                remaining_bytes,
65            } => write!(
66                f,
67                "truncated IVF frame data at frame {frame_index}: expected {frame_size} bytes, only {remaining_bytes} remain"
68            ),
69        }
70    }
71}
72
73impl std::error::Error for IvfError {}
74
75pub struct IvfReader<'a> {
76    data: &'a [u8],
77    header: IvfHeader,
78}
79
80impl<'a> IvfReader<'a> {
81    pub fn new(data: &'a [u8]) -> Result<Self, IvfError> {
82        if data.len() < 32 {
83            return Err(IvfError::TruncatedHeader);
84        }
85        if !data.starts_with(b"DKIF") {
86            return Err(IvfError::InvalidSignature);
87        }
88
89        let header_size = u16::from_le_bytes([data[6], data[7]]);
90        if header_size < 32 || header_size as usize > data.len() {
91            return Err(IvfError::InvalidHeaderSize(header_size));
92        }
93
94        let header = IvfHeader {
95            version: u16::from_le_bytes([data[4], data[5]]),
96            header_size,
97            codec: [data[8], data[9], data[10], data[11]],
98            width: u16::from_le_bytes([data[12], data[13]]),
99            height: u16::from_le_bytes([data[14], data[15]]),
100            timebase_denominator: u32::from_le_bytes([data[16], data[17], data[18], data[19]]),
101            timebase_numerator: u32::from_le_bytes([data[20], data[21], data[22], data[23]]),
102            num_frames: u32::from_le_bytes([data[24], data[25], data[26], data[27]]),
103            unused: u32::from_le_bytes([data[28], data[29], data[30], data[31]]),
104        };
105
106        Ok(Self { data, header })
107    }
108
109    pub fn header(&self) -> &IvfHeader {
110        &self.header
111    }
112
113    pub fn frames(&self) -> IvfFrames<'a> {
114        IvfFrames {
115            data: self.data,
116            offset: self.header.header_size as usize,
117            frame_index: 0,
118            failed: false,
119        }
120    }
121}
122
123pub struct IvfFrames<'a> {
124    data: &'a [u8],
125    offset: usize,
126    frame_index: usize,
127    failed: bool,
128}
129
130impl<'a> Iterator for IvfFrames<'a> {
131    type Item = Result<IvfFrame<'a>, IvfError>;
132
133    fn next(&mut self) -> Option<Self::Item> {
134        if self.failed || self.offset >= self.data.len() {
135            return None;
136        }
137
138        if self.offset + 12 > self.data.len() {
139            self.failed = true;
140            return Some(Err(IvfError::TruncatedFrameHeader {
141                frame_index: self.frame_index,
142            }));
143        }
144
145        let frame_size = u32::from_le_bytes([
146            self.data[self.offset],
147            self.data[self.offset + 1],
148            self.data[self.offset + 2],
149            self.data[self.offset + 3],
150        ]) as usize;
151        
152        let pts = u64::from_le_bytes([
153            self.data[self.offset + 4],
154            self.data[self.offset + 5],
155            self.data[self.offset + 6],
156            self.data[self.offset + 7],
157            self.data[self.offset + 8],
158            self.data[self.offset + 9],
159            self.data[self.offset + 10],
160            self.data[self.offset + 11],
161        ]);
162
163        let data_offset = self.offset + 12;
164
165        if data_offset + frame_size > self.data.len() {
166            self.failed = true;
167            return Some(Err(IvfError::TruncatedFrameData {
168                frame_index: self.frame_index,
169                frame_size,
170                remaining_bytes: self.data.len().saturating_sub(data_offset),
171            }));
172        }
173
174        let frame = IvfFrame {
175            index: self.frame_index,
176            pts,
177            data: &self.data[data_offset..data_offset + frame_size],
178        };
179
180        self.offset = data_offset + frame_size;
181        self.frame_index += 1;
182        Some(Ok(frame))
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    fn build_ivf(frame_payloads: &[&[u8]]) -> Vec<u8> {
191        let mut data = Vec::new();
192        data.extend_from_slice(b"DKIF");
193        data.extend_from_slice(&0u16.to_le_bytes());
194        data.extend_from_slice(&32u16.to_le_bytes());
195        data.extend_from_slice(b"AV01");
196        data.extend_from_slice(&640u16.to_le_bytes());
197        data.extend_from_slice(&480u16.to_le_bytes());
198        data.extend_from_slice(&30u32.to_le_bytes());
199        data.extend_from_slice(&1u32.to_le_bytes());
200        data.extend_from_slice(&(frame_payloads.len() as u32).to_le_bytes());
201        data.extend_from_slice(&0u32.to_le_bytes());
202
203        for (index, payload) in frame_payloads.iter().enumerate() {
204            data.extend_from_slice(&(payload.len() as u32).to_le_bytes());
205            data.extend_from_slice(&(index as u64).to_le_bytes());
206            data.extend_from_slice(payload);
207        }
208
209        data
210    }
211
212    #[test]
213    fn test_parse_ivf_header() {
214        let data = build_ivf(&[b"\x12\x00"]);
215        let ivf = IvfReader::new(&data).expect("valid IVF header");
216
217        assert_eq!(ivf.header().codec, *b"AV01");
218        assert_eq!(ivf.header().codec_string(), "AV01");
219        assert_eq!(ivf.header().width, 640);
220        assert_eq!(ivf.header().height, 480);
221        assert_eq!(ivf.header().num_frames, 1);
222    }
223
224    #[test]
225    fn test_iterate_ivf_frames() {
226        let data = build_ivf(&[b"\x12\x00", b"\x0A\x01\x02"]);
227        let ivf = IvfReader::new(&data).expect("valid IVF data");
228        let frames = ivf
229            .frames()
230            .collect::<Result<Vec<_>, _>>()
231            .expect("valid IVF frames");
232
233        assert_eq!(frames.len(), 2);
234        assert_eq!(frames[0].index, 0);
235        assert_eq!(frames[0].pts, 0);
236        assert_eq!(frames[0].data, b"\x12\x00");
237        assert_eq!(frames[1].index, 1);
238        assert_eq!(frames[1].pts, 1);
239        assert_eq!(frames[1].data, b"\x0A\x01\x02");
240    }
241}
av1_obu_parser/lib.rs

av1_obu_parser/
lib.rs
