nihav-commonfmt/src/demuxers/avi.rs

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