GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2026-09-12 18:44:25 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix(video): 统一关键帧参数集处理并缓存启动帧

Browse Source 
在共享管线中将长度前缀 H264 转为 Annex-B，要求编码器标记与 IDR/IRAP 内容一致，并补齐参数集后才标记可独立解码的关键帧。

向新订阅者提供缓存启动帧，启动管线时清理缓存，并移除 RustDesk 会话内重复的 SPS/PPS 处理。补充格式归一化、关键帧判定和订阅测试。
This commit is contained in:
mofeng-git
2026-09-05 14:10:53 +08:00
parent a8ee9552a0
commit 79a4dcf2b0
3 changed files with 222 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
src/rustdesk/frame_adapters.rs
View File

@@ -31,8 +31,6 @@ pub struct VideoFrameAdapter {
codec: VideoCodec,
seq: u32,
timestamp_base: u64,
h264_sps: Option<Bytes>,
h264_pps: Option<Bytes>,
}
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)
}

66
src/video/codec/h264_bitstream.rs
View File

@@ -252,13 +252,37 @@ pub fn avcc_to_annex_b(data: &[u8]) -> Option<Vec<u8>> {
}
}
pub fn normalize_for_webrtc(data: &[u8]) -> Vec<u8> {
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<u8> {
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));
}
}

170
src/video/pipeline/shared.rs
View File

@@ -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<CachedH26xParameterSets>,
/// Most recent random-access frame with all decoder parameter sets.
bootstrap_frame: ParkingRwLock<Option<Arc<EncodedVideoFrame>>>,
/// 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<EncodedVideoFrame>) {
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(