From 79a4dcf2b072572cc7df8974494675263a756357 Mon Sep 17 00:00:00 2001 From: mofeng-git Date: Sat, 5 Sep 2026 14:10:53 +0800 Subject: [PATCH] =?UTF-8?q?fix(video):=20=E7=BB=9F=E4=B8=80=E5=85=B3?= =?UTF-8?q?=E9=94=AE=E5=B8=A7=E5=8F=82=E6=95=B0=E9=9B=86=E5=A4=84=E7=90=86?= =?UTF-8?q?=E5=B9=B6=E7=BC=93=E5=AD=98=E5=90=AF=E5=8A=A8=E5=B8=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 在共享管线中将长度前缀 H264 转为 Annex-B，要求编码器标记与 IDR/IRAP 内容一致，并补齐参数集后才标记可独立解码的关键帧。 向新订阅者提供缓存启动帧，启动管线时清理缓存，并移除 RustDesk 会话内重复的 SPS/PPS 处理。补充格式归一化、关键帧判定和订阅测试。 --- src/rustdesk/frame_adapters.rs | 44 -------- src/video/codec/h264_bitstream.rs | 66 +++++++++++- src/video/pipeline/shared.rs | 170 ++++++++++++++++++++++++++++-- 3 files changed, 222 insertions(+), 58 deletions(-) diff --git a/src/rustdesk/frame_adapters.rs b/src/rustdesk/frame_adapters.rs index de57bc4c..7a3ed733 100644 --- a/src/rustdesk/frame_adapters.rs +++ b/src/rustdesk/frame_adapters.rs @@ -31,8 +31,6 @@ pub struct VideoFrameAdapter { codec: VideoCodec, seq: u32, timestamp_base: u64, - h264_sps: Option, - h264_pps: Option, } impl VideoFrameAdapter { @@ -41,8 +39,6 @@ impl VideoFrameAdapter { codec, seq: 0, timestamp_base: 0, - h264_sps: None, - h264_pps: None, } } @@ -56,7 +52,6 @@ impl VideoFrameAdapter { is_keyframe: bool, timestamp_ms: u64, ) -> Message { - let data = self.prepare_h264_frame(data, is_keyframe); if self.seq == 0 { self.timestamp_base = timestamp_ms; } @@ -86,45 +81,6 @@ impl VideoFrameAdapter { msg } - fn prepare_h264_frame(&mut self, data: Bytes, is_keyframe: bool) -> Bytes { - if self.codec != VideoCodec::H264 { - return data; - } - - // Parameter sets are relevant only on random-access frames. Avoid a - // full Annex-B/AVCC scan on every delta frame in every client session. - if !is_keyframe { - return data; - } - - let (sps, pps) = crate::video::codec::h264_bitstream::extract_sps_pps(&data); - let mut has_sps = false; - let mut has_pps = false; - - if let Some(sps) = sps { - self.h264_sps = Some(Bytes::from(sps)); - has_sps = true; - } - if let Some(pps) = pps { - self.h264_pps = Some(Bytes::from(pps)); - has_pps = true; - } - - if is_keyframe && (!has_sps || !has_pps) { - if let (Some(sps), Some(pps)) = (self.h264_sps.as_ref(), self.h264_pps.as_ref()) { - let mut out = Vec::with_capacity(8 + sps.len() + pps.len() + data.len()); - out.extend_from_slice(&[0, 0, 0, 1]); - out.extend_from_slice(sps); - out.extend_from_slice(&[0, 0, 0, 1]); - out.extend_from_slice(pps); - out.extend_from_slice(&data); - return Bytes::from(out); - } - } - - data - } - pub fn encode_frame(&mut self, data: &[u8], is_keyframe: bool, timestamp_ms: u64) -> Message { self.encode_frame_from_bytes(Bytes::copy_from_slice(data), is_keyframe, timestamp_ms) } diff --git a/src/video/codec/h264_bitstream.rs b/src/video/codec/h264_bitstream.rs index eac14041..d47357e5 100644 --- a/src/video/codec/h264_bitstream.rs +++ b/src/video/codec/h264_bitstream.rs @@ -252,13 +252,37 @@ pub fn avcc_to_annex_b(data: &[u8]) -> Option> { } } -pub fn normalize_for_webrtc(data: &[u8]) -> Vec { - if is_annex_b(data) { - return strip_aud_nal_units(data); +/// Normalize a length-prefixed H.264 access unit to Annex-B when necessary. +/// +/// FFmpeg normally exposes elementary H.264 from hardware encoders as +/// Annex-B, but some V4L2 M2M drivers return AVCC-style packets. Consumers +/// such as RustDesk do not receive codec extradata from our protocol adapter, +/// so passing those packets through unchanged leaves the decoder unable to +/// find NAL unit boundaries. +pub fn normalize_annex_b(data: bytes::Bytes) -> bytes::Bytes { + // A four-byte start code is unambiguous for real encoder output. A + // three-byte prefix is not: an AVCC NAL of 256..511 bytes also begins + // with 00 00 01. Validate AVCC before accepting that shorter prefix. + if data.starts_with(&[0, 0, 0, 1]) { + return data; } - if let Some(annex_b) = avcc_to_annex_b(data) { - return strip_aud_nal_units(&annex_b); + if let Some(annex_b) = avcc_to_annex_b(data.as_ref()) { + return bytes::Bytes::from(annex_b); + } + + data +} + +pub fn normalize_for_webrtc(data: &[u8]) -> Vec { + if !data.starts_with(&[0, 0, 0, 1]) { + if let Some(annex_b) = avcc_to_annex_b(data) { + return strip_aud_nal_units(&annex_b); + } + } + + if is_annex_b(data) { + return strip_aud_nal_units(data); } data.to_vec() @@ -296,4 +320,36 @@ mod tests { Some("42402a".to_string()) ); } + + #[test] + fn converts_avcc_access_unit_to_annex_b() { + let avcc = [ + 0, 0, 0, 4, 0x67, 0x42, 0x40, 0x1f, // SPS + 0, 0, 0, 2, 0x68, 0xce, // PPS + 0, 0, 0, 3, 0x65, 0x88, 0x84, // IDR + ]; + + let annex_b = normalize_annex_b(bytes::Bytes::copy_from_slice(&avcc)); + assert!(is_annex_b(&annex_b)); + assert!(has_sps_pps(&annex_b)); + assert!(is_keyframe(&annex_b)); + } + + #[test] + fn leaves_annex_b_packet_unchanged() { + let annex_b = bytes::Bytes::from_static(&[0, 0, 0, 1, 0x65, 0x88, 0x84]); + let normalized = normalize_annex_b(annex_b.clone()); + assert_eq!(normalized, annex_b); + } + + #[test] + fn recognizes_avcc_length_that_looks_like_three_byte_start_code() { + let mut avcc = vec![0, 0, 1, 0]; + avcc.push(0x65); + avcc.resize(4 + 256, 0x88); + + let annex_b = normalize_annex_b(bytes::Bytes::from(avcc)); + assert_eq!(&annex_b[..5], &[0, 0, 0, 1, 0x65]); + assert!(is_keyframe(&annex_b)); + } } diff --git a/src/video/pipeline/shared.rs b/src/video/pipeline/shared.rs index ddfe33a0..52328fee 100644 --- a/src/video/pipeline/shared.rs +++ b/src/video/pipeline/shared.rs @@ -71,7 +71,10 @@ pub struct EncodedVideoFrame { pub data: Bytes, /// Presentation timestamp in milliseconds pub pts_ms: i64, - /// Whether this is a keyframe + /// Whether this frame can initialize a decoder without earlier frames. + /// + /// For H.264/H.265 this is stricter than the encoder packet flag: the + /// payload must be IDR/IRAP and include all required parameter sets. pub is_keyframe: bool, /// Frame sequence number pub sequence: u64, @@ -306,6 +309,8 @@ pub struct SharedVideoPipeline { /// Atomic flag for keyframe request (avoids lock contention) keyframe_requested: AtomicBool, parameter_sets: ParkingMutex, + /// Most recent random-access frame with all decoder parameter sets. + bootstrap_frame: ParkingRwLock>>, /// Pipeline start time for monotonic PTS calculation (microseconds from process start). /// Uses AtomicI64 instead of Mutex for lock-free access. pipeline_start_time_us: AtomicI64, @@ -346,6 +351,7 @@ impl SharedVideoPipeline { sequence: AtomicU64::new(0), keyframe_requested: AtomicBool::new(false), parameter_sets: ParkingMutex::new(CachedH26xParameterSets::default()), + bootstrap_frame: ParkingRwLock::new(None), pipeline_start_time_us: AtomicI64::new(0), pending_sync_geometry: ParkingMutex::new(None), device_lost_reason: ParkingMutex::new(None), @@ -404,6 +410,9 @@ impl SharedVideoPipeline { // Keep at most one pending frame so a slow WebRTC writer cannot make // the encoder wait or accumulate seconds of latency. let (tx, rx) = mpsc::channel(1); + if let Some(frame) = self.bootstrap_frame.read().clone() { + let _ = tx.try_send(frame); + } self.subscribers.write().push(tx); rx } @@ -526,8 +535,15 @@ impl SharedVideoPipeline { ) -> (Bytes, bool) { match codec { VideoEncoderType::H264 => { + let was_annex_b = h264_bitstream::is_annex_b(data.as_ref()); + let data = h264_bitstream::normalize_annex_b(data); + if !was_annex_b && h264_bitstream::is_annex_b(data.as_ref()) { + debug!("[Pipeline] Converted length-prefixed H264 packet to Annex-B"); + } let (sps, pps) = h264_bitstream::extract_sps_pps(data.as_ref()); - let is_keyframe = ffmpeg_keyframe || h264_bitstream::is_keyframe(data.as_ref()); + // Require metadata and payload to agree before advertising a + // decoder bootstrap frame. + let is_idr = ffmpeg_keyframe && h264_bitstream::is_keyframe(data.as_ref()); let mut cache = self.parameter_sets.lock(); if let Some(sps) = sps.as_ref() { cache.h264_sps = Some(sps.clone()); @@ -536,8 +552,11 @@ impl SharedVideoPipeline { cache.h264_pps = Some(pps.clone()); } - if !is_keyframe || (sps.is_some() && pps.is_some()) { - return (data, is_keyframe); + if !is_idr { + return (data, false); + } + if sps.is_some() && pps.is_some() { + return (data, true); } match (&cache.h264_sps, &cache.h264_pps) { @@ -553,12 +572,13 @@ impl SharedVideoPipeline { debug!("[Pipeline] Prepended cached SPS/PPS to H264 IDR"); (Bytes::from(output), true) } - _ => (data, true), + // An IDR without SPS/PPS is not a decoder bootstrap frame. + _ => (data, false), } } VideoEncoderType::H265 => { let (vps, sps, pps) = h265_bitstream::extract_vps_sps_pps(data.as_ref()); - let is_keyframe = ffmpeg_keyframe || h265_bitstream::is_keyframe(data.as_ref()); + let is_irap = ffmpeg_keyframe && h265_bitstream::is_keyframe(data.as_ref()); let mut cache = self.parameter_sets.lock(); if let Some(vps) = vps.as_ref() { cache.h265_vps = Some(vps.clone()); @@ -570,8 +590,11 @@ impl SharedVideoPipeline { cache.h265_pps = Some(pps.clone()); } - if !is_keyframe || (vps.is_some() && sps.is_some() && pps.is_some()) { - return (data, is_keyframe); + if !is_irap { + return (data, false); + } + if vps.is_some() && sps.is_some() && pps.is_some() { + return (data, true); } match (&cache.h265_vps, &cache.h265_sps, &cache.h265_pps) { @@ -591,7 +614,7 @@ impl SharedVideoPipeline { debug!("[Pipeline] Prepended cached VPS/SPS/PPS to H265 IRAP"); (Bytes::from(output), true) } - _ => (data, true), + _ => (data, false), } } _ => (data, ffmpeg_keyframe), @@ -599,6 +622,10 @@ impl SharedVideoPipeline { } fn broadcast_encoded(&self, frame: Arc) { + if frame.is_keyframe { + *self.bootstrap_frame.write() = Some(frame.clone()); + } + let subscribers = { let guard = self.subscribers.read(); if guard.is_empty() { @@ -638,6 +665,9 @@ impl SharedVideoPipeline { return Ok(()); } + *self.parameter_sets.lock() = CachedH26xParameterSets::default(); + *self.bootstrap_frame.write() = None; + let mut config = self.config.read().await.clone(); let parallel_mjpeg_decode = should_parallel_decode_mjpeg(&config); { @@ -1680,6 +1710,128 @@ mod tests { assert_eq!(h265.output_codec, VideoEncoderType::H265); } + #[test] + fn h264_keyframe_requires_idr_and_parameter_sets() { + let pipeline = SharedVideoPipeline::new(SharedVideoPipelineConfig::h264( + Resolution::HD720, + BitratePreset::Balanced, + )) + .unwrap(); + + let predicted = Bytes::from_static(&[0, 0, 0, 1, 0x41, 0xc0]); + let (_, key) = + pipeline.inspect_and_parameterize_packet(VideoEncoderType::H264, predicted, true); + assert!( + !key, + "a driver flag must not turn a P-frame into a keyframe" + ); + + let parameter_sets = + Bytes::from_static(&[0, 0, 0, 1, 0x67, 0x42, 0x40, 0x1f, 0, 0, 0, 1, 0x68, 0xce]); + let (_, key) = + pipeline.inspect_and_parameterize_packet(VideoEncoderType::H264, parameter_sets, false); + assert!(!key, "parameter sets alone are not a keyframe"); + + let idr = Bytes::from_static(&[0, 0, 0, 1, 0x65, 0x88]); + let (_, key) = pipeline.inspect_and_parameterize_packet(VideoEncoderType::H264, idr, false); + assert!(!key, "an IDR without a driver key flag is not trusted"); + + let idr = Bytes::from_static(&[0, 0, 0, 1, 0x65, 0x88]); + let (bootstrap, key) = + pipeline.inspect_and_parameterize_packet(VideoEncoderType::H264, idr, true); + assert!( + key, + "matching driver metadata and IDR payload should bootstrap" + ); + assert!(h264_bitstream::has_sps_pps(bootstrap.as_ref())); + assert!(h264_bitstream::is_keyframe(bootstrap.as_ref())); + } + + #[test] + fn h265_keyframe_requires_irap_and_parameter_sets() { + let pipeline = SharedVideoPipeline::new(SharedVideoPipelineConfig::h265( + Resolution::HD720, + BitratePreset::Balanced, + )) + .unwrap(); + + let trail = Bytes::from_static(&[0, 0, 0, 1, 1 << 1, 1, 0xaa]); + let (_, key) = + pipeline.inspect_and_parameterize_packet(VideoEncoderType::H265, trail, true); + assert!( + !key, + "a driver flag must not turn a trailing frame into a keyframe" + ); + + let parameter_sets = Bytes::from_static(&[ + 0, + 0, + 0, + 1, + 32 << 1, + 1, + 0xaa, + 0, + 0, + 0, + 1, + 33 << 1, + 1, + 0xbb, + 0, + 0, + 0, + 1, + 34 << 1, + 1, + 0xcc, + ]); + let (_, key) = + pipeline.inspect_and_parameterize_packet(VideoEncoderType::H265, parameter_sets, false); + assert!(!key, "parameter sets alone are not a keyframe"); + + let irap = Bytes::from_static(&[0, 0, 0, 1, 19 << 1, 1, 0xdd]); + let (_, key) = + pipeline.inspect_and_parameterize_packet(VideoEncoderType::H265, irap, false); + assert!(!key, "an IRAP without a driver key flag is not trusted"); + + let irap = Bytes::from_static(&[0, 0, 0, 1, 19 << 1, 1, 0xdd]); + let (bootstrap, key) = + pipeline.inspect_and_parameterize_packet(VideoEncoderType::H265, irap, true); + assert!( + key, + "matching driver metadata and IRAP payload should bootstrap" + ); + assert!(h265_bitstream::has_vps_sps_pps(bootstrap.as_ref())); + assert!(h265_bitstream::is_keyframe(bootstrap.as_ref())); + } + + #[test] + fn new_subscriber_receives_cached_bootstrap_frame() { + let pipeline = SharedVideoPipeline::new(SharedVideoPipelineConfig::h264( + Resolution::HD720, + BitratePreset::Balanced, + )) + .unwrap(); + let bootstrap = Arc::new(EncodedVideoFrame { + data: Bytes::from_static(&[ + 0, 0, 0, 1, 0x67, 0x42, 0x40, 0x1f, 0, 0, 0, 1, 0x68, 0xce, 0, 0, 0, 1, 0x65, 0x88, + ]), + pts_ms: 0, + is_keyframe: true, + sequence: 1, + duration: Duration::from_millis(33), + codec: VideoEncoderType::H264, + }); + + pipeline.broadcast_encoded(bootstrap.clone()); + let mut subscriber = pipeline.subscribe(); + let received = subscriber + .try_recv() + .expect("cached bootstrap frame should seed the subscriber queue"); + assert!(Arc::ptr_eq(&received, &bootstrap)); + } + #[test] fn stop_request_does_not_publish_worker_exit() { let pipeline = SharedVideoPipeline::new(SharedVideoPipelineConfig::h264(