src/demuxers/avi.rs

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