nihav-commonfmt/src/demuxers/avi.rs

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