diff --git a/src/disc/bluray.rs b/src/disc/bluray.rs index fcb2308..a0edc02 100644 --- a/src/disc/bluray.rs +++ b/src/disc/bluray.rs @@ -187,6 +187,7 @@ impl Disc { chapters, extents, content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), }) } diff --git a/src/disc/dvd.rs b/src/disc/dvd.rs index 8f13438..c219973 100644 --- a/src/disc/dvd.rs +++ b/src/disc/dvd.rs @@ -111,6 +111,7 @@ impl Disc { chapters: Vec::new(), extents, content_format: ContentFormat::MpegPs, + codec_privates: Vec::new(), }); } } diff --git a/src/disc/mod.rs b/src/disc/mod.rs index 4e00694..fec1886 100644 --- a/src/disc/mod.rs +++ b/src/disc/mod.rs @@ -116,6 +116,9 @@ pub struct DiscTitle { pub extents: Vec, /// Content format for this title pub content_format: ContentFormat, + /// Codec initialization data per stream (SPS/PPS, etc). + /// Index matches `streams`. None for streams without codec init data. + pub codec_privates: Vec>>, } /// A clip reference within a title. @@ -797,6 +800,7 @@ impl DiscTitle { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), } } @@ -1502,6 +1506,7 @@ mod tests { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), } } diff --git a/src/mux/m2ts.rs b/src/mux/m2ts.rs index ec4d6eb..01031f9 100644 --- a/src/mux/m2ts.rs +++ b/src/mux/m2ts.rs @@ -90,7 +90,8 @@ impl M2tsStream { let file_size = reader.seek(SeekFrom::End(0))?; reader.seek(SeekFrom::Start(0))?; - // Try FMKV metadata header + // Try FMKV metadata header — save position so we can seek back on failure + let start = reader.stream_position()?; if let Ok(Some(m)) = meta::read_header(&mut reader) { let header_end = reader.stream_position()?; let content_size = file_size.saturating_sub(header_end); @@ -113,8 +114,8 @@ impl M2tsStream { }); } - // Fallback: scan PMT for streams, PTS for duration - reader.seek(SeekFrom::Start(0))?; + // No FMKV header — seek back and try PMT scan + reader.seek(SeekFrom::Start(start))?; let mut buf = vec![0u8; SCAN_SIZE]; let n = reader.read(&mut buf)?; diff --git a/src/mux/meta.rs b/src/mux/meta.rs index 2c1c187..80b5070 100644 --- a/src/mux/meta.rs +++ b/src/mux/meta.rs @@ -5,7 +5,7 @@ use crate::disc::{AudioStream, ColorSpace, DiscTitle, Stream, SubtitleStream, VideoStream}; use serde::{Deserialize, Serialize}; -use std::io::{self, Read, Seek, SeekFrom, Write}; +use std::io::{self, Read, Write}; /// Magic bytes: "FMKV" + version 1 + 2 reserved bytes. const MAGIC: [u8; 8] = [b'F', b'M', b'K', b'V', 0x00, 0x01, 0x00, 0x00]; @@ -77,26 +77,14 @@ pub enum MetaStream { } impl M2tsMeta { - /// Build metadata from a disc Title with optional codec_private per stream. - pub fn from_title_with_privates(title: &DiscTitle, codec_privates: &[Option>]) -> Self { - let mut meta = Self::from_title(title); - for (i, s) in meta.streams.iter_mut().enumerate() { - if let MetaStream::Video { codec_private, .. } = s { - if let Some(Some(cp)) = codec_privates.get(i) { - use base64::Engine; - *codec_private = Some(base64::engine::general_purpose::STANDARD.encode(cp)); - } - } - } - meta - } - - /// Build metadata from a disc Title. + /// Build metadata from a DiscTitle. Codec privates come from title.codec_privates. pub fn from_title(title: &DiscTitle) -> Self { + use base64::Engine; let streams = title .streams .iter() - .map(|s| match s { + .enumerate() + .map(|(i, s)| match s { Stream::Video(v) => MetaStream::Video { pid: v.pid, codec: v.codec.id().into(), @@ -105,7 +93,9 @@ impl M2tsMeta { hdr: v.hdr.id().into(), label: v.label.clone(), secondary: v.secondary, - codec_private: None, + codec_private: title.codec_privates.get(i) + .and_then(|cp| cp.as_ref()) + .map(|cp| base64::engine::general_purpose::STANDARD.encode(cp)), }, Stream::Audio(a) => MetaStream::Audio { pid: a.pid, @@ -200,6 +190,7 @@ impl M2tsMeta { chapters: Vec::new(), extents: Vec::new(), content_format: crate::disc::ContentFormat::BdTs, + codec_privates: self.codec_privates(), } } @@ -237,27 +228,21 @@ pub fn write_header(w: &mut impl Write, meta: &M2tsMeta) -> io::Result<()> { Ok(()) } -/// Try to read a metadata header from the start of an m2ts file. -/// Returns None for bare m2ts files (no header). -/// On success, leaves reader positioned at the first TS packet. -/// On failure, seeks back to the start. -pub fn read_header(r: &mut R) -> io::Result> { - let start = r.stream_position()?; +/// Try to read an FMKV metadata header. +/// Returns None if magic bytes don't match. Consumes header bytes on success. +/// Caller handles seek-back on failure if needed (e.g. for fallback PMT scan). +pub fn read_header(r: &mut impl Read) -> io::Result> { + const MAX_JSON_SIZE: usize = 10 * 1024 * 1024; // 10 MB let mut magic = [0u8; 8]; if r.read_exact(&mut magic).is_err() { - r.seek(SeekFrom::Start(start))?; return Ok(None); } if magic[..4] != MAGIC[..4] { - // Not a freemkv m2ts — seek back - r.seek(SeekFrom::Start(start))?; return Ok(None); } - const MAX_JSON_SIZE: usize = 10 * 1024 * 1024; // 10 MB - let mut len_buf = [0u8; 4]; r.read_exact(&mut len_buf)?; let json_len = u32::from_be_bytes(len_buf) as usize; @@ -275,42 +260,6 @@ pub fn read_header(r: &mut R) -> io::Result> { let raw_len = 8 + 4 + json_len; let padded_len = raw_len.div_ceil(PACKET_SIZE) * PACKET_SIZE; let padding = padded_len - raw_len; - if padding > 0 { - r.seek(SeekFrom::Current(padding as i64))?; - } - - Ok(Some(meta)) -} - -/// Read a metadata header from a forward-only stream (no Seek required). -/// Returns None if the magic bytes don't match. Consumes the header bytes. -pub fn read_header_from_stream(r: &mut impl Read) -> io::Result> { - let mut magic = [0u8; 8]; - r.read_exact(&mut magic)?; - - if magic[..4] != MAGIC[..4] { - return Ok(None); - } - - const MAX_JSON_SIZE: usize = 10 * 1024 * 1024; - - let mut len_buf = [0u8; 4]; - r.read_exact(&mut len_buf)?; - let json_len = u32::from_be_bytes(len_buf) as usize; - if json_len > MAX_JSON_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, "FMKV JSON too large")); - } - - let mut json_buf = vec![0u8; json_len]; - r.read_exact(&mut json_buf)?; - - let meta: M2tsMeta = serde_json::from_slice(&json_buf) - .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; - - // Skip padding - let raw_len = 8 + 4 + json_len; - let padded_len = raw_len.div_ceil(PACKET_SIZE) * PACKET_SIZE; - let padding = padded_len - raw_len; if padding > 0 { let mut skip = vec![0u8; padding]; r.read_exact(&mut skip)?; diff --git a/src/mux/mkvout.rs b/src/mux/mkvout.rs index 1454dbc..f0324da 100644 --- a/src/mux/mkvout.rs +++ b/src/mux/mkvout.rs @@ -15,13 +15,10 @@ pub struct MkvOutputStream { } impl MkvOutputStream { - /// Create an MKV output stream. - /// `codec_privates` provides initialization data per track (from InputStream). - /// Tracks without codec_private get None. + /// Create an MKV output stream. Codec privates come from title.codec_privates. pub fn create( writer: Box, title: &DiscTitle, - codec_privates: &[Option>], ) -> io::Result { let mut tracks = Vec::new(); for (idx, s) in title.streams.iter().enumerate() { @@ -30,7 +27,7 @@ impl MkvOutputStream { crate::disc::Stream::Audio(a) => MkvTrack::audio(a), crate::disc::Stream::Subtitle(s) => MkvTrack::subtitle(s), }; - if let Some(cp) = codec_privates.get(idx).and_then(|c| c.as_ref()) { + if let Some(cp) = title.codec_privates.get(idx).and_then(|c| c.as_ref()) { track.codec_private = Some(cp.clone()); } tracks.push(track); diff --git a/src/mux/network.rs b/src/mux/network.rs index d614f39..9bcfa09 100644 --- a/src/mux/network.rs +++ b/src/mux/network.rs @@ -58,8 +58,8 @@ impl NetworkStream { stream.set_nodelay(true)?; let mut reader = BufReader::with_capacity(NET_BUF_SIZE, stream); - // Read FMKV metadata header (inline, since TcpStream doesn't impl Seek) - let disc_title = meta::read_header_from_stream(&mut reader)? + // Read FMKV metadata header + let disc_title = meta::read_header(&mut reader)? .ok_or_else(|| { io::Error::new( io::ErrorKind::InvalidData, @@ -149,6 +149,7 @@ mod tests { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), } } diff --git a/src/mux/pesout.rs b/src/mux/pesout.rs index 4c7ae3c..5b4269f 100644 --- a/src/mux/pesout.rs +++ b/src/mux/pesout.rs @@ -13,23 +13,17 @@ pub struct M2tsOutputStream { } impl M2tsOutputStream { - pub fn create( - path: &str, - title: &DiscTitle, - codec_privates: &[Option>], - ) -> io::Result { + pub fn create(path: &str, title: &DiscTitle) -> io::Result { let file = std::fs::File::create(path) .map_err(|e| io::Error::new(e.kind(), format!("m2ts://{}: {}", path, e)))?; let mut writer = io::BufWriter::with_capacity(4 * 1024 * 1024, file); - // Write FMKV metadata header with codec_privates so M2tsStream::open can read them back if !title.streams.is_empty() { - let m = super::meta::M2tsMeta::from_title_with_privates(title, codec_privates); + let m = super::meta::M2tsMeta::from_title(title); super::meta::write_header(&mut writer, &m)?; } let pids = extract_pids(title); let mut muxer = TsMuxer::new(writer, &pids); - // Pass codec_privates to TsMuxer for Annex B parameter set injection - for (i, cp) in codec_privates.iter().enumerate() { + for (i, cp) in title.codec_privates.iter().enumerate() { if let Some(data) = cp { muxer.set_codec_private(i, data.clone()); } @@ -45,7 +39,7 @@ impl crate::pes::Stream for M2tsOutputStream { fn write(&mut self, frame: &PesFrame) -> io::Result<()> { self.muxer.write_frame(frame.track, frame.pts, &frame.data) } - fn finish(&mut self) -> io::Result<()> { self.muxer.finish_ref() } + fn finish(&mut self) -> io::Result<()> { self.muxer.finish() } fn info(&self) -> &DiscTitle { &self.title } } @@ -99,7 +93,6 @@ impl NetworkOutputStream { pub fn connect(addr: &str, title: &DiscTitle) -> io::Result { let stream = std::net::TcpStream::connect(addr)?; let mut writer = io::BufWriter::with_capacity(256 * 1024, stream); - // Send FMKV metadata header immediately so receiver can read it if !title.streams.is_empty() { let m = super::meta::M2tsMeta::from_title(title); super::meta::write_header(&mut writer, &m)?; diff --git a/src/mux/resolve.rs b/src/mux/resolve.rs index bae4b21..3758db9 100644 --- a/src/mux/resolve.rs +++ b/src/mux/resolve.rs @@ -370,7 +370,6 @@ pub fn input(url: &str, opts: &InputOptions) -> io::Result>], ) -> io::Result> { let parsed = parse_url(url); match parsed { @@ -380,11 +379,11 @@ pub fn output( .map_err(|e| io::Error::new(e.kind(), format!("mkv://{}: {}", path.display(), e)))?; let writer: Box = Box::new(std::io::BufWriter::with_capacity(IO_BUF_SIZE, file)); - Ok(Box::new(super::mkvout::MkvOutputStream::create(writer, title, codec_privates)?)) + Ok(Box::new(super::mkvout::MkvOutputStream::create(writer, title)?)) } StreamUrl::M2ts { ref path } => { validate_file_path(path, "m2ts")?; - Ok(Box::new(super::pesout::M2tsOutputStream::create(&path.to_string_lossy(), title, codec_privates)?)) + Ok(Box::new(super::pesout::M2tsOutputStream::create(&path.to_string_lossy(), title)?)) } StreamUrl::Network { ref addr } => { validate_network_addr(addr)?; diff --git a/src/mux/tsmux.rs b/src/mux/tsmux.rs index 8d1fee6..ba56911 100644 --- a/src/mux/tsmux.rs +++ b/src/mux/tsmux.rs @@ -139,11 +139,7 @@ impl TsMuxer { Ok(()) } - pub fn finish(mut self) -> io::Result<()> { - self.writer.flush() - } - - pub fn finish_ref(&mut self) -> io::Result<()> { + pub fn finish(&mut self) -> io::Result<()> { self.writer.flush() } } diff --git a/tests/disc_tests.rs b/tests/disc_tests.rs index e1e2733..d488ac2 100644 --- a/tests/disc_tests.rs +++ b/tests/disc_tests.rs @@ -171,6 +171,7 @@ fn title_with_video( chapters: Vec::new(), extents: Vec::new(), content_format, + codec_privates: Vec::new(), } } diff --git a/tests/streams.rs b/tests/streams.rs index 0c397e0..6decb00 100644 --- a/tests/streams.rs +++ b/tests/streams.rs @@ -51,6 +51,7 @@ fn sample_disc_title() -> DiscTitle { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), } } @@ -452,6 +453,7 @@ fn meta_codec_roundtrip() { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), }; let meta = M2tsMeta::from_title(&dt); @@ -486,6 +488,7 @@ fn meta_empty_streams() { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), }; let meta = M2tsMeta::from_title(&dt); @@ -505,6 +508,7 @@ fn meta_all_stream_types() { clips: Vec::new(), chapters: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), streams: vec![ Stream::Video(VideoStream { pid: 0x1011, @@ -651,6 +655,7 @@ fn mkvstream_roundtrip_bdts() { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), }; let output = Cursor::new(Vec::new()); @@ -731,6 +736,7 @@ fn mkvstream_meta_preserves_all_streams() { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), }; let output = Cursor::new(Vec::new()); @@ -791,6 +797,7 @@ fn mkvstream_e2e_h264_produces_valid_mkv() { chapters: Vec::new(), extents: Vec::new(), content_format: ContentFormat::BdTs, + codec_privates: Vec::new(), }; // Build synthetic BD-TS packets containing valid H.264 NALs