]> git.nihav.org Git - nihav.git/blob - src/demuxers/avi.rs
more work on supporting decoders in framework
[nihav.git] / 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 }