avi: end demuxing when index is encountered
[nihav.git] / nihav-commonfmt / src / demuxers / avi.rs
1 use nihav_core::demuxers::*;
2 use nihav_registry::register;
3 use nihav_core::demuxers::DemuxerError::*;
4
5 macro_rules! mktag {
6 ($a:expr, $b:expr, $c:expr, $d:expr) => ({
7 (($a as u32) << 24) | (($b as u32) << 16) | (($c as u32) << 8) | ($d as u32)
8 });
9 ($arr:expr) => ({
10 (($arr[0] as u32) << 24) | (($arr[1] as u32) << 16) | (($arr[2] as u32) << 8) | ($arr[3] as u32)
11 });
12 }
13
14 struct StreamState {
15 strm_no: u8,
16 got_strf: bool,
17 strm_type: Option<StreamType>,
18 }
19
20 impl StreamState {
21 fn new() -> Self {
22 StreamState { strm_no: 0, got_strf: true, strm_type: None }
23 }
24 fn reset(&mut self) {
25 self.strm_type = None;
26 self.got_strf = true;
27 self.strm_no += 1;
28 }
29 fn valid_state(&self) -> bool {
30 match self.strm_type {
31 None => self.got_strf,
32 _ => false,
33 }
34 }
35 }
36
37 struct PalInfo {
38 pal: Arc<[u8; 1024]>,
39 changed: bool,
40 stream_no: usize,
41 }
42
43 #[allow(dead_code)]
44 struct AVIDemuxer<'a> {
45 src: &'a mut ByteReader<'a>,
46 cur_frame: Vec<u64>,
47 num_streams: u8,
48 size: usize,
49 movi_size: usize,
50 movi_pos: u64,
51 movi_orig: usize,
52 sstate: StreamState,
53 tb_num: u32,
54 tb_den: u32,
55 key_offs: Vec<u64>,
56 pal: Vec<PalInfo>,
57 }
58
59 #[derive(Debug,Clone,Copy,PartialEq)]
60 enum RIFFTag {
61 Chunk(u32),
62 List(u32,u32),
63 }
64
65 struct RIFFParser {
66 tag: RIFFTag,
67 parse: fn(&mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize>,
68 }
69
70 impl<'a> DemuxCore<'a> for AVIDemuxer<'a> {
71 fn open(&mut self, strmgr: &mut StreamManager, seek_index: &mut SeekIndex) -> DemuxerResult<()> {
72 self.read_header(strmgr, seek_index)?;
73 Ok(())
74 }
75
76 fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> {
77 if self.movi_size == 0 { return Err(EOF); }
78 let mut tag: [u8; 4] = [0; 4];
79 loop {
80 if (self.src.tell() & 1) == 1 {
81 self.src.read_skip(1)?;
82 self.movi_size -= 1;
83 if self.movi_size == 0 { return Err(EOF); }
84 }
85 let is_keyframe = self.key_offs.binary_search(&self.src.tell()).is_ok();
86 self.src.read_buf(&mut tag)?;
87 let size = self.src.read_u32le()? as usize;
88 if mktag!(tag) == mktag!(b"JUNK") {
89 self.movi_size -= size + 8;
90 self.src.read_skip(size)?;
91 if self.movi_size == 0 { return Err(EOF); }
92 continue;
93 }
94 if mktag!(tag) == mktag!(b"LIST") {
95 self.movi_size -= 12;
96 self.src.read_skip(4)?;
97 if self.movi_size == 0 { return Err(EOF); }
98 continue;
99 }
100 if (tag[0] == b'i' && tag[1] == b'x') || (&tag == b"idx1") {
101 return Err(EOF);
102 }
103 if tag[0] < b'0' || tag[0] > b'9' || tag[1] < b'0' || tag[1] > b'9' {
104 return Err(InvalidData);
105 }
106 let stream_no = (tag[0] - b'0') * 10 + (tag[1] - b'0');
107 if tag[2] == b'p' && tag[3] == b'c' {
108 self.parse_palette_change(stream_no as usize, size)?;
109 self.movi_size -= size;
110 if self.movi_size == 0 { return Err(EOF); }
111 continue;
112 }
113 let str = strmgr.get_stream(stream_no as usize);
114 if str.is_none() { return Err(InvalidData); }
115 let stream = str.unwrap();
116 if size == 0 {
117 self.movi_size -= 8;
118 if self.movi_size == 0 { return Err(EOF); }
119 continue;
120 }
121 let (tb_num, tb_den) = stream.get_timebase();
122 let ts = NATimeInfo::new(Some(self.cur_frame[stream_no as usize]), None, None, tb_num, tb_den);
123 let mut pkt = self.src.read_packet(stream, ts, is_keyframe, size)?;
124 for pe in self.pal.iter_mut() {
125 if pe.stream_no == (stream_no as usize) {
126 pkt.add_side_data(NASideData::Palette(pe.changed, pe.pal.clone()));
127 pe.changed = false;
128 break;
129 }
130 }
131 self.cur_frame[stream_no as usize] += 1;
132 self.movi_size -= size + 8;
133
134 return Ok(pkt);
135 }
136 }
137
138 fn seek(&mut self, time: u64, seek_index: &SeekIndex) -> DemuxerResult<()> {
139 let ret = seek_index.find_pos(time);
140 if ret.is_none() {
141 return Err(DemuxerError::SeekError);
142 }
143 let seek_info = ret.unwrap();
144
145 if seek_info.pos < self.movi_pos { return Err(DemuxerError::SeekError); }
146 let skip_size = (seek_info.pos - self.movi_pos) as usize;
147 if skip_size > self.movi_size { return Err(DemuxerError::SeekError); }
148 self.movi_size = self.movi_orig - skip_size;
149
150 self.cur_frame[seek_info.str_id as usize] = seek_info.pts;
151 self.src.seek(SeekFrom::Start(seek_info.pos))?;
152
153 Ok(())
154 }
155 }
156
157 impl<'a> AVIDemuxer<'a> {
158 fn new(io: &'a mut ByteReader<'a>) -> Self {
159 AVIDemuxer {
160 cur_frame: Vec::new(),
161 num_streams: 0,
162 src: io,
163 size: 0,
164 movi_size: 0,
165 movi_pos: 0,
166 movi_orig: 0,
167 sstate: StreamState::new(),
168 tb_num: 0,
169 tb_den: 0,
170 key_offs: Vec::new(),
171 pal: Vec::new(),
172 }
173 }
174
175 fn parse_chunk(&mut self, strmgr: &mut StreamManager, end_tag: RIFFTag, csize: usize, depth: u16) -> DemuxerResult<(usize, bool)> {
176 if csize < 8 { return Err(InvalidData); }
177 if depth > 42 { return Err(InvalidData); }
178
179 let tag = self.src.read_u32be()?;
180 let size = self.src.read_u32le()? as usize;
181 if size > csize { return Err(InvalidData); }
182 if RIFFTag::Chunk(tag) == end_tag {
183 return Ok((size, true));
184 }
185 let is_list = is_list_tag(tag);
186 let ltag = if is_list { self.src.read_u32be()? } else { 0 };
187 if RIFFTag::List(tag, ltag) == end_tag {
188 return Ok((size, true));
189 }
190
191 for chunk in CHUNKS.iter() {
192 if RIFFTag::Chunk(tag) == chunk.tag {
193 let psize = (chunk.parse)(self, strmgr, size)?;
194 if psize != size { return Err(InvalidData); }
195 if (psize & 1) == 1 { self.src.read_skip(1)?; }
196 return Ok((size + 8, false));
197 }
198 if RIFFTag::List(tag, ltag) == chunk.tag {
199 let mut rest_size = size - 4;
200 let psize = (chunk.parse)(self, strmgr, rest_size)?;
201 if psize > rest_size { return Err(InvalidData); }
202 rest_size -= psize;
203 while rest_size > 0 {
204 let (psize, _) = self.parse_chunk(strmgr, end_tag, rest_size, depth+1)?;
205 if psize > rest_size { return Err(InvalidData); }
206 rest_size -= psize;
207 if ((psize & 1) == 1) && (rest_size > 0) {
208 rest_size -= 1;
209 }
210 }
211
212 return Ok((size + 8, false));
213 }
214 }
215 if !is_list {
216 self.src.read_skip(size)?;
217 } else {
218 if size < 4 { return Err(InvalidData); }
219 self.src.read_skip(size - 4)?;
220 }
221 if (size & 1) == 1 { self.src.read_skip(1)?; }
222 Ok((size + 8, false))
223 }
224
225 fn read_header(&mut self, strmgr: &mut StreamManager, seek_idx: &mut SeekIndex) -> DemuxerResult<()> {
226 let riff_tag = self.src.read_u32be()?;
227 let size = self.src.read_u32le()? as usize;
228 let avi_tag = self.src.read_u32be()?;
229 let mut matches = false;
230 for rt in RIFF_TAGS.iter() {
231 if rt[0] == riff_tag && rt[1] == avi_tag {
232 matches = true;
233 break;
234 }
235 }
236 if !matches {
237 return Err(InvalidData);
238 }
239 self.size = size;
240 let mut rest_size = size;
241 loop {
242 let (csz, end) = self.parse_chunk(strmgr, RIFFTag::List(mktag!(b"LIST"), mktag!(b"movi")), rest_size,0)?;
243 if end {
244 self.movi_size = csz - 4;
245 self.movi_orig = self.movi_size;
246 self.movi_pos = self.src.tell();
247 break;
248 }
249 rest_size -= csz;
250 }
251 if !seek_idx.skip_index {
252 self.src.read_skip(self.movi_size)?;
253 while rest_size > 0 {
254 let ret = self.parse_chunk(strmgr, RIFFTag::Chunk(mktag!(b"idx1")), rest_size,0);
255 if ret.is_err() { break; }
256 let (csz, end) = ret.unwrap();
257 if end {
258 let _res = parse_idx1(&mut self.src, strmgr, seek_idx, csz, self.movi_pos, &mut self.key_offs);
259 break;
260 }
261 rest_size -= csz;
262 }
263 }
264 if self.movi_pos != 0 {
265 self.src.seek(SeekFrom::Start(self.movi_pos))?;
266 } else {
267 return Err(InvalidData);
268 }
269 if !self.sstate.valid_state() || self.sstate.strm_no != self.num_streams {
270 return Err(InvalidData);
271 }
272 Ok(())
273 }
274
275 fn read_extradata(&mut self, size: usize) -> DemuxerResult<Option<Vec<u8>>> {
276 if size == 0 { return Ok(None); }
277 let mut edvec: Vec<u8> = vec![0; size];
278 self.src.read_buf(&mut edvec)?;
279 Ok(Some(edvec))
280 }
281
282 fn parse_palette_change(&mut self, stream_no: usize, size: usize) -> DemuxerResult<()> {
283 for pe in self.pal.iter_mut() {
284 if pe.stream_no == stream_no {
285 let start_clr = self.src.read_byte()? as usize;
286 let len = self.src.read_byte()? as usize;
287 let _flags = self.src.read_u16le()?;
288 validate!(start_clr + len <= 256);
289 validate!(len * 4 + 4 == size);
290 let mut newpal = *pe.pal;
291 for i in start_clr..(start_clr + len) {
292 newpal[i * 4] = self.src.read_byte()?;
293 newpal[i * 4 + 1] = self.src.read_byte()?;
294 newpal[i * 4 + 2] = self.src.read_byte()?;
295 newpal[i * 4 + 3] = self.src.read_byte()?;
296 }
297 pe.pal = Arc::new(newpal);
298 pe.changed = true;
299 return Ok(());
300 }
301 }
302 self.src.read_skip(size)?;
303 Ok(())
304 }
305 }
306
307 const RIFF_TAGS: &[[u32; 2]] = &[
308 [ mktag!(b"RIFF"), mktag!(b"AVI ") ],
309 [ mktag!(b"RIFF"), mktag!(b"AVIX") ],
310 [ mktag!(b"ON2 "), mktag!(b"ON2f") ],
311 ];
312
313 const CHUNKS: [RIFFParser; 7] = [
314 RIFFParser { tag: RIFFTag::List(mktag!(b"LIST"), mktag!(b"hdrl")), parse: parse_hdrl },
315 RIFFParser { tag: RIFFTag::List(mktag!(b"LIST"), mktag!(b"strl")), parse: parse_strl },
316 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"avih")), parse: parse_avih },
317 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"ON2h")), parse: parse_avih },
318 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"strf")), parse: parse_strf },
319 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"strh")), parse: parse_strh },
320 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"JUNK")), parse: parse_junk },
321 ];
322
323 fn is_list_tag(tag: u32) -> bool {
324 for chunk in CHUNKS.iter() {
325 if let RIFFTag::List(ltag, _) = chunk.tag {
326 if tag == ltag {
327 return true;
328 }
329 }
330 }
331 false
332 }
333
334 #[allow(unused_variables)]
335 fn parse_hdrl(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
336 Ok(0)
337 }
338
339 #[allow(unused_variables)]
340 fn parse_strl(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
341 Ok(0)
342 }
343
344 #[allow(unused_variables)]
345 fn parse_strh(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
346 if size < 0x38 { return Err(InvalidData); }
347 let tag = dmx.src.read_u32be()?; //stream type
348 let fcc = dmx.src.read_u32be()?; //handler(fourcc)
349 dmx.src.read_u32le()?; //flags
350 dmx.src.read_skip(2)?; //priority
351 dmx.src.read_skip(2)?; //language
352 dmx.src.read_skip(4)?; //initial frames
353 dmx.tb_num = dmx.src.read_u32le()?; //scale
354 dmx.tb_den = dmx.src.read_u32le()?; //rate
355 dmx.src.read_skip(4)?; //start
356 dmx.src.read_skip(4)?; //length
357 dmx.src.read_skip(4)?; //buf size
358 dmx.src.read_skip(4)?; //quality
359 dmx.src.read_skip(4)?; //sample size
360 let a = dmx.src.read_u16le()?;
361 let b = dmx.src.read_u16le()?;
362 let c = dmx.src.read_u16le()?;
363 let d = dmx.src.read_u16le()?;
364
365 dmx.src.read_skip(size - 0x38)?;
366
367 if !dmx.sstate.valid_state() || dmx.sstate.strm_no >= dmx.num_streams {
368 return Err(InvalidData);
369 }
370 if tag == mktag!(b"vids") {
371 dmx.sstate.strm_type = Some(StreamType::Video);
372 } else if tag == mktag!(b"auds") {
373 dmx.sstate.strm_type = Some(StreamType::Audio);
374 } else {
375 dmx.sstate.strm_type = Some(StreamType::Data);
376 }
377 dmx.sstate.got_strf = false;
378
379 Ok(size)
380 }
381
382 fn parse_strf(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
383 if dmx.sstate.strm_type.is_none() { return Err(InvalidData); }
384 match dmx.sstate.strm_type.unwrap() {
385 StreamType::Video => parse_strf_vids(dmx, strmgr, size),
386 StreamType::Audio => parse_strf_auds(dmx, strmgr, size),
387 _ => parse_strf_xxxx(dmx, strmgr, size),
388 }
389 }
390
391 #[allow(unused_variables)]
392 fn parse_strf_vids(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
393 if size < 40 { return Err(InvalidData); }
394 let bi_size = dmx.src.read_u32le()?;
395 if (bi_size as usize) > size { return Err(InvalidData); }
396 let width = dmx.src.read_u32le()?;
397 let height = dmx.src.read_u32le()? as i32;
398 let planes = dmx.src.read_u16le()?;
399 let bitcount = dmx.src.read_u16le()?;
400 let mut compression: [u8; 4] = [0; 4];
401 dmx.src.read_buf(&mut compression)?;
402 let img_size = dmx.src.read_u32le()?;
403 let xdpi = dmx.src.read_u32le()?;
404 let ydpi = dmx.src.read_u32le()?;
405 let colors = dmx.src.read_u32le()?;
406 validate!(colors <= 256);
407 let imp_colors = dmx.src.read_u32le()?;
408
409 let flip = height < 0;
410 let format = if bitcount > 8 { RGB24_FORMAT } else { PAL8_FORMAT };
411 let vhdr = NAVideoInfo::new(width as usize, if flip { -height as usize } else { height as usize}, flip, PAL8_FORMAT);
412 let vci = NACodecTypeInfo::Video(vhdr);
413 let edata = dmx.read_extradata(size - 40)?;
414 if colors > 0 {
415 if let Some(ref buf) = edata {
416 let mut pal = [0u8; 1024];
417 for (dpal, spal) in pal.chunks_mut(4).take(colors as usize).zip(buf.chunks(4)) {
418 dpal[0] = spal[0];
419 dpal[1] = spal[1];
420 dpal[2] = spal[2];
421 dpal[3] = spal[3];
422 }
423 let pal = PalInfo { pal: Arc::new(pal), changed: true, stream_no: strmgr.get_num_streams() };
424 dmx.pal.push(pal);
425 }
426 }
427 let cname = match register::find_codec_from_avi_fourcc(&compression) {
428 None => "unknown",
429 Some(name) => name,
430 };
431 let vinfo = NACodecInfo::new(cname, vci, edata);
432 let res = strmgr.add_stream(NAStream::new(StreamType::Video, u32::from(dmx.sstate.strm_no), vinfo, dmx.tb_num, dmx.tb_den));
433 if res.is_none() { return Err(MemoryError); }
434 dmx.sstate.reset();
435 Ok(size)
436 }
437
438 #[allow(unused_variables)]
439 fn parse_strf_auds(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
440 if size < 16 { return Err(InvalidData); }
441 let w_format_tag = dmx.src.read_u16le()?;
442 let channels = dmx.src.read_u16le()?;
443 let samplespersec = dmx.src.read_u32le()?;
444 let avgbytespersec = dmx.src.read_u32le()?;
445 let block_align = dmx.src.read_u16le()?;
446 let bits_per_sample = dmx.src.read_u16le()?;
447
448 let soniton = NASoniton::new(bits_per_sample as u8, SONITON_FLAG_SIGNED);
449 let ahdr = NAAudioInfo::new(samplespersec, channels as u8, soniton, block_align as usize);
450 let edata = dmx.read_extradata(size - 16)?;
451 let cname = match register::find_codec_from_wav_twocc(w_format_tag) {
452 None => "unknown",
453 Some(name) => name,
454 };
455 let ainfo = NACodecInfo::new(cname, NACodecTypeInfo::Audio(ahdr), edata);
456 let res = strmgr.add_stream(NAStream::new(StreamType::Audio, u32::from(dmx.sstate.strm_no), ainfo, dmx.tb_num, dmx.tb_den));
457 if res.is_none() { return Err(MemoryError); }
458 dmx.sstate.reset();
459 Ok(size)
460 }
461
462 fn parse_strf_xxxx(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
463 let edata = dmx.read_extradata(size)?;
464 let info = NACodecInfo::new("unknown", NACodecTypeInfo::None, edata);
465 let res = strmgr.add_stream(NAStream::new(StreamType::Data, u32::from(dmx.sstate.strm_no), info, dmx.tb_num, dmx.tb_den));
466 if res.is_none() { return Err(MemoryError); }
467 dmx.sstate.reset();
468 Ok(size)
469 }
470
471 #[allow(unused_variables)]
472 fn parse_avih(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
473 if size < 0x38 { return Err(InvalidData); }
474 let timebase = dmx.src.read_u32le()?; //microsec per frame
475 dmx.src.read_skip(4)?; //max frame size
476 dmx.src.read_skip(4)?; //padding
477 dmx.src.read_u32le()?; //flags
478 let frames = dmx.src.read_u32le()?; //frames
479 dmx.src.read_skip(4)?; //initial frames
480 let streams = dmx.src.read_u32le()?; //streams
481 if streams > 100 { return Err(InvalidData); }
482 dmx.num_streams = streams as u8;
483
484 dmx.src.read_skip(4)?; //buf size
485 let width = dmx.src.read_u32le()?; //width
486 let height = dmx.src.read_u32le()? as i32; //height
487 dmx.src.read_skip(16)?; //reserved
488
489 dmx.cur_frame.resize(streams as usize, 0);
490 dmx.src.read_skip(size - 0x38)?;
491 Ok(size)
492 }
493
494 #[allow(unused_variables)]
495 fn parse_junk(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
496 dmx.src.read_skip(size)?;
497 Ok(size)
498 }
499
500 fn parse_idx1(src: &mut ByteReader, strmgr: &mut StreamManager, seek_idx: &mut SeekIndex, size: usize, movi_pos: u64, key_offs: &mut Vec<u64>) -> DemuxerResult<usize> {
501 validate!((size & 15) == 0);
502 let mut tag = [0u8; 4];
503 let num_entries = size >> 4;
504 let mut counter = [0u64; 100];
505 for _ in 0..num_entries {
506 src.read_buf(&mut tag)?;
507 let flags = src.read_u32le()?;
508 let offset = src.read_u32le()? as u64;
509 let _length = src.read_u32le()?;
510
511 if tag[0] < b'0' || tag[0] > b'9' || tag[1] < b'0' || tag[1] > b'9' {
512 return Err(InvalidData);
513 }
514 let stream_no = ((tag[0] - b'0') * 10 + (tag[1] - b'0')) as usize;
515
516 if (flags & 0x10) != 0 {
517 if let Some(str) = strmgr.get_stream(stream_no) {
518 if str.get_media_type() == StreamType::Video {
519 let (tb_num, tb_den) = str.get_timebase();
520 let pts = counter[stream_no];
521 let time = NATimeInfo::ts_to_time(pts, 1000, tb_num, tb_den);
522 seek_idx.add_entry(stream_no as u32, SeekEntry { time, pts, pos: offset + movi_pos - 4 });
523 }
524 key_offs.push(offset);
525 }
526 }
527 counter[stream_no] += 1;
528 }
529 key_offs.sort_unstable();
530 Ok(size)
531 }
532
533 pub struct AVIDemuxerCreator { }
534
535 impl DemuxerCreator for AVIDemuxerCreator {
536 fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> {
537 Box::new(AVIDemuxer::new(br))
538 }
539 fn get_name(&self) -> &'static str { "avi" }
540 }
541
542 #[cfg(test)]
543 mod test {
544 use super::*;
545 use std::fs::File;
546
547 #[test]
548 fn test_avi_demux() {
549 let mut file = File::open("assets/Indeo/laser05.avi").unwrap();
550 let mut fr = FileReader::new_read(&mut file);
551 let mut br = ByteReader::new(&mut fr);
552 let mut dmx = AVIDemuxer::new(&mut br);
553 let mut sm = StreamManager::new();
554 let mut si = SeekIndex::new();
555 dmx.open(&mut sm, &mut si).unwrap();
556
557 loop {
558 let pktres = dmx.get_frame(&mut sm);
559 if let Err(e) = pktres {
560 if e == DemuxerError::EOF { break; }
561 panic!("error");
562 }
563 let pkt = pktres.unwrap();
564 println!("Got {}", pkt);
565 }
566 }
567 }