avi: fix palette order
[nihav.git] / nihav-commonfmt / src / demuxers / avi.rs
1 use nihav_core::demuxers::*;
2 use nihav_registry::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 struct PalInfo {
38 pal: Arc<[u8; 1024]>,
39 changed: bool,
40 stream_no: usize,
41 }
42
43 #[allow(dead_code)]
44 struct AVIDemuxer<'a> {
45 src: &'a mut ByteReader<'a>,
46 cur_frame: Vec<u64>,
47 num_streams: u8,
48 size: usize,
49 movi_size: usize,
50 movi_pos: u64,
51 movi_orig: usize,
52 sstate: StreamState,
53 tb_num: u32,
54 tb_den: u32,
55 key_offs: Vec<u64>,
56 pal: Vec<PalInfo>,
57 }
58
59 #[derive(Debug,Clone,Copy,PartialEq)]
60 enum RIFFTag {
61 Chunk(u32),
62 List(u32,u32),
63 }
64
65 struct RIFFParser {
66 tag: RIFFTag,
67 parse: fn(&mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize>,
68 }
69
70 impl<'a> DemuxCore<'a> for AVIDemuxer<'a> {
71 fn open(&mut self, strmgr: &mut StreamManager, seek_index: &mut SeekIndex) -> DemuxerResult<()> {
72 self.read_header(strmgr, seek_index)?;
73 Ok(())
74 }
75
76 fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> {
77 if self.movi_size == 0 { return Err(EOF); }
78 let mut tag: [u8; 4] = [0; 4];
79 loop {
80 if (self.src.tell() & 1) == 1 {
81 self.src.read_skip(1)?;
82 self.movi_size -= 1;
83 if self.movi_size == 0 { return Err(EOF); }
84 }
85 let is_keyframe = self.key_offs.binary_search(&self.src.tell()).is_ok();
86 self.src.read_buf(&mut tag)?;
87 let size = self.src.read_u32le()? as usize;
88 if mktag!(tag) == mktag!(b"JUNK") {
89 self.movi_size -= size + 8;
90 self.src.read_skip(size)?;
91 if self.movi_size == 0 { return Err(EOF); }
92 continue;
93 }
94 if mktag!(tag) == mktag!(b"LIST") {
95 self.movi_size -= 12;
96 self.src.read_skip(4)?;
97 if self.movi_size == 0 { return Err(EOF); }
98 continue;
99 }
100 if (tag[0] == b'i' && tag[1] == b'x') || (&tag == b"idx1") {
101 return Err(EOF);
102 }
103 if tag[0] < b'0' || tag[0] > b'9' || tag[1] < b'0' || tag[1] > b'9' {
104 return Err(InvalidData);
105 }
106 let stream_no = (tag[0] - b'0') * 10 + (tag[1] - b'0');
107 if tag[2] == b'p' && tag[3] == b'c' {
108 self.parse_palette_change(stream_no as usize, size)?;
109 self.movi_size -= size;
110 if self.movi_size == 0 { return Err(EOF); }
111 continue;
112 }
113 let str = strmgr.get_stream(stream_no as usize);
114 if str.is_none() { return Err(InvalidData); }
115 let stream = str.unwrap();
116 if size == 0 {
117 self.movi_size -= 8;
118 if self.movi_size == 0 { return Err(EOF); }
119 continue;
120 }
121 let (tb_num, tb_den) = stream.get_timebase();
122 let ts = NATimeInfo::new(Some(self.cur_frame[stream_no as usize]), None, None, tb_num, tb_den);
123 let mut pkt = self.src.read_packet(stream, ts, is_keyframe, size)?;
124 for pe in self.pal.iter_mut() {
125 if pe.stream_no == (stream_no as usize) {
126 pkt.add_side_data(NASideData::Palette(pe.changed, pe.pal.clone()));
127 pe.changed = false;
128 break;
129 }
130 }
131 self.cur_frame[stream_no as usize] += 1;
132 self.movi_size -= size + 8;
133
134 return Ok(pkt);
135 }
136 }
137
138 fn seek(&mut self, time: u64, seek_index: &SeekIndex) -> DemuxerResult<()> {
139 let ret = seek_index.find_pos(time);
140 if ret.is_none() {
141 return Err(DemuxerError::SeekError);
142 }
143 let seek_info = ret.unwrap();
144
145 if seek_info.pos < self.movi_pos { return Err(DemuxerError::SeekError); }
146 let skip_size = (seek_info.pos - self.movi_pos) as usize;
147 if skip_size > self.movi_size { return Err(DemuxerError::SeekError); }
148 self.movi_size = self.movi_orig - skip_size;
149
150 self.cur_frame[seek_info.str_id as usize] = seek_info.pts;
151 self.src.seek(SeekFrom::Start(seek_info.pos))?;
152
153 Ok(())
154 }
155 }
156
157 impl<'a> AVIDemuxer<'a> {
158 fn new(io: &'a mut ByteReader<'a>) -> Self {
159 AVIDemuxer {
160 cur_frame: Vec::new(),
161 num_streams: 0,
162 src: io,
163 size: 0,
164 movi_size: 0,
165 movi_pos: 0,
166 movi_orig: 0,
167 sstate: StreamState::new(),
168 tb_num: 0,
169 tb_den: 0,
170 key_offs: Vec::new(),
171 pal: Vec::new(),
172 }
173 }
174
175 fn parse_chunk(&mut self, strmgr: &mut StreamManager, end_tag: RIFFTag, csize: usize, depth: u16) -> DemuxerResult<(usize, bool)> {
176 if csize < 8 { return Err(InvalidData); }
177 if depth > 42 { return Err(InvalidData); }
178
179 let tag = self.src.read_u32be()?;
180 let size = self.src.read_u32le()? as usize;
181 if size > csize { return Err(InvalidData); }
182 if RIFFTag::Chunk(tag) == end_tag {
183 return Ok((size, true));
184 }
185 let is_list = is_list_tag(tag);
186 let ltag = if is_list { self.src.read_u32be()? } else { 0 };
187 if RIFFTag::List(tag, ltag) == end_tag {
188 return Ok((size, true));
189 }
190
191 for chunk in CHUNKS.iter() {
192 if RIFFTag::Chunk(tag) == chunk.tag {
193 let psize = (chunk.parse)(self, strmgr, size)?;
194 if psize != size { return Err(InvalidData); }
195 if (psize & 1) == 1 { self.src.read_skip(1)?; }
196 return Ok((size + 8, false));
197 }
198 if RIFFTag::List(tag, ltag) == chunk.tag {
199 let mut rest_size = size - 4;
200 let psize = (chunk.parse)(self, strmgr, rest_size)?;
201 if psize > rest_size { return Err(InvalidData); }
202 rest_size -= psize;
203 while rest_size > 0 {
204 let (psize, _) = self.parse_chunk(strmgr, end_tag, rest_size, depth+1)?;
205 if psize > rest_size { return Err(InvalidData); }
206 rest_size -= psize;
207 if ((psize & 1) == 1) && (rest_size > 0) {
208 rest_size -= 1;
209 }
210 }
211
212 return Ok((size + 8, false));
213 }
214 }
215 if !is_list {
216 self.src.read_skip(size)?;
217 } else {
218 if size < 4 { return Err(InvalidData); }
219 self.src.read_skip(size - 4)?;
220 }
221 if (size & 1) == 1 { self.src.read_skip(1)?; }
222 Ok((size + 8, false))
223 }
224
225 fn read_header(&mut self, strmgr: &mut StreamManager, seek_idx: &mut SeekIndex) -> DemuxerResult<()> {
226 let riff_tag = self.src.read_u32be()?;
227 let size = self.src.read_u32le()? as usize;
228 let avi_tag = self.src.read_u32be()?;
229 let mut matches = false;
230 for rt in RIFF_TAGS.iter() {
231 if rt[0] == riff_tag && rt[1] == avi_tag {
232 matches = true;
233 break;
234 }
235 }
236 if !matches {
237 return Err(InvalidData);
238 }
239 self.size = size;
240 let mut rest_size = size;
241 loop {
242 let (csz, end) = self.parse_chunk(strmgr, RIFFTag::List(mktag!(b"LIST"), mktag!(b"movi")), rest_size,0)?;
243 if end {
244 self.movi_size = csz - 4;
245 self.movi_orig = self.movi_size;
246 self.movi_pos = self.src.tell();
247 break;
248 }
249 rest_size -= csz;
250 }
251 if !seek_idx.skip_index {
252 self.src.read_skip(self.movi_size)?;
253 while rest_size > 0 {
254 let ret = self.parse_chunk(strmgr, RIFFTag::Chunk(mktag!(b"idx1")), rest_size,0);
255 if ret.is_err() { break; }
256 let (csz, end) = ret.unwrap();
257 if end {
258 let _res = parse_idx1(&mut self.src, strmgr, seek_idx, csz, self.movi_pos, &mut self.key_offs);
259 break;
260 }
261 rest_size -= csz;
262 }
263 }
264 if self.movi_pos != 0 {
265 self.src.seek(SeekFrom::Start(self.movi_pos))?;
266 } else {
267 return Err(InvalidData);
268 }
269 if !self.sstate.valid_state() || self.sstate.strm_no != self.num_streams {
270 return Err(InvalidData);
271 }
272 Ok(())
273 }
274
275 fn read_extradata(&mut self, size: usize) -> DemuxerResult<Option<Vec<u8>>> {
276 if size == 0 { return Ok(None); }
277 let mut edvec: Vec<u8> = vec![0; size];
278 self.src.read_buf(&mut edvec)?;
279 Ok(Some(edvec))
280 }
281
282 fn parse_palette_change(&mut self, stream_no: usize, size: usize) -> DemuxerResult<()> {
283 for pe in self.pal.iter_mut() {
284 if pe.stream_no == stream_no {
285 let start_clr = self.src.read_byte()? as usize;
286 let len = self.src.read_byte()? as usize;
287 let _flags = self.src.read_u16le()?;
288 validate!(start_clr + len <= 256);
289 validate!(len * 4 + 4 == size);
290 let mut newpal = *pe.pal;
291 for i in start_clr..(start_clr + len) {
292 newpal[i * 4] = self.src.read_byte()?;
293 newpal[i * 4 + 1] = self.src.read_byte()?;
294 newpal[i * 4 + 2] = self.src.read_byte()?;
295 newpal[i * 4 + 3] = 0;
296 self.src.read_byte()?; // flags
297 }
298 pe.pal = Arc::new(newpal);
299 pe.changed = true;
300 return Ok(());
301 }
302 }
303 self.src.read_skip(size)?;
304 Ok(())
305 }
306 }
307
308 const RIFF_TAGS: &[[u32; 2]] = &[
309 [ mktag!(b"RIFF"), mktag!(b"AVI ") ],
310 [ mktag!(b"RIFF"), mktag!(b"AVIX") ],
311 [ mktag!(b"ON2 "), mktag!(b"ON2f") ],
312 ];
313
314 const CHUNKS: [RIFFParser; 7] = [
315 RIFFParser { tag: RIFFTag::List(mktag!(b"LIST"), mktag!(b"hdrl")), parse: parse_hdrl },
316 RIFFParser { tag: RIFFTag::List(mktag!(b"LIST"), mktag!(b"strl")), parse: parse_strl },
317 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"avih")), parse: parse_avih },
318 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"ON2h")), parse: parse_avih },
319 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"strf")), parse: parse_strf },
320 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"strh")), parse: parse_strh },
321 RIFFParser { tag: RIFFTag::Chunk(mktag!(b"JUNK")), parse: parse_junk },
322 ];
323
324 fn is_list_tag(tag: u32) -> bool {
325 for chunk in CHUNKS.iter() {
326 if let RIFFTag::List(ltag, _) = chunk.tag {
327 if tag == ltag {
328 return true;
329 }
330 }
331 }
332 false
333 }
334
335 #[allow(unused_variables)]
336 fn parse_hdrl(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
337 Ok(0)
338 }
339
340 #[allow(unused_variables)]
341 fn parse_strl(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
342 Ok(0)
343 }
344
345 #[allow(unused_variables)]
346 fn parse_strh(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
347 if size < 0x38 { return Err(InvalidData); }
348 let tag = dmx.src.read_u32be()?; //stream type
349 let fcc = dmx.src.read_u32be()?; //handler(fourcc)
350 dmx.src.read_u32le()?; //flags
351 dmx.src.read_skip(2)?; //priority
352 dmx.src.read_skip(2)?; //language
353 dmx.src.read_skip(4)?; //initial frames
354 dmx.tb_num = dmx.src.read_u32le()?; //scale
355 dmx.tb_den = dmx.src.read_u32le()?; //rate
356 dmx.src.read_skip(4)?; //start
357 dmx.src.read_skip(4)?; //length
358 dmx.src.read_skip(4)?; //buf size
359 dmx.src.read_skip(4)?; //quality
360 dmx.src.read_skip(4)?; //sample size
361 let a = dmx.src.read_u16le()?;
362 let b = dmx.src.read_u16le()?;
363 let c = dmx.src.read_u16le()?;
364 let d = dmx.src.read_u16le()?;
365
366 dmx.src.read_skip(size - 0x38)?;
367
368 if !dmx.sstate.valid_state() || dmx.sstate.strm_no >= dmx.num_streams {
369 return Err(InvalidData);
370 }
371 if tag == mktag!(b"vids") {
372 dmx.sstate.strm_type = Some(StreamType::Video);
373 } else if tag == mktag!(b"auds") {
374 dmx.sstate.strm_type = Some(StreamType::Audio);
375 } else {
376 dmx.sstate.strm_type = Some(StreamType::Data);
377 }
378 dmx.sstate.got_strf = false;
379
380 Ok(size)
381 }
382
383 fn parse_strf(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
384 if dmx.sstate.strm_type.is_none() { return Err(InvalidData); }
385 match dmx.sstate.strm_type.unwrap() {
386 StreamType::Video => parse_strf_vids(dmx, strmgr, size),
387 StreamType::Audio => parse_strf_auds(dmx, strmgr, size),
388 _ => parse_strf_xxxx(dmx, strmgr, size),
389 }
390 }
391
392 #[allow(unused_variables)]
393 fn parse_strf_vids(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
394 if size < 40 { return Err(InvalidData); }
395 let bi_size = dmx.src.read_u32le()?;
396 if (bi_size as usize) > size { return Err(InvalidData); }
397 let width = dmx.src.read_u32le()?;
398 let height = dmx.src.read_u32le()? as i32;
399 let planes = dmx.src.read_u16le()?;
400 let bitcount = dmx.src.read_u16le()?;
401 let mut compression: [u8; 4] = [0; 4];
402 dmx.src.read_buf(&mut compression)?;
403 let img_size = dmx.src.read_u32le()?;
404 let xdpi = dmx.src.read_u32le()?;
405 let ydpi = dmx.src.read_u32le()?;
406 let colors = dmx.src.read_u32le()?;
407 validate!(colors <= 256);
408 let imp_colors = dmx.src.read_u32le()?;
409
410 let flip = height < 0;
411 let format = if bitcount > 8 { RGB24_FORMAT } else { PAL8_FORMAT };
412 let vhdr = NAVideoInfo::new(width as usize, if flip { -height as usize } else { height as usize}, flip, PAL8_FORMAT);
413 let vci = NACodecTypeInfo::Video(vhdr);
414 let edata = dmx.read_extradata(size - 40)?;
415 if colors > 0 {
416 if let Some(ref buf) = edata {
417 let mut pal = [0u8; 1024];
418 for (dpal, spal) in pal.chunks_mut(4).take(colors as usize).zip(buf.chunks(4)) {
419 dpal[0] = spal[2];
420 dpal[1] = spal[1];
421 dpal[2] = spal[0];
422 dpal[3] = 0;
423 }
424 let pal = PalInfo { pal: Arc::new(pal), changed: true, stream_no: strmgr.get_num_streams() };
425 dmx.pal.push(pal);
426 }
427 }
428 let cname = match register::find_codec_from_avi_fourcc(&compression) {
429 None => "unknown",
430 Some(name) => name,
431 };
432 let vinfo = NACodecInfo::new(cname, vci, edata);
433 let res = strmgr.add_stream(NAStream::new(StreamType::Video, u32::from(dmx.sstate.strm_no), vinfo, dmx.tb_num, dmx.tb_den));
434 if res.is_none() { return Err(MemoryError); }
435 dmx.sstate.reset();
436 Ok(size)
437 }
438
439 #[allow(unused_variables)]
440 fn parse_strf_auds(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
441 if size < 16 { return Err(InvalidData); }
442 let w_format_tag = dmx.src.read_u16le()?;
443 let channels = dmx.src.read_u16le()?;
444 let samplespersec = dmx.src.read_u32le()?;
445 let avgbytespersec = dmx.src.read_u32le()?;
446 let block_align = dmx.src.read_u16le()?;
447 let bits_per_sample = dmx.src.read_u16le()?;
448
449 let soniton = NASoniton::new(bits_per_sample as u8, SONITON_FLAG_SIGNED);
450 let ahdr = NAAudioInfo::new(samplespersec, channels as u8, soniton, block_align as usize);
451 let edata = dmx.read_extradata(size - 16)?;
452 let cname = match register::find_codec_from_wav_twocc(w_format_tag) {
453 None => "unknown",
454 Some(name) => name,
455 };
456 let ainfo = NACodecInfo::new(cname, NACodecTypeInfo::Audio(ahdr), edata);
457 let res = strmgr.add_stream(NAStream::new(StreamType::Audio, u32::from(dmx.sstate.strm_no), ainfo, dmx.tb_num, dmx.tb_den));
458 if res.is_none() { return Err(MemoryError); }
459 dmx.sstate.reset();
460 Ok(size)
461 }
462
463 fn parse_strf_xxxx(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
464 let edata = dmx.read_extradata(size)?;
465 let info = NACodecInfo::new("unknown", NACodecTypeInfo::None, edata);
466 let res = strmgr.add_stream(NAStream::new(StreamType::Data, u32::from(dmx.sstate.strm_no), info, dmx.tb_num, dmx.tb_den));
467 if res.is_none() { return Err(MemoryError); }
468 dmx.sstate.reset();
469 Ok(size)
470 }
471
472 #[allow(unused_variables)]
473 fn parse_avih(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
474 if size < 0x38 { return Err(InvalidData); }
475 let timebase = dmx.src.read_u32le()?; //microsec per frame
476 dmx.src.read_skip(4)?; //max frame size
477 dmx.src.read_skip(4)?; //padding
478 dmx.src.read_u32le()?; //flags
479 let frames = dmx.src.read_u32le()?; //frames
480 dmx.src.read_skip(4)?; //initial frames
481 let streams = dmx.src.read_u32le()?; //streams
482 if streams > 100 { return Err(InvalidData); }
483 dmx.num_streams = streams as u8;
484
485 dmx.src.read_skip(4)?; //buf size
486 let width = dmx.src.read_u32le()?; //width
487 let height = dmx.src.read_u32le()? as i32; //height
488 dmx.src.read_skip(16)?; //reserved
489
490 dmx.cur_frame.resize(streams as usize, 0);
491 dmx.src.read_skip(size - 0x38)?;
492 Ok(size)
493 }
494
495 #[allow(unused_variables)]
496 fn parse_junk(dmx: &mut AVIDemuxer, strmgr: &mut StreamManager, size: usize) -> DemuxerResult<usize> {
497 dmx.src.read_skip(size)?;
498 Ok(size)
499 }
500
501 fn parse_idx1(src: &mut ByteReader, strmgr: &mut StreamManager, seek_idx: &mut SeekIndex, size: usize, movi_pos: u64, key_offs: &mut Vec<u64>) -> DemuxerResult<usize> {
502 validate!((size & 15) == 0);
503 let mut tag = [0u8; 4];
504 let num_entries = size >> 4;
505 let mut counter = [0u64; 100];
506 for _ in 0..num_entries {
507 src.read_buf(&mut tag)?;
508 let flags = src.read_u32le()?;
509 let offset = src.read_u32le()? as u64;
510 let _length = src.read_u32le()?;
511
512 if tag[0] < b'0' || tag[0] > b'9' || tag[1] < b'0' || tag[1] > b'9' {
513 return Err(InvalidData);
514 }
515 let stream_no = ((tag[0] - b'0') * 10 + (tag[1] - b'0')) as usize;
516
517 if (flags & 0x10) != 0 {
518 if let Some(str) = strmgr.get_stream(stream_no) {
519 if str.get_media_type() == StreamType::Video {
520 let (tb_num, tb_den) = str.get_timebase();
521 let pts = counter[stream_no];
522 let time = NATimeInfo::ts_to_time(pts, 1000, tb_num, tb_den);
523 seek_idx.add_entry(stream_no as u32, SeekEntry { time, pts, pos: offset + movi_pos - 4 });
524 }
525 key_offs.push(offset);
526 }
527 }
528 counter[stream_no] += 1;
529 }
530 key_offs.sort_unstable();
531 Ok(size)
532 }
533
534 pub struct AVIDemuxerCreator { }
535
536 impl DemuxerCreator for AVIDemuxerCreator {
537 fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> {
538 Box::new(AVIDemuxer::new(br))
539 }
540 fn get_name(&self) -> &'static str { "avi" }
541 }
542
543 #[cfg(test)]
544 mod test {
545 use super::*;
546 use std::fs::File;
547
548 #[test]
549 fn test_avi_demux() {
550 let mut file = File::open("assets/Indeo/laser05.avi").unwrap();
551 let mut fr = FileReader::new_read(&mut file);
552 let mut br = ByteReader::new(&mut fr);
553 let mut dmx = AVIDemuxer::new(&mut br);
554 let mut sm = StreamManager::new();
555 let mut si = SeekIndex::new();
556 dmx.open(&mut sm, &mut si).unwrap();
557
558 loop {
559 let pktres = dmx.get_frame(&mut sm);
560 if let Err(e) = pktres {
561 if e == DemuxerError::EOF { break; }
562 panic!("error");
563 }
564 let pkt = pktres.unwrap();
565 println!("Got {}", pkt);
566 }
567 }
568 }