simple-instant-stream/crates/server/src/codec/h265.rs

use bytes::Bytes;

use crate::codec::{CodecParser, VideoFrame};

pub struct H265CodecParser {
    vps: Option<Vec<u8>>,
    sps: Option<Vec<u8>>,
    pps: Option<Vec<u8>>,
}

impl H265CodecParser {
    pub fn new() -> Self {
        Self {
            vps: None,
            sps: None,
            pps: None,
        }
    }

    // HEVCDecoderConfigurationRecord layout (ISO 14496-15 §8.3.3):
    // [0]     configurationVersion (always 1)
    // [1..2]  general_profile_space(2b) | general_tier_flag(1b) | general_profile_idc(5b)
    // [2..6]  general_profile_compatibility_flags
    // [6..12] general_constraint_indicator_flags
    // [12]    general_level_idc
    // [13..15] min_spatial_segmentation_idc (lower 12 bits)
    // [15]    parallelismType (lower 2 bits)
    // [16]    chroma_format_idc (lower 2 bits)
    // [17]    bit_depth_luma_minus8 (lower 3 bits)
    // [18]    bit_depth_chroma_minus8 (lower 3 bits)
    // [19..21] avgFrameRate
    // [21]    constantFrameRate(2b) | numTemporalLayers(3b) | temporalIdNested(1b) | lengthSizeMinusOne(2b)
    // [22]    numOfArrays
    // [23..]  arrays: [ array_completeness(1b) | reserved(1b) | NAL_unit_type(6b), numNalus(2b),
    //                   [ naluLength(2b), nalu(naluLength) ] ]
    fn parse_sequence_header(&mut self, payload: &[u8]) {
        if payload.len() < 23 {
            return;
        }

        let num_arrays = payload[22] as usize;
        let mut i = 23;

        for _ in 0..num_arrays {
            if i + 3 > payload.len() {
                return;
            }
            let nal_type = payload[i] & 0x3F;
            let num_nalus = u16::from_be_bytes([payload[i + 1], payload[i + 2]]) as usize;
            i += 3;

            for _ in 0..num_nalus {
                if i + 2 > payload.len() {
                    return;
                }
                let nalu_len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize;
                i += 2;
                if i + nalu_len > payload.len() {
                    return;
                }

                match nal_type {
                    32 => self.vps = Some(payload[i..i + nalu_len].to_vec()),
                    33 => self.sps = Some(payload[i..i + nalu_len].to_vec()),
                    34 => self.pps = Some(payload[i..i + nalu_len].to_vec()),
                    _ => {}
                }

                i += nalu_len;
            }
        }
    }

    fn hvcc_to_annexb(&self, payload: &[u8], is_keyframe: bool) -> Option<Vec<u8>> {
        let mut out = Vec::with_capacity(payload.len());

        if is_keyframe {
            if let (Some(vps), Some(sps), Some(pps)) = (&self.vps, &self.sps, &self.pps) {
                out.extend_from_slice(&[0, 0, 0, 1]);
                out.extend_from_slice(vps);
                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);
            }
        }

        let mut i = 0;
        while i + 4 <= payload.len() {
            let nalu_len = u32::from_be_bytes(payload[i..i + 4].try_into().unwrap()) as usize;
            i += 4;
            if i + nalu_len > payload.len() {
                break;
            }
            out.extend_from_slice(&[0, 0, 0, 1]);
            out.extend_from_slice(&payload[i..i + nalu_len]);
            i += nalu_len;
        }

        if out.is_empty() { None } else { Some(out) }
    }
}

impl CodecParser for H265CodecParser {
    fn parse(&mut self, data: &[u8], timestamp_ms: u32) -> Option<VideoFrame> {
        // H.265 only comes via enhanced RTMP (bit 7 set, FourCC "hvc1")
        if data.len() < 5 || data[0] & 0x80 == 0 {
            return None;
        }

        let is_keyframe = (data[0] >> 4) & 0x07 == 1;
        let packet_type = data[0] & 0x0F;
        // bytes 1-4 are FourCC "hvc1" — already validated by rtmp.rs

        match packet_type {
            0 => {
                // SequenceStart: bytes[5..] = HEVCDecoderConfigurationRecord
                self.parse_sequence_header(data.get(5..)?);
                None
            }
            1 => {
                // CodedFrames: bytes 5-7 = SI24 composition time offset, bytes 8+ = HVCC.
                // RTP timestamps must be presentation time (RFC 7798), so emit
                // PTS = DTS + CTS. Encoders use this packet type exactly when CTS != 0
                // (B-frames present); stamping DTS instead makes playback stutter.
                if data.len() < 8 {
                    return None;
                }
                let cts = i32::from_be_bytes([0, data[5], data[6], data[7]]) << 8 >> 8;
                let pts_ms = (timestamp_ms as i64 + cts as i64).max(0) as u32;
                let annexb = self.hvcc_to_annexb(data.get(8..)?, is_keyframe)?;
                Some(VideoFrame { data: Bytes::from(annexb), is_keyframe, timestamp_ms: pts_ms })
            }
            3 => {
                // CodedFramesX: no CTS, bytes 5+ = HVCC
                let annexb = self.hvcc_to_annexb(data.get(5..)?, is_keyframe)?;
                Some(VideoFrame { data: Bytes::from(annexb), is_keyframe, timestamp_ms })
            }
            _ => None,
        }
    }
}