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