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 (tb_num, tb_den) = stream.get_timebase();
 219                     let pts = (u64::from(ext_time) << 24) | u64::from(time);
 220                     let ts = NATimeInfo::new(Some(pts), None, None, tb_num, tb_den);
 221                     self.apkts.push(self.src.read_packet(stream, ts, true, data_size)?);
 222                 }
 223             },
 224             9 => {
 225                 let hdr                 = self.src.read_byte()?;
 226                 let ftype = match hdr >> 4 {
 227                         1 => FrameType::I,
 228                         2 => FrameType::P,
 229                         3 => FrameType::P, // droppable
 230                         4 => FrameType::Other, // generated key frame
 231                         5 => FrameType::Other, // video info/command frame
 232                         _ => return Err(DemuxerError::InvalidData),
 233                     };
 234                 let codec_tag = hdr & 0xF;
 235                 if let Some(id) = self.vtag {
 236                     validate!(id == codec_tag);
 237                 } else {
 238                     let cname = get_vcodec_name(codec_tag)?;
 239                     let is_avc = codec_tag == AVC_ID;
 240                     if is_avc {
 241                         let pkt_type            = self.src.read_byte()?;
 242                         validate!(pkt_type == 0);
 243                                                   self.src.read_u24be()?;
 244                     }
 245                     let mut edata = None;
 246                     let (width, height) = match codec_tag {
 247                             2 => {
 248                                 let mut buf = [0; 9];
 249                                                   self.src.peek_buf(&mut buf)?;
 250                                 let mut br = BitReader::new(&buf, BitReaderMode::BE);
 251                                                 br.skip(30).unwrap_or(());
 252                                 let sfmt      = br.read(3).unwrap_or(7);
 253                                 match sfmt {
 254                                     0 => {
 255                                         let w           = br.read(8).unwrap_or(0) as usize;
 256                                         let h           = br.read(8).unwrap_or(0) as usize;
 257                                         (w, h)
 258                                     },
 259                                     1 => {
 260                                         let w           = br.read(16).unwrap_or(0) as usize;
 261                                         let h           = br.read(16).unwrap_or(0) as usize;
 262                                         (w, h)
 263                                     },
 264                                     2 => (352, 288),
 265                                     3 => (176, 144),
 266                                     4 => (128, 96),
 267                                     5 => (320, 240),
 268                                     6 => (160, 120),
 269                                     _ => (0, 0),
 270                                 }
 271                             },
 272                             3 | 6 => {
 273                                 let mut buf = [0; 4];
 274                                                   self.src.peek_buf(&mut buf)?;
 275                                 let w = (read_u16be(&buf[0..])? & 0xFFF) as usize;
 276                                 let h = (read_u16be(&buf[2..])? & 0xFFF) as usize;
 277                                 (w, h)
 278                             },
 279                             4 => {
 280                                 let mut buf = [0; 7];
 281                                                   self.src.peek_buf(&mut buf)?;
 282                                 let off = if (buf[1] & 1) != 0 || (buf[2] & 6) == 0 { 5 } else { 3 };
 283                                 validate!(buf[off] != 0 && buf[off + 1] != 0);
 284                                 let w = usize::from(buf[off + 1]) * 16 - usize::from(buf[0] >> 4);
 285                                 let h = usize::from(buf[off])     * 16 - usize::from(buf[0] & 0xF);
 286
 287                                 edata = Some(vec![buf[0]]);
 288
 289                                 (w, h)
 290                             },
 291                             5 => {
 292                                 let mut buf = [0; 10];
 293                                                   self.src.peek_buf(&mut buf)?;
 294                                 let off = if (buf[4] & 1) != 0 || (buf[5] & 6) == 0 { 8 } else { 6 };
 295                                 validate!(buf[off] != 0 && buf[off + 1] != 0);
 296                                 let w = usize::from(buf[off + 1]) * 16 - usize::from(buf[0] >> 4);
 297                                 let h = usize::from(buf[off])     * 16 - usize::from(buf[0] & 0xF);
 298
 299                                 edata = Some(vec![buf[0]]);
 300
 301                                 (w, h)
 302                             },
 303                             7 => {
 304                                 let pkt_hdr_size = (self.src.tell() - pkt_start) as usize;
 305                                 validate!(data_size >= pkt_hdr_size);
 306                                 data_size -= pkt_hdr_size;
 307                                 let mut data = vec![0; data_size + 4];
 308                                 data[..4].copy_from_slice(b"avcC");
 309                                                   self.src.read_buf(&mut data[4..])?;
 310                                 edata = Some(data);
 311                                 (self.width, self.height)
 312                             },
 313                             _ => unreachable!(),
 314                         };
 315
 316                     let vhdr = NAVideoInfo::new(width, height, false, YUV420_FORMAT);
 317                     let vci = NACodecTypeInfo::Video(vhdr);
 318                     let vinfo = NACodecInfo::new(cname, vci, edata);
 319                     if let Some(id) = strmgr.add_stream(NAStream::new(StreamType::Video, 0, vinfo, 1, 1000, 0)) {
 320                         self.vstream = id;
 321                     } else {
 322                         return Err(DemuxerError::MemoryError);
 323                     }
 324                     self.vtag = Some(codec_tag);
 325                     if is_avc {
 326                         return Ok(());
 327                     }
 328                 }
 329                 let mut cts = 0;
 330                 match codec_tag {
 331                     4 | 5 => {
 332                                           self.src.read_skip(1)?;
 333                     },
 334                     7 => {
 335                         let pkt_type            = self.src.read_byte()?;
 336                         if pkt_type == 1 {
 337                             cts                 = ((self.src.read_u24be()? << 8) as i32) >> 8;
 338                         } else if pkt_type == 2 {
 339                             let pkt_hdr_size = (self.src.tell() - pkt_start) as usize;
 340                             validate!(data_size >= pkt_hdr_size);
 341                             data_size -= pkt_hdr_size;
 342                                                   self.src.read_skip(data_size)?;
 343                             return Ok(());
 344                         }
 345                     },
 346                     _ => {},
 347                 };
 348
 349                 let pkt_hdr_size = (self.src.tell() - pkt_start) as usize;
 350                 validate!(data_size >= pkt_hdr_size);
 351                 data_size -= pkt_hdr_size;
 352
 353                 if data_size > 0 {
 354                     let stream = strmgr.get_stream(self.vstream).unwrap();
 355                     let (tb_num, tb_den) = stream.get_timebase();
 356                     let pts = (u64::from(ext_time) << 24) | u64::from(time);
 357                     let dts = ((pts as i64) + i64::from(cts)).max(0) as u64;
 358                     let ts = NATimeInfo::new(Some(pts), Some(dts), None, tb_num, tb_den);
 359                     self.vpkts.push(self.src.read_packet(stream, ts, ftype == FrameType::I, data_size)?);
 360                 }
 361             },
 362             18 => {
 363                 let end = self.src.tell() + (data_size as u64);
 364                 let ntype           = self.src.read_byte()?;
 365                 validate!(ntype == 2);
 366                 let nlen            = self.src.read_u16be()? as usize;
 367                 validate!(nlen > 0);
 368                 let mut name = vec![0; nlen];
 369                                       self.src.read_buf(&mut name)?;
 370                 if &name == b"onMetaData" {
 371                     let otype               = self.src.read_byte()?;
 372                     validate!(otype == 8);
 373                     let _size               = self.src.read_u32be()?;
 374                     while self.src.tell() < end {
 375                         let nlen            = self.src.read_u16be()? as usize;
 376                         if nlen == 0 {
 377                             let emarker     = self.src.peek_byte()?;
 378                             if emarker == 9 {
 379                                               self.src.read_skip(1)?;
 380                                 break;
 381                             }
 382                         }
 383                         let mut name = vec![0; nlen];
 384                                               self.src.read_buf(&mut name)?;
 385                         let vtype           = self.src.read_byte()?;
 386                         match vtype {
 387                             0 => {
 388                                 let val     = self.src.read_f64be()?;
 389                                 match name.as_slice() {
 390                                     b"duration" => self.duration = (val * 1000.0) as u64,
 391                                     b"width"    => self.width  = val as usize,
 392                                     b"height"   => self.height = val as usize,
 393                                     b"videocodecid" => {
 394                                         let codec_tag = val as u8;
 395                                         if self.vtag.is_none() && codec_tag != AVC_ID && self.width != 0 && self.height != 0 {
 396                                             let cname = get_vcodec_name(codec_tag)?;
 397                                             let edata = if cname.starts_with("vp6") {
 398                                                     let ebyte = ((16 - (self.width & 0xF)) & 0xF) * 16 + ((16 - (self.height & 0xF)) & 0xF);
 399                                                     Some(vec![ebyte as u8])
 400                                                 } else { None };
 401                                             let vhdr = NAVideoInfo::new(self.width, self.height, false, YUV420_FORMAT);
 402                                             let vci = NACodecTypeInfo::Video(vhdr);
 403                                             let vinfo = NACodecInfo::new(cname, vci, edata);
 404                                             if let Some(id) = strmgr.add_stream(NAStream::new(StreamType::Video, 0, vinfo, 1, 1000, 0)) {
 405                                                 self.vstream = id;
 406                                             } else {
 407                                                 return Err(DemuxerError::MemoryError);
 408                                             }
 409                                             self.vtag = Some(codec_tag);
 410                                         }
 411                                     },
 412                                     _ => {},
 413                                 };
 414                             },
 415                             1 => {
 416                                 let _val    = self.src.read_byte()?;
 417                             },
 418                             2 => {
 419                                 let len     = self.src.read_u16be()? as usize;
 420                                 let mut val = vec![0; len];
 421                                               self.src.read_buf(&mut val)?;
 422                             },
 423                             3 => {
 424                                 break;//unimplemented!();
 425                             },
 426                             5 => {},
 427                             6 => {},
 428                             7 => {
 429                                               self.src.read_u16be()?;
 430                             },
 431                             8 => {
 432                                 unimplemented!();
 433                             },
 434                             10 => {
 435                                 unimplemented!();
 436                             },
 437                             11 => {
 438                                               self.src.read_f64be()?;
 439                                               self.src.read_u16be()?;
 440                             },
 441                             12 => {
 442                                 let len     = self.src.read_u16be()? as usize;
 443                                 let mut val = vec![0; len];
 444                                               self.src.read_buf(&mut val)?;
 445                             },
 446                             _ => break,
 447                         };
 448                     }
 449                 }
 450                 validate!(self.src.tell() <= end);
 451                 let to_skip = (end - self.src.tell()) as usize;
 452                                           self.src.read_skip(to_skip)?;
 453             },
 454             _ => {
 455                                           self.src.read_skip(data_size)?;
 456             },
 457         };
 458         Ok(())
 459     }
 460 }
 461
 462 pub struct FLVDemuxerCreator { }
 463
 464 impl DemuxerCreator for FLVDemuxerCreator {
 465     fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> {
 466         Box::new(FLVDemuxer::new(br))
 467     }
 468     fn get_name(&self) -> &'static str { "flv" }
 469 }
 470
 471 #[cfg(test)]
 472 mod test {
 473     use super::*;
 474     use std::fs::File;
 475
 476     #[test]
 477     fn test_flv_demux() {
 478         let mut file = File::open("assets/Flash/input.flv").unwrap();
 479         let mut fr = FileReader::new_read(&mut file);
 480         let mut br = ByteReader::new(&mut fr);
 481         let mut dmx = FLVDemuxer::new(&mut br);
 482         let mut sm = StreamManager::new();
 483         let mut si = SeekIndex::new();
 484         dmx.open(&mut sm, &mut si).unwrap();
 485
 486         loop {
 487             let pktres = dmx.get_frame(&mut sm);
 488             if let Err(e) = pktres {
 489                 if e == DemuxerError::EOF { break; }
 490                 panic!("error");
 491             }
 492             let pkt = pktres.unwrap();
 493             println!("Got {}", pkt);
 494         }
 495     }
 496     #[test]
 497     fn test_flv_demux_back() {
 498         let mut file = File::open("assets/Flash/input.flv").unwrap();
 499         let mut fr = FileReader::new_read(&mut file);
 500         let mut br = ByteReader::new(&mut fr);
 501         let mut dmx = FLVDemuxer::new(&mut br);
 502         let mut sm = StreamManager::new();
 503         let mut si = SeekIndex::new();
 504         dmx.open(&mut sm, &mut si).unwrap();
 505         dmx.src.seek(SeekFrom::End(-4)).unwrap();
 506         dmx.seek(NATimePoint::Milliseconds(7500), &si).unwrap();
 507
 508         loop {
 509             let pktres = dmx.get_frame(&mut sm);
 510             if let Err(e) = pktres {
 511                 if e == DemuxerError::EOF { break; }
 512                 panic!("error");
 513             }
 514             let pkt = pktres.unwrap();
 515             println!("Got {}", pkt);
 516         }
 517     }
 518 }