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