diff --git a/davey/src/cryptor/codec_utils.rs b/davey/src/cryptor/codec_utils.rs index 86da879..3d7f899 100644 --- a/davey/src/cryptor/codec_utils.rs +++ b/davey/src/cryptor/codec_utils.rs @@ -1,6 +1,12 @@ +use core::matches; + use tracing::warn; -use super::{Codec, frame_processors::OutboundFrameProcessor}; +use super::{ + Codec, + frame_processors::OutboundFrameProcessor, + leb128::{leb128_size, read_leb128, write_leb128}, +}; const NALU_SHORT_START_SEQUENCE_SIZE: usize = 3; @@ -273,6 +279,101 @@ pub fn process_frame_vp9(processor: &mut OutboundFrameProcessor, frame: &[u8]) - true } +#[allow( + clippy::unusual_byte_groupings, + reason = "following upstream libdave byte groupings" +)] +pub fn process_frame_av1(processor: &mut OutboundFrameProcessor, frame: &[u8]) -> bool { + const OBU_HEADER_HAS_EXTENSION_MASK: u8 = 0b0_0000_100; + const OBU_HEADER_HAS_SIZE_MASK: u8 = 0b0_0000_010; + const OBU_HEADER_TYPE_MASK: u8 = 0b0_1111_000; + const OBU_TYPE_TEMPORAL_DELIMITER: u8 = 2; + const OBU_TYPE_TILE_LIST: u8 = 8; + const OBU_TYPE_PADDING: u8 = 15; + const OBU_EXTENSION_SIZE_BYTES: usize = 1; + + let mut i = 0; + while i < frame.len() { + // Read the OBU header. + let obu_header_index = i; + let mut obu_header = frame[i]; + i += size_of_val(&obu_header); + + let obu_has_extension: bool = (obu_header & OBU_HEADER_HAS_EXTENSION_MASK) != 0; + let obu_has_size: bool = (obu_header & OBU_HEADER_HAS_SIZE_MASK) != 0; + let obu_type: u8 = (obu_header & OBU_HEADER_TYPE_MASK) >> 3; + + if obu_has_extension { + // Skip extension byte + i += OBU_EXTENSION_SIZE_BYTES; + } + + if i >= frame.len() { + warn!("Malformed AV1 frame: header overflows frame"); + return false; + } + + let obu_payload_size: usize; + if obu_has_size { + let Some((obu_payload_size_explicit, leb128_size)) = read_leb128(&frame[i..]) else { + warn!("Malformed AV1 frame: invalid LEB128 size"); + return false; + }; + obu_payload_size = obu_payload_size_explicit as usize; + i += leb128_size; + } else { + // If the size is not present, the OBU extends to the end of the frame. + obu_payload_size = frame.len() - i; + } + + let obu_payload_index = i; + + if i + obu_payload_size > frame.len() { + warn!("Malformed AV1 frame: payload overflows frame"); + return false; + } + + i += obu_payload_size; + + // We only copy the OBUs that will not get dropped by the packetizer + if matches!( + obu_type, + OBU_TYPE_TEMPORAL_DELIMITER | OBU_TYPE_TILE_LIST | OBU_TYPE_PADDING + ) { + continue; + } + + // if this is the last OBU, we may need to flip the "has size" bit + // which allows us to append necessary protocol data to the frame + let rewritten_without_size = i == frame.len() && obu_has_size; + if rewritten_without_size { + obu_header &= !OBU_HEADER_HAS_SIZE_MASK; + } + + // write the OBU header unencrypted + processor.add_unencrypted_bytes(&[obu_header]); + if obu_has_extension { + // write the extension byte unencrypted + processor.add_unencrypted_bytes( + &frame[(obu_header_index + size_of_val(&obu_header))..][..OBU_EXTENSION_SIZE_BYTES], + ); + } + + // write the OBU payload size unencrypted if it was present and we didn't rewrite + // without it + if obu_has_size && !rewritten_without_size { + let mut leb128_buffer = vec![0u8; leb128_size(obu_payload_size as u64)]; + write_leb128(obu_payload_size as u64, leb128_buffer.as_mut_slice()); + processor.add_unencrypted_bytes(&leb128_buffer); + } + + // add the OBU payload, encrypted + processor.add_encrypted_bytes(&frame[obu_payload_index..][..obu_payload_size]); + } + + true +} + pub fn validate_encrypted_frame(processor: &OutboundFrameProcessor, frame: &[u8]) -> bool { let codec = &processor.frame_codec; if *codec != Codec::H264 && *codec != Codec::H265 { diff --git a/davey/src/cryptor/frame_processors.rs b/davey/src/cryptor/frame_processors.rs index d1c0775..b9c2e0f 100644 --- a/davey/src/cryptor/frame_processors.rs +++ b/davey/src/cryptor/frame_processors.rs @@ -349,9 +349,10 @@ impl OutboundFrameProcessor { Codec::H265 => process_frame_h265(self, frame), Codec::VP8 => process_frame_vp8(self, frame), Codec::VP9 => process_frame_vp9(self, frame), + Codec::AV1 => process_frame_av1(self, frame), _ => { // TODO we dont need to but maybe add more codecs later - unimplemented!("h264, h265, vp8, vp9 and opus are the only supported codecs currently") + unimplemented!("h264, h265, vp8, vp9, av1 and opus are the only supported codecs currently") } };