fix(webrtc): preserve H.26x headers for new sessions

Detect random-access frames from Annex-B data and cache codec headers.
Prepend missing headers to later H.264 IDR and H.265 IRAP frames.
Forward standalone headers while WebRTC waits for a decodable frame.

Signed-off-by: BigfootACA <bigfoot@radxa.com>

src/video/codec/h265_bitstream.rs

@@ -0,0 +1,112 @@
+const VPS_NAL_TYPE: u8 = 32;
+const SPS_NAL_TYPE: u8 = 33;
+const PPS_NAL_TYPE: u8 = 34;
+
+fn find_start_code(data: &[u8], from: usize) -> Option<(usize, usize)> {
+    let mut offset = from;
+    while offset + 3 <= data.len() {
+        if offset + 4 <= data.len() && data[offset..offset + 4] == [0, 0, 0, 1] {
+            return Some((offset, 4));
+        }
+        if data[offset..offset + 3] == [0, 0, 1] {
+            return Some((offset, 3));
+        }
+        offset += 1;
+    }
+    None
+}
+
+fn for_each_nal(data: &[u8], mut visit: impl FnMut(u8, &[u8])) {
+    let mut cursor = 0;
+    while let Some((start, start_code_len)) = find_start_code(data, cursor) {
+        let nal_start = start + start_code_len;
+        if nal_start + 2 > data.len() {
+            break;
+        }
+        let next_start = find_start_code(data, nal_start)
+            .map(|(offset, _)| offset)
+            .unwrap_or(data.len());
+        let mut nal_end = next_start;
+        while nal_end > nal_start && data[nal_end - 1] == 0 {
+            nal_end -= 1;
+        }
+        if nal_end >= nal_start + 2 {
+            visit((data[nal_start] >> 1) & 0x3f, &data[nal_start..nal_end]);
+        }
+        if next_start == data.len() {
+            break;
+        }
+        cursor = next_start;
+    }
+}
+
+pub fn is_keyframe(data: &[u8]) -> bool {
+    let mut keyframe = false;
+    for_each_nal(data, |nal_type, _| {
+        if (16..=23).contains(&nal_type) {
+            keyframe = true;
+        }
+    });
+    keyframe
+}
+
+pub fn extract_vps_sps_pps(data: &[u8]) -> (Option<Vec<u8>>, Option<Vec<u8>>, Option<Vec<u8>>) {
+    let mut vps = None;
+    let mut sps = None;
+    let mut pps = None;
+    for_each_nal(data, |nal_type, nal| match nal_type {
+        VPS_NAL_TYPE => vps = Some(nal.to_vec()),
+        SPS_NAL_TYPE => sps = Some(nal.to_vec()),
+        PPS_NAL_TYPE => pps = Some(nal.to_vec()),
+        _ => {}
+    });
+    (vps, sps, pps)
+}
+
+pub fn has_vps_sps_pps(data: &[u8]) -> bool {
+    let (vps, sps, pps) = extract_vps_sps_pps(data);
+    vps.is_some() && sps.is_some() && pps.is_some()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn detects_irap_but_not_trail_frame() {
+        assert!(is_keyframe(&[0, 0, 0, 1, 19 << 1, 1, 0xaa]));
+        assert!(is_keyframe(&[0, 0, 1, 21 << 1, 1, 0xbb]));
+        assert!(!is_keyframe(&[0, 0, 0, 1, 1 << 1, 1, 0xcc]));
+    }
+
+    #[test]
+    fn extracts_parameter_sets() {
+        let data = [
+            0,
+            0,
+            0,
+            1,
+            32 << 1,
+            1,
+            0xaa,
+            0,
+            0,
+            1,
+            33 << 1,
+            1,
+            0xbb,
+            0,
+            0,
+            0,
+            1,
+            34 << 1,
+            1,
+            0xcc,
+        ];
+        let (vps, sps, pps) = extract_vps_sps_pps(&data);
+        assert_eq!(vps.unwrap(), [32 << 1, 1, 0xaa]);
+        assert_eq!(sps.unwrap(), [33 << 1, 1, 0xbb]);
+        assert_eq!(pps.unwrap(), [34 << 1, 1, 0xcc]);
+        assert!(has_vps_sps_pps(&data));
+    }
+}

src/video/codec/mod.rs

@@ -8,6 +8,7 @@ pub mod convert;
 pub mod h264;
 pub mod h264_bitstream;
 pub mod h265;
+pub mod h265_bitstream;
 pub mod jpeg;
 pub mod registry;
 pub mod self_check;

src/video/pipeline/shared.rs

@@ -48,9 +48,9 @@ use crate::video::capture::status::{
     signal_status_from_capture_kind, CaptureIoErrorKind,
 };
 use crate::video::capture::{BridgeContext, CaptureReadError, CaptureStream};
-use crate::video::codec::h264_bitstream;
 use crate::video::codec::registry::{EncoderBackend, VideoEncoderType};
 use crate::video::codec::MjpegToNv12Decoder;
+use crate::video::codec::{h264_bitstream, h265_bitstream};
 use crate::video::device::parse_bridge_kind;
 use crate::video::device::VideoControlMode;
 use crate::video::format::{PixelFormat, Resolution};
@@ -276,6 +276,15 @@ pub struct SharedVideoPipelineStats {
     pub current_fps: f32,
 }
 
+#[derive(Default)]
+struct CachedH26xParameterSets {
+    h264_sps: Option<Vec<u8>>,
+    h264_pps: Option<Vec<u8>>,
+    h265_vps: Option<Vec<u8>>,
+    h265_sps: Option<Vec<u8>>,
+    h265_pps: Option<Vec<u8>>,
+}
+
 /// Universal shared video pipeline
 pub struct SharedVideoPipeline {
     config: RwLock<SharedVideoPipelineConfig>,
@@ -296,6 +305,7 @@ pub struct SharedVideoPipeline {
     sequence: AtomicU64,
     /// Atomic flag for keyframe request (avoids lock contention)
     keyframe_requested: AtomicBool,
+    parameter_sets: ParkingMutex<CachedH26xParameterSets>,
     /// 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,
@@ -335,6 +345,7 @@ impl SharedVideoPipeline {
             running_flag: AtomicBool::new(false),
             sequence: AtomicU64::new(0),
             keyframe_requested: AtomicBool::new(false),
+            parameter_sets: ParkingMutex::new(CachedH26xParameterSets::default()),
             pipeline_start_time_us: AtomicI64::new(0),
             pending_sync_geometry: ParkingMutex::new(None),
             device_lost_reason: ParkingMutex::new(None),
@@ -507,6 +518,86 @@ impl SharedVideoPipeline {
         let _ = self.h264_profile_level_id.send(Some(profile_level_id));
     }
 
+    fn inspect_and_parameterize_packet(
+        &self,
+        codec: VideoEncoderType,
+        data: Bytes,
+        ffmpeg_keyframe: bool,
+    ) -> (Bytes, bool) {
+        match codec {
+            VideoEncoderType::H264 => {
+                let (sps, pps) = h264_bitstream::extract_sps_pps(data.as_ref());
+                let is_keyframe = 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());
+                }
+                if let Some(pps) = pps.as_ref() {
+                    cache.h264_pps = Some(pps.clone());
+                }
+
+                if !is_keyframe || (sps.is_some() && pps.is_some()) {
+                    return (data, is_keyframe);
+                }
+
+                match (&cache.h264_sps, &cache.h264_pps) {
+                    (Some(cached_sps), Some(cached_pps)) => {
+                        let mut output = Vec::with_capacity(
+                            data.len() + cached_sps.len() + cached_pps.len() + 8,
+                        );
+                        output.extend_from_slice(&[0, 0, 0, 1]);
+                        output.extend_from_slice(cached_sps);
+                        output.extend_from_slice(&[0, 0, 0, 1]);
+                        output.extend_from_slice(cached_pps);
+                        output.extend_from_slice(data.as_ref());
+                        debug!("[Pipeline] Prepended cached SPS/PPS to H264 IDR");
+                        (Bytes::from(output), true)
+                    }
+                    _ => (data, true),
+                }
+            }
+            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 mut cache = self.parameter_sets.lock();
+                if let Some(vps) = vps.as_ref() {
+                    cache.h265_vps = Some(vps.clone());
+                }
+                if let Some(sps) = sps.as_ref() {
+                    cache.h265_sps = Some(sps.clone());
+                }
+                if let Some(pps) = pps.as_ref() {
+                    cache.h265_pps = Some(pps.clone());
+                }
+
+                if !is_keyframe || (vps.is_some() && sps.is_some() && pps.is_some()) {
+                    return (data, is_keyframe);
+                }
+
+                match (&cache.h265_vps, &cache.h265_sps, &cache.h265_pps) {
+                    (Some(cached_vps), Some(cached_sps), Some(cached_pps)) => {
+                        let mut output = Vec::with_capacity(
+                            data.len()
+                                + cached_vps.len()
+                                + cached_sps.len()
+                                + cached_pps.len()
+                                + 12,
+                        );
+                        for parameter_set in [cached_vps, cached_sps, cached_pps] {
+                            output.extend_from_slice(&[0, 0, 0, 1]);
+                            output.extend_from_slice(parameter_set);
+                        }
+                        output.extend_from_slice(data.as_ref());
+                        debug!("[Pipeline] Prepended cached VPS/SPS/PPS to H265 IRAP");
+                        (Bytes::from(output), true)
+                    }
+                    _ => (data, true),
+                }
+            }
+            _ => (data, ffmpeg_keyframe),
+        }
+    }
+
     fn broadcast_encoded(&self, frame: Arc<EncodedVideoFrame>) {
         let subscribers = {
             let guard = self.subscribers.read();
@@ -1193,9 +1284,11 @@ impl SharedVideoPipeline {
             })?;
 
             if let Some((data, is_keyframe)) = packet {
+                let (data, is_keyframe) =
+                    self.inspect_and_parameterize_packet(codec, Bytes::from(data), is_keyframe);
                 let sequence = self.sequence.fetch_add(1, Ordering::Relaxed) + 1;
                 return Ok(vec![EncodedVideoFrame {
-                    data: Bytes::from(data),
+                    data,
                     pts_ms,
                     is_keyframe,
                     sequence,
@@ -1275,14 +1368,15 @@ impl SharedVideoPipeline {
 
                 let mut encoded_frames = Vec::with_capacity(frames.len());
                 for encoded in frames {
-                    let is_keyframe = encoded.key == 1;
+                    let (data, is_keyframe) =
+                        self.inspect_and_parameterize_packet(codec, encoded.data, encoded.key == 1);
                     let sequence = self.sequence.fetch_add(1, Ordering::Relaxed) + 1;
                     if codec == VideoEncoderType::H264 {
-                        self.update_h264_profile_level_id(&encoded.data);
+                        self.update_h264_profile_level_id(&data);
                     }
 
                     encoded_frames.push(EncodedVideoFrame {
-                        data: encoded.data,
+                        data,
                         pts_ms,
                         is_keyframe,
                         sequence,

src/webrtc/universal_session.rs

@@ -36,7 +36,7 @@ use crate::audio::OpusFrame;
 use crate::error::{AppError, Result};
 use crate::hid::datachannel::{parse_hid_message, HidChannelEvent};
 use crate::hid::HidController;
-use crate::video::codec::h264_bitstream;
+use crate::video::codec::{h264_bitstream, h265_bitstream};
 use crate::video::types::{
     BitratePreset, EncodedVideoFrame, PixelFormat, Resolution, VideoEncoderType,
 };
@@ -638,11 +638,16 @@ impl UniversalSession {
                                     next_keyframe_retry = Instant::now();
                                 }
 
-                                // Some H264 encoders output SPS/PPS in a separate non-keyframe AU
-                                // before IDR. Keep this frame so browser can decode the next IDR.
-                                let forward_h264_parameter_frame = waiting_for_keyframe
-                                    && expected_codec == VideoEncoderType::H264
-                                    && h264_bitstream::has_sps_pps(encoded_frame.data.as_ref());
+                                let forward_parameter_frame = waiting_for_keyframe
+                                    && match expected_codec {
+                                        VideoEncoderType::H264 => h264_bitstream::has_sps_pps(
+                                            encoded_frame.data.as_ref(),
+                                        ),
+                                        VideoEncoderType::H265 => h265_bitstream::has_vps_sps_pps(
+                                            encoded_frame.data.as_ref(),
+                                        ),
+                                        _ => false,
+                                    };
 
                                 let now = Instant::now();
                                 if keyframe_requests < KEYFRAME_RETRY_LIMIT
@@ -659,7 +664,7 @@ impl UniversalSession {
                                         );
                                     }
                                 }
-                                if !forward_h264_parameter_frame {
+                                if !forward_parameter_frame {
                                     continue;
                                 }
                             }