Files
simple-instant-stream/src/media.rs
T
2026-06-14 01:06:25 +01:00

148 lines
4.1 KiB
Rust

// Claude slop... im not skilled amount to do this bullshit
use bytes::Bytes;
pub struct VideoFrame {
pub data: Bytes,
pub is_keyframe: bool,
pub timestamp_ms: u32,
}
pub struct AudioFrame {
pub data: Bytes,
pub timestamp_ms: u32,
}
pub struct H264Parser {
sps: Option<Vec<u8>>,
pps: Option<Vec<u8>>,
}
impl H264Parser {
pub fn new() -> Self {
Self {
sps: None,
pps: None,
}
}
/// Parse an RTMP VideoDataReceived payload. Returns None for sequence
/// header packets (which carry SPS/PPS but no displayable frame).
pub fn parse(&mut self, bytes: &[u8], timestamp_ms: u32) -> Option<VideoFrame> {
if bytes.len() < 5 {
return None;
}
let frame_type = (bytes[0] >> 4) & 0x0F;
let codec_id = bytes[0] & 0x0F;
if codec_id != 7 {
return None; // not H.264
}
let avc_packet_type = bytes[1];
// bytes[2..5] are the composition time offset — not needed for sending
let payload = &bytes[5..];
match avc_packet_type {
0 => {
self.parse_sequence_header(payload);
None
}
1 => {
let is_keyframe = frame_type == 1;
let data = self.avcc_to_annexb(payload, is_keyframe)?;
Some(VideoFrame {
data: Bytes::from(data),
is_keyframe,
timestamp_ms,
})
}
_ => None,
}
}
fn parse_sequence_header(&mut self, payload: &[u8]) {
// AVCDecoderConfigurationRecord layout:
// [0] configurationVersion
// [1] AVCProfileIndication
// [2] profile_compatibility
// [3] AVCLevelIndication
// [4] 0xFF (lower 2 bits = lengthSizeMinusOne, always 3 meaning 4-byte lengths)
// [5] 0xE0 | numSPS
// [6..] SPS entries: 2-byte length + bytes
// then: numPPS, PPS entries: 2-byte length + bytes
if payload.len() < 7 {
return;
}
let mut i = 5;
let num_sps = (payload[i] & 0x1F) as usize;
i += 1;
for _ in 0..num_sps {
if i + 2 > payload.len() {
return;
}
let len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize;
i += 2;
if i + len > payload.len() {
return;
}
self.sps = Some(payload[i..i + len].to_vec());
i += len;
}
if i >= payload.len() {
return;
}
let num_pps = payload[i] as usize;
i += 1;
for _ in 0..num_pps {
if i + 2 > payload.len() {
return;
}
let len = u16::from_be_bytes([payload[i], payload[i + 1]]) as usize;
i += 2;
if i + len > payload.len() {
return;
}
self.pps = Some(payload[i..i + len].to_vec());
i += len;
}
}
fn avcc_to_annexb(&self, payload: &[u8], is_keyframe: bool) -> Option<Vec<u8>> {
let mut out = Vec::new();
// Prepend SPS+PPS before every keyframe so str0m's packetizer
// can bundle them into a STAP-A alongside the IDR NALU.
if is_keyframe {
if let (Some(sps), Some(pps)) = (&self.sps, &self.pps) {
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);
}
}
// Convert each length-prefixed NALU to an Annex B start-code NALU.
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) }
}
}