feat: add AV1 encryptor

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 {

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")
       }
     };