AVI demuxer
[nihav.git] / src / demuxers / avi.rs
CommitLineData
ca41f125
KS
1use super::*;
2use super::DemuxerError::*;
3use io::byteio::*;
4use frame::*;
5use formats::*;
6
7macro_rules! mktag {
8 ($a:expr, $b:expr, $c:expr, $d:expr) => ({
9 (($a as u32) << 24) | (($b as u32) << 16) | (($c as u32) << 8) | ($d as u32)
10 });
11 ($arr:expr) => ({
12 (($arr[0] as u32) << 24) | (($arr[1] as u32) << 16) | (($arr[2] as u32) << 8) | ($arr[3] as u32)
13 });
14}
15
16struct StreamState {
17 strm_no: u8,
18 got_strf: bool,
19 strm_type: Option<StreamType>,
20}
21
22impl StreamState {
23 fn new() -> Self {
24 StreamState { strm_no: 0, got_strf: true, strm_type: None }
25 }
26 fn reset(&mut self) {
27 self.strm_type = None;
28 self.got_strf = true;
29 self.strm_no += 1;
30 }
31 fn valid_state(&self) -> bool {
32 match self.strm_type {
33 None => self.got_strf,
34 _ => false,
35 }
36 }
37}
38
39#[allow(dead_code)]
40pub struct AVIDemuxer<'a> {
41 opened: bool,
42 src: &'a mut ByteReader<'a>,
43 cur_frame: Vec<u64>,
44 num_streams: u8,
45 dmx: Demuxer,
46 size: usize,
47 movi_size: usize,
48 sstate: StreamState,
49}
50
51#[derive(Debug,Clone,Copy)]
52enum RIFFTag {
53 Chunk(u32),
54 List(u32,u32),
55}
56
57struct RIFFParser {
58 tag: RIFFTag,
59 parse: fn(&mut AVIDemuxer, size: usize) -> DemuxerResult<usize>,
60}
61
62impl<'a> Demux<'a> for AVIDemuxer<'a> {
63 #[allow(unused_variables)]
64 fn open(&mut self) -> DemuxerResult<()> {
65 self.read_header()?;
66 self.opened = true;
67 Ok(())
68 }
69
70 fn get_frame(&mut self) -> DemuxerResult<NAPacket> {
71 if !self.opened { return Err(NoSuchInput); }
72 if self.movi_size == 0 { return Err(EOF); }
73 let mut tag: [u8; 4] = [0; 4];
74 loop {
75 self.src.read_buf(&mut tag)?;
76 let size = self.src.read_u32le()? as usize;
77 if mktag!(tag) == mktag!(b"JUNK") {
78 self.movi_size -= size + 8;
79 self.src.read_skip(size)?;
80 continue;
81 }
82 if tag[0] < b'0' || tag[0] > b'9' || tag[1] < b'0' || tag[1] > b'9' {
83 return Err(InvalidData);
84 }
85 let stream_no = (tag[0] - b'0') * 10 + (tag[1] - b'0');
86 let str = self.dmx.get_stream(stream_no as usize);
87 if let None = str { return Err(InvalidData); }
88 let stream = str.unwrap();
89 let pkt = self.src.read_packet(stream, Some(self.cur_frame[stream_no as usize]), None, None, false, size)?;
90 self.cur_frame[stream_no as usize] += 1;
91 self.movi_size -= size + 8;
92
93 return Ok(pkt);
94 }
95 }
96
97 #[allow(unused_variables)]
98 fn seek(&mut self, time: u64) -> DemuxerResult<()> {
99 if !self.opened { return Err(NoSuchInput); }
100 Err(NotImplemented)
101 }
102}
103
104impl<'a> AVIDemuxer<'a> {
105 pub fn new(io: &'a mut ByteReader<'a>) -> Self {
106 AVIDemuxer {
107 cur_frame: Vec::new(),
108 num_streams: 0,
109 opened: false,
110 src: io,
111 size: 0,
112 movi_size: 0,
113 sstate: StreamState::new(),
114 dmx: Demuxer::new()
115 }
116 }
117
118 fn parse_chunk(&mut self, end_tag: RIFFTag, csize: usize, depth: u16) -> DemuxerResult<(usize, bool)> {
119 if csize < 8 { return Err(InvalidData); }
120 if depth > 42 { return Err(InvalidData); }
121
122 let tag = self.src.read_u32be()?;
123 let size = self.src.read_u32le()? as usize;
124 if size > csize { return Err(InvalidData); }
125 if let RIFFTag::Chunk(etag) = end_tag {
126 if etag == tag { return Ok((size, true)); }
127 }
128 let ltag = if is_list_tag(tag) { self.src.read_u32be()? } else { 0 };
129 if let RIFFTag::List(etag1, etag2) = end_tag {
130 if etag1 == tag && etag2 == ltag {
131 return Ok((size, true));
132 }
133 }
134
135 for i in 0..CHUNKS.len() {
136 if let RIFFTag::Chunk(ctag) = CHUNKS[i].tag {
137 if ctag == tag {
138 let psize = (CHUNKS[i].parse)(self, size)?;
139 if psize != size { return Err(InvalidData); }
140 return Ok((size + 8, false));
141 }
142 }
143 if let RIFFTag::List(ctag, clst) = CHUNKS[i].tag {
144 if ctag == tag && clst == ltag {
145 let mut rest_size = size - 4;
146 let psize = (CHUNKS[i].parse)(self, rest_size)?;
147 if psize > rest_size { return Err(InvalidData); }
148 rest_size -= psize;
149 while rest_size > 0 {
150 let (psize, _) = self.parse_chunk(end_tag, rest_size, depth+1)?;
151 if psize > rest_size { return Err(InvalidData); }
152 rest_size -= psize;
153 }
154
155 return Ok((size + 8, false));
156 }
157 }
158 }
159 self.src.read_skip(size)?;
160 return Ok((size + 8, false));
161 }
162
163 fn read_header(&mut self) -> DemuxerResult<()> {
164 let riff_tag = self.src.read_u32be()?;
165 let size = self.src.read_u32le()? as usize;
166 let avi_tag = self.src.read_u32be()?;
167 if riff_tag != mktag!(b"RIFF") || avi_tag != mktag!(b"AVI ") {
168 return Err(InvalidData);
169 }
170 self.size = size;
171 let mut rest_size = size;
172 loop {
173 let (csz, end) = self.parse_chunk(RIFFTag::List(mktag!(b"LIST"), mktag!(b"movi")), rest_size,0)?;
174 if end { self.movi_size = csz - 4; break; }
175 rest_size -= csz;
176 }
177 if !self.sstate.valid_state() || self.sstate.strm_no != self.num_streams {
178 return Err(InvalidData);
179 }
180 Ok(())
181 }
182
183 fn read_extradata(&mut self, size: usize) -> DemuxerResult<Option<Vec<u8>>> {
184 if size == 0 { return Ok(None); }
185 let mut edvec: Vec<u8> = Vec::with_capacity(size);
186 edvec.resize(size, 0);
187 self.src.read_buf(&mut edvec)?;
188 Ok(Some(edvec))
189 }
190}
191
192const CHUNKS: [RIFFParser; 6] = [
193 RIFFParser { tag: RIFFTag::List(mktag!(b"LIST"), mktag!(b"hdrl")), parse: parse_hdrl },
194 RIFFParser { tag: RIFFTag::List(mktag!(b"LIST"), mktag!(b"strl")), parse: parse_strl },
195 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"avih")), parse: parse_avih },
196 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"strf")), parse: parse_strf },
197 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"strh")), parse: parse_strh },
198 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"JUNK")), parse: parse_junk },
199];
200
201fn is_list_tag(tag: u32) -> bool {
202 for i in 0..CHUNKS.len() {
203 if let RIFFTag::List(ltag, _) = CHUNKS[i].tag {
204 if tag == ltag {
205 return true;
206 }
207 }
208 }
209 false
210}
211
212#[allow(unused_variables)]
213fn parse_hdrl(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
214 Ok(0)
215}
216
217#[allow(unused_variables)]
218fn parse_strl(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
219 Ok(0)
220}
221
222#[allow(unused_variables)]
223fn parse_strh(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
224 if size < 0x38 { return Err(InvalidData); }
225 let tag = dmx.src.read_u32be()?; //stream type
226 let fcc = dmx.src.read_u32be()?; //handler(fourcc)
227 dmx.src.read_u32le()?; //flags
228 dmx.src.read_skip(2)?; //priority
229 dmx.src.read_skip(2)?; //language
230 dmx.src.read_skip(4)?; //initial frames
231 let scale = dmx.src.read_u32le()?; //scale
232 let rate = dmx.src.read_u32le()?; //rate
233 dmx.src.read_skip(4)?; //start
234 dmx.src.read_skip(4)?; //length
235 dmx.src.read_skip(4)?; //buf size
236 dmx.src.read_skip(4)?; //quality
237 dmx.src.read_skip(4)?; //sample size
238 let a = dmx.src.read_u16le()?;
239 let b = dmx.src.read_u16le()?;
240 let c = dmx.src.read_u16le()?;
241 let d = dmx.src.read_u16le()?;
242
243 dmx.src.read_skip(size - 0x38)?;
244
245 if !dmx.sstate.valid_state() || dmx.sstate.strm_no >= dmx.num_streams {
246 return Err(InvalidData);
247 }
248 if tag == mktag!(b"vids") {
249 dmx.sstate.strm_type = Some(StreamType::Video);
250 } else if tag == mktag!(b"auds") {
251 dmx.sstate.strm_type = Some(StreamType::Audio);
252 } else {
253 dmx.sstate.strm_type = Some(StreamType::Data);
254 }
255 dmx.sstate.got_strf = false;
256
257 Ok(size)
258}
259
260fn parse_strf(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
261 if let None = dmx.sstate.strm_type { return Err(InvalidData); }
262 match dmx.sstate.strm_type.unwrap() {
263 StreamType::Video => parse_strf_vids(dmx, size),
264 StreamType::Audio => parse_strf_auds(dmx, size),
265 _ => parse_strf_xxxx(dmx, size),
266 }
267}
268
269#[allow(unused_variables)]
270fn parse_strf_vids(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
271 if size < 40 { return Err(InvalidData); }
272 let bi_size = dmx.src.read_u32le()?;
273 if (bi_size as usize) > size { return Err(InvalidData); }
274 let width = dmx.src.read_u32le()?;
275 let height = dmx.src.read_u32le()? as i32;
276 let planes = dmx.src.read_u16le()?;
277 let bitcount = dmx.src.read_u16le()?;
278 let compression = dmx.src.read_u32be()?;
279 let img_size = dmx.src.read_u32le()?;
280 let xdpi = dmx.src.read_u32le()?;
281 let ydpi = dmx.src.read_u32le()?;
282 let colors = dmx.src.read_u32le()?;
283 let imp_colors = dmx.src.read_u32le()?;
284
285 let flip = height < 0;
286 let format = if bitcount > 8 { RGB24_FORMAT } else { PAL8_FORMAT };
287 let vhdr = NAVideoInfo::new(width, if flip { -height as u32 } else { height as u32}, flip, PAL8_FORMAT);
288 let vci = NACodecTypeInfo::Video(vhdr);
289 let edata = dmx.read_extradata(size - 40)?;
290 let vinfo = NACodecInfo::new(vci, edata);
291 let res = dmx.dmx.add_stream(NAStream::new(StreamType::Video, dmx.sstate.strm_no as u32, vinfo));
292 if let None = res { return Err(MemoryError); }
293 dmx.sstate.reset();
294 Ok(size)
295}
296
297#[allow(unused_variables)]
298fn parse_strf_auds(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
299 if size < 16 { return Err(InvalidData); }
300 let w_format_tag = dmx.src.read_u16le()?;
301 let channels = dmx.src.read_u16le()?;
302 let samplespersec = dmx.src.read_u32le()?;
303 let avgbytespersec = dmx.src.read_u32le()?;
304 let block_align = dmx.src.read_u16le()?;
305 let bits_per_sample = dmx.src.read_u16le()?;
306
307 let soniton = NASoniton::new(bits_per_sample as u8, SONITON_FLAG_SIGNED);
308 let ahdr = NAAudioInfo::new(samplespersec, channels as u8, soniton, block_align as usize);
309 let edata = dmx.read_extradata(size - 16)?;
310 let ainfo = NACodecInfo::new(NACodecTypeInfo::Audio(ahdr), edata);
311 let res = dmx.dmx.add_stream(NAStream::new(StreamType::Audio, dmx.sstate.strm_no as u32, ainfo));
312 if let None = res { return Err(MemoryError); }
313 dmx.sstate.reset();
314 Ok(size)
315}
316
317fn parse_strf_xxxx(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
318 let edata = dmx.read_extradata(size)?;
319 let info = NACodecInfo::new(NACodecTypeInfo::None, edata);
320 let res = dmx.dmx.add_stream(NAStream::new(StreamType::Data, dmx.sstate.strm_no as u32, info));
321 if let None = res { return Err(MemoryError); }
322 dmx.sstate.reset();
323 Ok(size)
324}
325
326#[allow(unused_variables)]
327fn parse_avih(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
328 if size < 0x38 { return Err(InvalidData); }
329 let timebase = dmx.src.read_u32le()?; //microsec per frame
330 dmx.src.read_skip(4)?; //max frame size
331 dmx.src.read_skip(4)?; //padding
332 dmx.src.read_u32le()?; //flags
333 let frames = dmx.src.read_u32le()?; //frames
334 dmx.src.read_skip(4)?; //initial frames
335 let streams = dmx.src.read_u32le()?; //streams
336 if streams > 100 { return Err(InvalidData); }
337 dmx.num_streams = streams as u8;
338
339 dmx.src.read_skip(4)?; //buf size
340 let width = dmx.src.read_u32le()?; //width
341 let height = dmx.src.read_u32le()? as i32; //height
342 dmx.src.read_skip(16)?; //reserved
343
344 dmx.cur_frame.resize(streams as usize, 0);
345 dmx.src.read_skip(size - 0x38)?;
346 Ok(size)
347}
348
349fn parse_junk(dmx: &mut AVIDemuxer, size: usize) -> DemuxerResult<usize> {
350 dmx.src.read_skip(size)?;
351 Ok(size)
352}
353
354#[cfg(test)]
355mod test {
356 use super::*;
357 use std::fs::File;
358
359 #[test]
360 fn test_avi_demux() {
361 let mut file = File::open("assets/laser05.avi").unwrap();
362 let mut fr = FileReader::new_read(&mut file);
363 let mut br = ByteReader::new(&mut fr);
364 let mut dmx = AVIDemuxer::new(&mut br);
365 dmx.open().unwrap();
366
367 loop {
368 let pktres = dmx.get_frame();
369 if let Err(e) = pktres {
370 if (e as i32) == (DemuxerError::EOF as i32) { break; }
371 panic!("error");
372 }
373 let pkt = pktres.unwrap();
374 println!("Got {}", pkt);
375 }
376 }
377}