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