nihav-flash/src/muxers/flv.rs

   1 use nihav_core::muxers::*;
   2
   3 const FLV_VCODECS: &[(&str, u8)] = &[
   4     ("flv263", 2),
   5     ("flashsv", 3),
   6     ("vp6f", 4),
   7     ("vp6a", 5),
   8     ("flashsv2", 6),
   9     ("h264", AVC_ID)
  10 ];
  11
  12 const NO_CODEC: u8 = 0;
  13 const AVC_ID: u8 = 7;
  14 const AAC_ID: u8 = 10;
  15
  16 fn find_audio_tag(cname: &'static str, rate: u32, channels: u8) -> MuxerResult<u8> {
  17     if channels != 1 && channels != 2 {
  18         return Err(MuxerError::InvalidArgument);
  19     }
  20     let tag = match cname {
  21             "flv-adpcm" => 1,
  22             "pcm" => 3,
  23             "mp3" => if rate != 8000 { 2 } else { return Ok(14); },
  24             "asao" => {
  25                 if channels != 1 {
  26                     return Err(MuxerError::InvalidArgument);
  27                 }
  28                 match rate {
  29                     16000 => return Ok(4),
  30                     8000 => return Ok(5),
  31                     _ => 6,
  32                 }
  33             },
  34             "alaw" => 7,
  35             "ulaw" => 8,
  36             "aac" => AAC_ID,
  37             "speex" => 11,
  38             _ => return Err(MuxerError::InvalidArgument),
  39         };
  40     match rate {
  41         5500 | 11025 | 22050 | 44100 => {},
  42         _ => return Err(MuxerError::InvalidArgument),
  43     };
  44     Ok(tag)
  45 }
  46
  47 trait FLVPropertyWriter {
  48     fn write_property_num(&mut self, name: &str, val: f64) -> MuxerResult<()>;
  49     fn write_property_bool(&mut self, name: &str, val: bool) -> MuxerResult<()>;
  50 }
  51
  52 impl<'a> FLVPropertyWriter for ByteWriter<'a> {
  53     fn write_property_num(&mut self, name: &str, val: f64) -> MuxerResult<()> {
  54         self.write_u16be(name.len() as u16)?;
  55         self.write_buf(name.as_bytes())?;
  56         self.write_byte(0)?;
  57         self.write_f64be(val)?;
  58         Ok(())
  59     }
  60     fn write_property_bool(&mut self, name: &str, val: bool) -> MuxerResult<()> {
  61         self.write_u16be(name.len() as u16)?;
  62         self.write_buf(name.as_bytes())?;
  63         self.write_byte(1)?;
  64         self.write_byte(val as u8)?;
  65         Ok(())
  66     }
  67 }
  68
  69 macro_rules! write_packet {
  70     ($self: expr, $pkt_type: expr, $ts: expr, $code: block) => {
  71         let start = $self.bw.tell();
  72         $self.bw.write_byte($pkt_type)?;
  73         $self.bw.write_u24be(0)?;
  74         $self.bw.write_u24be($ts)?;
  75         $self.bw.write_byte(($ts >> 24) as u8)?;
  76         $self.bw.write_u24be(0)?;
  77
  78         $code
  79
  80         let end = $self.bw.tell();
  81         let size = end - start - 11;
  82         $self.bw.seek(SeekFrom::Start(start + 1))?;
  83         $self.bw.write_u24be(size as u32)?;
  84         $self.bw.seek(SeekFrom::Start(end))?;
  85         $self.bw.write_u32be((size + 11) as u32)?;
  86     }
  87 }
  88
  89 struct FLVMuxer<'a> {
  90     bw:     &'a mut ByteWriter<'a>,
  91     atag:   u8,
  92     ahdr:   u8,
  93     vtag:   u8,
  94     vp6b:   u8,
  95     time:   u32,
  96     dpos:   u64,
  97 }
  98
  99 impl<'a> FLVMuxer<'a> {
 100     fn new(bw: &'a mut ByteWriter<'a>) -> Self {
 101         Self {
 102             bw,
 103             atag:   NO_CODEC,
 104             ahdr:   0,
 105             vtag:   NO_CODEC,
 106             vp6b:   0,
 107             time:   0,
 108             dpos:   0,
 109         }
 110     }
 111     fn write_metadata(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
 112         write_packet!(self, 18, 0, {
 113             self.bw.write_buf(b"\x02\x00\x0AonMetaData\x08\x00\x00\x00\x00")?;
 114             for stream in strmgr.iter() {
 115                 match stream.get_info().get_properties() {
 116                     NACodecTypeInfo::Video(ref vinfo) => {
 117                         self.bw.write_property_num("width", vinfo.width as f64)?;
 118                         self.bw.write_property_num("height", vinfo.height as f64)?;
 119                         self.bw.write_property_num("videocodecid", self.vtag as f64)?;
 120                     },
 121                     NACodecTypeInfo::Audio(ref ainfo) => {
 122                         self.bw.write_property_num("audiosamplerate", ainfo.sample_rate as f64)?;
 123                         self.bw.write_property_bool("stereo", ainfo.channels == 2)?;
 124                         self.bw.write_property_num("audiocodecid", self.atag as f64)?;
 125                     },
 126                     _ => {},
 127                 };
 128             }
 129             self.bw.write_property_num("duration", 0.0)?;
 130             self.dpos = self.bw.tell() - 8;
 131             self.bw.write_u16be(0)?;
 132             self.bw.write_byte(9)?;
 133         });
 134
 135         Ok(())
 136     }
 137 }
 138
 139 impl<'a> MuxCore<'a> for FLVMuxer<'a> {
 140     fn create(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
 141         if strmgr.get_num_streams() == 0 || strmgr.get_num_streams() > 2 {
 142             return Err(MuxerError::InvalidArgument);
 143         }
 144
 145         let mut astream = None;
 146         let mut vstream = None;
 147         for stream in strmgr.iter() {
 148             let cname = stream.get_info().get_name();
 149             match stream.get_media_type() {
 150                 StreamType::Video => {
 151                     vstream = Some(stream.clone());
 152                     if self.vtag != NO_CODEC {
 153                         return Err(MuxerError::InvalidArgument);
 154                     }
 155                     for &(name, tag) in FLV_VCODECS.iter() {
 156                         if name == cname {
 157                             self.vtag = tag;
 158                             break;
 159                         }
 160                     }
 161                     if self.vtag == NO_CODEC {
 162                         return Err(MuxerError::UnsupportedFormat);
 163                     }
 164                 },
 165                 StreamType::Audio => {
 166                     astream = Some(stream.clone());
 167                     if self.atag != NO_CODEC {
 168                         return Err(MuxerError::InvalidArgument);
 169                     }
 170                     if let Some(ainfo) = stream.get_info().get_properties().get_audio_info() {
 171                         self.atag = find_audio_tag(cname, ainfo.sample_rate, ainfo.channels)?;
 172                         self.ahdr = (self.atag << 4) |
 173                             (match ainfo.sample_rate {
 174                                 5500 => 0,
 175                                 11025 => 1,
 176                                 22050 => 2,
 177                                 _     => 3,
 178                             } << 2) |
 179                             if ainfo.format.bits == 8 { 0 } else { 2 } |
 180                             if ainfo.channels == 1 { 0 } else { 1 };
 181                     } else {
 182                         return Err(MuxerError::InvalidArgument);
 183                     }
 184                 },
 185                 _ => return Err(MuxerError::UnsupportedFormat),
 186             };
 187         }
 188
 189         self.bw.write_buf(b"FLV\x01")?;
 190         let flags = 0x8 | if self.atag != NO_CODEC { 4 } else { 0 } | if self.vtag != NO_CODEC { 1 } else { 0 };
 191         self.bw.write_byte(flags)?;
 192         self.bw.write_u32be(9)?;
 193         self.bw.write_u32be(0)?;
 194
 195         self.write_metadata(strmgr)?;
 196
 197         if let (true, Some(ref stream)) = (self.vtag == 4 || self.vtag == 5, &vstream) {
 198             if let Some(edata) = stream.get_info().get_extradata() {
 199                 if !edata.is_empty() {
 200                     self.vp6b = edata[0];
 201                 }
 202             }
 203         }
 204
 205         if let (true, Some(stream)) = (self.vtag == AVC_ID, vstream) {
 206             if let Some(edata) = stream.get_info().get_extradata() {
 207                 validate!(edata.len() > 4);
 208                 write_packet!(self, 9, 0, {
 209                     self.bw.write_byte(0x57)?;
 210                     self.bw.write_byte(0x00)?;
 211                     self.bw.write_u24be(0)?;
 212                     self.bw.write_buf(&edata[4..])?;
 213                 });
 214             }
 215         }
 216         if let (true, Some(stream)) = (self.atag == AAC_ID, astream) {
 217             if let Some(edata) = stream.get_info().get_extradata() {
 218                 write_packet!(self, 8, 0, {
 219                     self.bw.write_byte(self.ahdr)?;
 220                     self.bw.write_byte(0x00)?;
 221                     self.bw.write_buf(&edata)?;
 222                 });
 223             }
 224         }
 225
 226         Ok(())
 227     }
 228     fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
 229         let stream = pkt.get_stream();
 230         let pts = pkt.get_pts().unwrap_or(0);
 231         let ms = NATimeInfo::ts_to_time(pts, 1000, pkt.ts.tb_num, pkt.ts.tb_den) as u32;
 232         self.time = self.time.max(ms);
 233         match stream.get_media_type() {
 234             StreamType::Video => {
 235                 write_packet!(self, 9, ms, {
 236                     let hdr = self.vtag | if pkt.keyframe { 0x10 } else { 0x20 };
 237                     self.bw.write_byte(hdr)?;
 238                     match self.vtag {
 239                         4 | 5 => {
 240                             self.bw.write_byte(self.vp6b)?;
 241                         },
 242                         AVC_ID => {
 243                             self.bw.write_byte(1)?;
 244                             let cms = NATimeInfo::ts_to_time(pkt.get_pts().unwrap_or(pts), 1000, pkt.ts.tb_num, pkt.ts.tb_den) as u32;
 245                             let cts = cms.wrapping_sub(ms) << 8 >> 8;
 246                             self.bw.write_u24be(cts)?;
 247                         },
 248                         _ => {},
 249                     };
 250                     self.bw.write_buf(&pkt.get_buffer())?;
 251                 });
 252             },
 253             StreamType::Audio => {
 254                 write_packet!(self, 8, ms, {
 255                     self.bw.write_byte(self.ahdr)?;
 256                     if self.atag == AAC_ID {
 257                         self.bw.write_byte(1)?;
 258                     }
 259                     self.bw.write_buf(&pkt.get_buffer())?;
 260                 });
 261             },
 262             _ => return Err(MuxerError::InvalidData),
 263         };
 264         Ok(())
 265     }
 266     fn flush(&mut self) -> MuxerResult<()> {
 267         Ok(())
 268     }
 269     fn end(&mut self) -> MuxerResult<()> {
 270         self.bw.seek(SeekFrom::Start(self.dpos))?;
 271         self.bw.write_f64be((self.time as f64) / 1000.0)?;
 272         Ok(())
 273     }
 274 }
 275
 276 impl<'a> NAOptionHandler for FLVMuxer<'a> {
 277     fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
 278     fn set_options(&mut self, _options: &[NAOption]) { }
 279     fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
 280 }
 281
 282 pub struct FLVMuxerCreator {}
 283
 284 impl MuxerCreator for FLVMuxerCreator {
 285     fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> {
 286         Box::new(FLVMuxer::new(bw))
 287     }
 288     fn get_name(&self) -> &'static str { "flv" }
 289     fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::SingleVideoAndAudio("any", "any") }
 290 }
 291
 292 #[cfg(test)]
 293 mod test {
 294     use nihav_core::codecs::*;
 295     use nihav_core::demuxers::*;
 296     use nihav_core::muxers::*;
 297     use nihav_codec_support::test::enc_video::*;
 298     use crate::*;
 299
 300     #[test]
 301     fn test_flv_muxer() {
 302         let mut dmx_reg = RegisteredDemuxers::new();
 303         flash_register_all_demuxers(&mut dmx_reg);
 304         let dec_config = DecoderTestParams {
 305                 demuxer:        "flv",
 306                 in_name:        "assets/Flash/input.flv",
 307                 limit:          None,
 308                 stream_type:    StreamType::None,
 309                 dmx_reg, dec_reg: RegisteredDecoders::new(),
 310             };
 311         let mut mux_reg = RegisteredMuxers::new();
 312         flash_register_all_muxers(&mut mux_reg);
 313         /*let enc_config = EncoderTestParams {
 314                 muxer:      "flv",
 315                 enc_name:   "",
 316                 out_name:   "muxed.flv",
 317                 mux_reg, enc_reg: RegisteredEncoders::new(),
 318             };
 319         test_remuxing(&dec_config, &enc_config);*/
 320         test_remuxing_md5(&dec_config, "flv", &mux_reg,
 321                           [0xc777b605, 0x5777919d, 0x47996fe8, 0xf5e8d64f]);
 322     }
 323 }