nihav-flash/src/demuxers/flv.rs

   1 use nihav_core::demuxers::*;
   2 use nihav_core::io::bitreader::*;
   3
   4 const AVC_ID: u8 = 7;
   5
   6 struct FLVDemuxer<'a> {
   7     src:        &'a mut ByteReader<'a>,
   8     vpkts:      Vec<NAPacket>,
   9     vtag:       Option<u8>,
  10     apkts:      Vec<NAPacket>,
  11     atag:       Option<u8>,
  12     vstream:    usize,
  13     astream:    usize,
  14     duration:   u64,
  15     width:      usize,
  16     height:     usize,
  17 }
  18
  19 fn get_vcodec_name(tag: u8) -> DemuxerResult<&'static str> {
  20     match tag {
  21         2 => Ok("flv263"),
  22         3 => Ok("flashsv"),
  23         4 => Ok("vp6f"),
  24         5 => Ok("vp6a"),
  25         6 => Ok("flashsv2"),
  26         7 => Ok("h264"),
  27         _ => Err(DemuxerError::InvalidData),
  28     }
  29 }
  30
  31 impl<'a> DemuxCore<'a> for FLVDemuxer<'a> {
  32     fn open(&mut self, strmgr: &mut StreamManager, seek_index: &mut SeekIndex) -> DemuxerResult<()> {
  33         let mut tag = [0; 3];
  34                                           self.src.read_buf(&mut tag)?;
  35         validate!(&tag == b"FLV");
  36         let ver                         = self.src.read_byte()?;
  37         validate!(ver == 0 || ver == 1);
  38         let hdr                         = self.src.read_byte()?;
  39         validate!((hdr & 0xF2) == 0);
  40         let has_audio = (hdr & 4) != 0;
  41         let has_video = (hdr & 1) != 0;
  42         validate!(has_video || has_audio);
  43         let hdr_size                    = self.src.read_u32be()?;
  44         validate!(hdr_size >= 9);
  45
  46         let first_prev_tag              = self.src.peek_u32be()?;
  47         validate!(first_prev_tag == 0);
  48
  49         while (self.vtag.is_some() != has_video) || (self.atag.is_some() != has_audio) {
  50             self.parse_tag(strmgr)?;
  51             if self.apkts.len() > 100 || self.vpkts.len() > 100 {
  52                 return Err(DemuxerError::InvalidData);
  53             }
  54         }
  55
  56         seek_index.mode = SeekIndexMode::Automatic;
  57
  58         Ok(())
  59     }
  60
  61     fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> {
  62         loop {
  63             if !self.vpkts.is_empty() && self.vpkts.len() >= self.apkts.len() {
  64                 return Ok(self.vpkts.remove(0));
  65             }
  66             if !self.apkts.is_empty() {
  67                 return Ok(self.apkts.remove(0));
  68             }
  69             self.parse_tag(strmgr)?;
  70         }
  71     }
  72     fn seek(&mut self, time: NATimePoint, _seek_index: &SeekIndex) -> DemuxerResult<()> {
  73         let dst_ms = match time {
  74                 NATimePoint::PTS(pts) => pts,
  75                 NATimePoint::Milliseconds(ms) => ms,
  76                 NATimePoint::None => return Err(DemuxerError::SeekError),
  77             };
  78         self.apkts.clear();
  79         self.vpkts.clear();
  80         let mut prev = None;
  81         loop {
  82             let ppos                    = self.src.read_u32be()?;
  83             let ret                     = self.src.read_byte();
  84             if let Err(ByteIOError::EOF) = ret {
  85                                           self.src.seek(SeekFrom::Current(-8 - i64::from(ppos)))?;
  86                 continue;
  87             }
  88             let data_size               = self.src.read_u24be()?;
  89             let time                    = self.src.read_u24be()?;
  90             let ext_time                = self.src.read_byte()?;
  91             let _stream_id              = self.src.read_u24be()?;
  92             let ts = (u64::from(ext_time) << 32) | u64::from(time);
  93             if dst_ms == ts {
  94                                           self.src.seek(SeekFrom::Current(-15))?;
  95                 return Ok(());
  96             }
  97             if let Some(p_ts) = prev {
  98                 if dst_ms > p_ts && dst_ms < ts {
  99                                           self.src.seek(SeekFrom::Current(-19 - i64::from(ppos)))?;
 100                     return Ok(());
 101                 }
 102             }
 103             prev = Some(ts);
 104             if dst_ms < ts {
 105                                           self.src.seek(SeekFrom::Current(-19 - i64::from(ppos)))?;
 106             } else {
 107                                           self.src.seek(SeekFrom::Current(i64::from(data_size)))?;
 108             }
 109         }
 110     }
 111     fn get_duration(&self) -> u64 { self.duration }
 112 }
 113
 114 impl<'a> NAOptionHandler for FLVDemuxer<'a> {
 115     fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
 116     fn set_options(&mut self, _options: &[NAOption]) { }
 117     fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
 118 }
 119
 120 impl<'a> FLVDemuxer<'a> {
 121     fn new(io: &'a mut ByteReader<'a>) -> Self {
 122         Self {
 123             src:        io,
 124             vpkts:      Vec::with_capacity(2),
 125             apkts:      Vec::with_capacity(2),
 126             vtag:       None,
 127             atag:       None,
 128             vstream:    0,
 129             astream:    0,
 130             duration:   0,
 131             width:      0,
 132             height:     0,
 133         }
 134     }
 135     fn parse_tag(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<()> {
 136         let _prev_tag_size              = self.src.read_u32be()?;
 137         let ret                         = self.src.read_byte();
 138         if let Err(ByteIOError::EOF) = ret {
 139             return Err(DemuxerError::EOF);
 140         }
 141
 142         let tag = ret?;
 143         let mut data_size               = self.src.read_u24be()? as usize;
 144         let time                        = self.src.read_u24be()?;
 145         let ext_time                    = self.src.read_byte()?;
 146         let stream_id                   = self.src.read_u24be()?;
 147         validate!(stream_id == 0);
 148         if data_size == 0 {
 149             return Ok(());
 150         }
 151         let pkt_start = self.src.tell();
 152         match tag {
 153             8 => {
 154                 let hdr                 = self.src.read_byte()?;
 155                 if let Some(tag) = self.atag {
 156                     validate!(tag == (hdr >> 4));
 157                 } else if data_size > 0 {
 158                     let cname = match hdr >> 4 {
 159                             0 | 3 => "pcm",
 160                             1 => "flv-adpcm",
 161                             2 | 14 => "mp3",
 162                             4..=6 => "asao",
 163                             7 => "alaw",
 164                             8 => "ulaw",
 165                             10 => "aac",
 166                             11 => "speex",
 167                             _ => return Err(DemuxerError::InvalidData),
 168                         };
 169                     let mut srate = match (hdr >> 2) & 0x3 {
 170                             0 => 5500,
 171                             1 => 11025,
 172                             2 => 22050,
 173                             _ => 44100,
 174                         };
 175                     let bits = if (hdr & 2) == 0 { 8 } else { 16 };
 176                     let mut channels = if (hdr & 1) == 0 { 1 } else { 2 };
 177                     let mut aac_edata = false;
 178                     match hdr >> 4 {
 179                          4 => { srate = 16000; channels = 1; },
 180                          5 => { srate = 8000; channels = 1; },
 181                         10 => { aac_edata = self.src.read_byte()? == 0; },
 182                         14 => srate = 8000,
 183                          _ => {},
 184                     };
 185                     let edata = if aac_edata {
 186                             let pkt_hdr_size = (self.src.tell() - pkt_start) as usize;
 187                             validate!(data_size >= pkt_hdr_size);
 188                             let mut data = vec![0; data_size - pkt_hdr_size];
 189                                               self.src.read_buf(&mut data)?;
 190                             Some(data)
 191                         } else {
 192                             None
 193                         };
 194                     let soniton = if bits == 16 { SND_S16P_FORMAT } else { SND_U8_FORMAT };
 195                     let ahdr = NAAudioInfo::new(srate, channels, soniton, 0);
 196                     let ci = NACodecTypeInfo::Audio(ahdr);
 197                     let ainfo = NACodecInfo::new(cname, ci, edata);
 198                     if let Some(id) = strmgr.add_stream(NAStream::new(StreamType::Audio, 1, ainfo, 1, 1000, 0)) {
 199                         self.astream = id;
 200                     } else {
 201                         return Err(DemuxerError::MemoryError);
 202                     }
 203                     self.atag = Some(hdr >> 4);
 204
 205                     if aac_edata {
 206                         return Ok(());
 207                     }
 208                 }
 209                 if (hdr >> 4) == 10 {
 210                     let pkt_type            = self.src.read_byte()?;
 211                     validate!(pkt_type == 1);
 212                 }
 213                 let pkt_hdr_size = (self.src.tell() - pkt_start) as usize;
 214                 validate!(data_size >= pkt_hdr_size);
 215                 data_size -= pkt_hdr_size;
 216                 if data_size > 0 {
 217                     let stream = strmgr.get_stream(self.astream).unwrap();
 218                     let pts = (u64::from(ext_time) << 24) | u64::from(time);
 219                     let ts = stream.make_ts(Some(pts), None, None);
 220                     self.apkts.push(self.src.read_packet(stream, ts, true, data_size)?);
 221                 }
 222             },
 223             9 => {
 224                 let hdr                 = self.src.read_byte()?;
 225                 let ftype = match hdr >> 4 {
 226                         1 => FrameType::I,
 227                         2 => FrameType::P,
 228                         3 => FrameType::P, // droppable
 229                         4 => FrameType::Other, // generated key frame
 230                         5 => FrameType::Other, // video info/command frame
 231                         _ => return Err(DemuxerError::InvalidData),
 232                     };
 233                 let codec_tag = hdr & 0xF;
 234                 if let Some(id) = self.vtag {
 235                     validate!(id == codec_tag);
 236                 } else {
 237                     let cname = get_vcodec_name(codec_tag)?;
 238                     let is_avc = codec_tag == AVC_ID;
 239                     if is_avc {
 240                         let pkt_type            = self.src.read_byte()?;
 241                         validate!(pkt_type == 0);
 242                                                   self.src.read_u24be()?;
 243                     }
 244                     let mut edata = None;
 245                     let (width, height) = match codec_tag {
 246                             2 => {
 247                                 let mut buf = [0; 9];
 248                                                   self.src.peek_buf(&mut buf)?;
 249                                 let mut br = BitReader::new(&buf, BitReaderMode::BE);
 250                                                 br.skip(30).unwrap_or(());
 251                                 let sfmt      = br.read(3).unwrap_or(7);
 252                                 match sfmt {
 253                                     0 => {
 254                                         let w           = br.read(8).unwrap_or(0) as usize;
 255                                         let h           = br.read(8).unwrap_or(0) as usize;
 256                                         (w, h)
 257                                     },
 258                                     1 => {
 259                                         let w           = br.read(16).unwrap_or(0) as usize;
 260                                         let h           = br.read(16).unwrap_or(0) as usize;
 261                                         (w, h)
 262                                     },
 263                                     2 => (352, 288),
 264                                     3 => (176, 144),
 265                                     4 => (128, 96),
 266                                     5 => (320, 240),
 267                                     6 => (160, 120),
 268                                     _ => (0, 0),
 269                                 }
 270                             },
 271                             3 | 6 => {
 272                                 let mut buf = [0; 4];
 273                                                   self.src.peek_buf(&mut buf)?;
 274                                 let w = (read_u16be(&buf[0..])? & 0xFFF) as usize;
 275                                 let h = (read_u16be(&buf[2..])? & 0xFFF) as usize;
 276                                 (w, h)
 277                             },
 278                             4 => {
 279                                 let mut buf = [0; 7];
 280                                                   self.src.peek_buf(&mut buf)?;
 281                                 let off = if (buf[1] & 1) != 0 || (buf[2] & 6) == 0 { 5 } else { 3 };
 282                                 validate!(buf[off] != 0 && buf[off + 1] != 0);
 283                                 let w = usize::from(buf[off + 1]) * 16 - usize::from(buf[0] >> 4);
 284                                 let h = usize::from(buf[off])     * 16 - usize::from(buf[0] & 0xF);
 285
 286                                 edata = Some(vec![buf[0]]);
 287
 288                                 (w, h)
 289                             },
 290                             5 => {
 291                                 let mut buf = [0; 10];
 292                                                   self.src.peek_buf(&mut buf)?;
 293                                 let off = if (buf[4] & 1) != 0 || (buf[5] & 6) == 0 { 8 } else { 6 };
 294                                 validate!(buf[off] != 0 && buf[off + 1] != 0);
 295                                 let w = usize::from(buf[off + 1]) * 16 - usize::from(buf[0] >> 4);
 296                                 let h = usize::from(buf[off])     * 16 - usize::from(buf[0] & 0xF);
 297
 298                                 edata = Some(vec![buf[0]]);
 299
 300                                 (w, h)
 301                             },
 302                             7 => {
 303                                 let pkt_hdr_size = (self.src.tell() - pkt_start) as usize;
 304                                 validate!(data_size >= pkt_hdr_size);
 305                                 data_size -= pkt_hdr_size;
 306                                 let mut data = vec![0; data_size + 4];
 307                                 data[..4].copy_from_slice(b"avcC");
 308                                                   self.src.read_buf(&mut data[4..])?;
 309                                 edata = Some(data);
 310                                 (self.width, self.height)
 311                             },
 312                             _ => unreachable!(),
 313                         };
 314
 315                     let vhdr = NAVideoInfo::new(width, height, false, YUV420_FORMAT);
 316                     let vci = NACodecTypeInfo::Video(vhdr);
 317                     let vinfo = NACodecInfo::new(cname, vci, edata);
 318                     if let Some(id) = strmgr.add_stream(NAStream::new(StreamType::Video, 0, vinfo, 1, 1000, 0)) {
 319                         self.vstream = id;
 320                     } else {
 321                         return Err(DemuxerError::MemoryError);
 322                     }
 323                     self.vtag = Some(codec_tag);
 324                     if is_avc {
 325                         return Ok(());
 326                     }
 327                 }
 328                 let mut cts = 0;
 329                 match codec_tag {
 330                     4 | 5 => {
 331                                           self.src.read_skip(1)?;
 332                     },
 333                     7 => {
 334                         let pkt_type            = self.src.read_byte()?;
 335                         if pkt_type == 1 {
 336                             cts                 = ((self.src.read_u24be()? << 8) as i32) >> 8;
 337                         } else if pkt_type == 2 {
 338                             let pkt_hdr_size = (self.src.tell() - pkt_start) as usize;
 339                             validate!(data_size >= pkt_hdr_size);
 340                             data_size -= pkt_hdr_size;
 341                                                   self.src.read_skip(data_size)?;
 342                             return Ok(());
 343                         }
 344                     },
 345                     _ => {},
 346                 };
 347
 348                 let pkt_hdr_size = (self.src.tell() - pkt_start) as usize;
 349                 validate!(data_size >= pkt_hdr_size);
 350                 data_size -= pkt_hdr_size;
 351
 352                 if data_size > 0 {
 353                     let stream = strmgr.get_stream(self.vstream).unwrap();
 354                     let pts = (u64::from(ext_time) << 24) | u64::from(time);
 355                     let dts = ((pts as i64) + i64::from(cts)).max(0) as u64;
 356                     let ts = stream.make_ts(Some(pts), Some(dts), None);
 357                     self.vpkts.push(self.src.read_packet(stream, ts, ftype == FrameType::I, data_size)?);
 358                 }
 359             },
 360             18 => {
 361                 let end = self.src.tell() + (data_size as u64);
 362                 let ntype           = self.src.read_byte()?;
 363                 validate!(ntype == 2);
 364                 let nlen            = self.src.read_u16be()? as usize;
 365                 validate!(nlen > 0);
 366                 let mut name = vec![0; nlen];
 367                                       self.src.read_buf(&mut name)?;
 368                 if &name == b"onMetaData" {
 369                     let otype               = self.src.read_byte()?;
 370                     validate!(otype == 8);
 371                     let _size               = self.src.read_u32be()?;
 372                     while self.src.tell() < end {
 373                         let nlen            = self.src.read_u16be()? as usize;
 374                         if nlen == 0 {
 375                             let emarker     = self.src.peek_byte()?;
 376                             if emarker == 9 {
 377                                               self.src.read_skip(1)?;
 378                                 break;
 379                             }
 380                         }
 381                         let mut name = vec![0; nlen];
 382                                               self.src.read_buf(&mut name)?;
 383                         let vtype           = self.src.read_byte()?;
 384                         match vtype {
 385                             0 => {
 386                                 let val     = self.src.read_f64be()?;
 387                                 match name.as_slice() {
 388                                     b"duration" => self.duration = (val * 1000.0) as u64,
 389                                     b"width"    => self.width  = val as usize,
 390                                     b"height"   => self.height = val as usize,
 391                                     b"videocodecid" => {
 392                                         let codec_tag = val as u8;
 393                                         if self.vtag.is_none() && codec_tag != AVC_ID && self.width != 0 && self.height != 0 {
 394                                             let cname = get_vcodec_name(codec_tag)?;
 395                                             let edata = if cname.starts_with("vp6") {
 396                                                     let ebyte = ((16 - (self.width & 0xF)) & 0xF) * 16 + ((16 - (self.height & 0xF)) & 0xF);
 397                                                     Some(vec![ebyte as u8])
 398                                                 } else { None };
 399                                             let vhdr = NAVideoInfo::new(self.width, self.height, false, YUV420_FORMAT);
 400                                             let vci = NACodecTypeInfo::Video(vhdr);
 401                                             let vinfo = NACodecInfo::new(cname, vci, edata);
 402                                             if let Some(id) = strmgr.add_stream(NAStream::new(StreamType::Video, 0, vinfo, 1, 1000, 0)) {
 403                                                 self.vstream = id;
 404                                             } else {
 405                                                 return Err(DemuxerError::MemoryError);
 406                                             }
 407                                             self.vtag = Some(codec_tag);
 408                                         }
 409                                     },
 410                                     _ => {},
 411                                 };
 412                             },
 413                             1 => {
 414                                 let _val    = self.src.read_byte()?;
 415                             },
 416                             2 => {
 417                                 let len     = self.src.read_u16be()? as usize;
 418                                 let mut val = vec![0; len];
 419                                               self.src.read_buf(&mut val)?;
 420                             },
 421                             3 => {
 422                                 break;//unimplemented!();
 423                             },
 424                             5 => {},
 425                             6 => {},
 426                             7 => {
 427                                               self.src.read_u16be()?;
 428                             },
 429                             8 => {
 430                                 unimplemented!();
 431                             },
 432                             10 => {
 433                                 unimplemented!();
 434                             },
 435                             11 => {
 436                                               self.src.read_f64be()?;
 437                                               self.src.read_u16be()?;
 438                             },
 439                             12 => {
 440                                 let len     = self.src.read_u16be()? as usize;
 441                                 let mut val = vec![0; len];
 442                                               self.src.read_buf(&mut val)?;
 443                             },
 444                             _ => break,
 445                         };
 446                     }
 447                 }
 448                 validate!(self.src.tell() <= end);
 449                 let to_skip = (end - self.src.tell()) as usize;
 450                                           self.src.read_skip(to_skip)?;
 451             },
 452             _ => {
 453                                           self.src.read_skip(data_size)?;
 454             },
 455         };
 456         Ok(())
 457     }
 458 }
 459
 460 pub struct FLVDemuxerCreator { }
 461
 462 impl DemuxerCreator for FLVDemuxerCreator {
 463     fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> {
 464         Box::new(FLVDemuxer::new(br))
 465     }
 466     fn get_name(&self) -> &'static str { "flv" }
 467 }
 468
 469 #[cfg(test)]
 470 mod test {
 471     use super::*;
 472     use std::fs::File;
 473
 474     // sample: https://samples.mplayerhq.hu/A-codecs/Nelly_Moser/input.flv
 475     #[test]
 476     fn test_flv_demux() {
 477         let mut file = File::open("assets/Flash/input.flv").unwrap();
 478         let mut fr = FileReader::new_read(&mut file);
 479         let mut br = ByteReader::new(&mut fr);
 480         let mut dmx = FLVDemuxer::new(&mut br);
 481         let mut sm = StreamManager::new();
 482         let mut si = SeekIndex::new();
 483         dmx.open(&mut sm, &mut si).unwrap();
 484
 485         loop {
 486             let pktres = dmx.get_frame(&mut sm);
 487             if let Err(e) = pktres {
 488                 if e == DemuxerError::EOF { break; }
 489                 panic!("error");
 490             }
 491             let pkt = pktres.unwrap();
 492             println!("Got {}", pkt);
 493         }
 494     }
 495     #[test]
 496     fn test_flv_demux_back() {
 497         let mut file = File::open("assets/Flash/input.flv").unwrap();
 498         let mut fr = FileReader::new_read(&mut file);
 499         let mut br = ByteReader::new(&mut fr);
 500         let mut dmx = FLVDemuxer::new(&mut br);
 501         let mut sm = StreamManager::new();
 502         let mut si = SeekIndex::new();
 503         dmx.open(&mut sm, &mut si).unwrap();
 504         dmx.src.seek(SeekFrom::End(-4)).unwrap();
 505         dmx.seek(NATimePoint::Milliseconds(7500), &si).unwrap();
 506
 507         loop {
 508             let pktres = dmx.get_frame(&mut sm);
 509             if let Err(e) = pktres {
 510                 if e == DemuxerError::EOF { break; }
 511                 panic!("error");
 512             }
 513             let pkt = pktres.unwrap();
 514             println!("Got {}", pkt);
 515         }
 516     }
 517 }