add session encryption

src/udp_source.rs

@@ -1,14 +1,16 @@
 use crate::circular_buffer_source::CircularBufferSource;
-use audiopus::coder::Decoder;
-use audiopus::{Channels, MutSignals, SampleRate};
-use log::{debug, error, info};
+use aws_lc_rs::aead::{AES_256_GCM, Aad, NONCE_LEN, Nonce, RandomizedNonceKey};
+use aws_lc_rs::error::Unspecified;
+use log::{debug, error, info, trace};
 use prost::{DecodeError, Message};
 use std::io::Write;
 use tap_interface::tap::message::TapMessage;
 use thiserror::Error;
 use tokio::net::UdpSocket;
+use tokio::select;
+use tokio_util::sync::CancellationToken;
 
-pub const SAMPLE_RATE: SampleRate = SampleRate::Hz48000;
+pub const SAMPLE_RATE: u32 = 48_000;
 
 #[derive(Debug, Error)]
 #[allow(clippy::enum_variant_names)]
@@ -16,7 +18,7 @@ pub enum Error {
     #[error("IO Error {0}")]
     IOError(#[from] std::io::Error),
     #[error("Failed to decode audio in packet: {0}")]
-    OpusError(#[from] audiopus::Error),
+    OpusError(#[from] opus::Error),
     #[error("Failed to decode packet: {0}")]
     DecodeError(#[from] DecodeError),
 }
@@ -24,17 +26,19 @@ pub enum Error {
 pub struct UdpSource {
     udp: UdpSocket,
     source: CircularBufferSource,
-    decoder: Decoder,
+    decoder: opus::Decoder,
     expected_seq: u32,
+    key: RandomizedNonceKey,
 }
 
 impl UdpSource {
-    pub fn new(udp_socket: UdpSocket, source: CircularBufferSource) -> Self {
+    pub fn new(udp_socket: UdpSocket, source: CircularBufferSource, key: &[u8]) -> Self {
         Self {
             udp: udp_socket,
             source,
-            decoder: Decoder::new(SAMPLE_RATE, Channels::Stereo).unwrap(),
+            decoder: opus::Decoder::new(SAMPLE_RATE, opus::Channels::Stereo).unwrap(),
             expected_seq: 0,
+            key: RandomizedNonceKey::new(&AES_256_GCM, key).unwrap(),
         }
     }
 
@@ -43,17 +47,15 @@ impl UdpSource {
         input_buffer: &[u8],
         output_buffer: &mut [f32],
     ) -> Result<usize, Error> {
-        let packet = audiopus::packet::Packet::try_from(input_buffer)?;
-
-        let signals = MutSignals::try_from(output_buffer)?;
-
-        let sample_size = self.decoder.decode_float(Some(packet), signals, false)?;
+        let sample_size = self
+            .decoder
+            .decode_float(input_buffer, output_buffer, false)?;
 
         Ok(sample_size * 2)
     }
 
     fn write_audio_data_to_source(&mut self, samples: &[f32]) -> Result<(), Error> {
-        let mut sample_bytes = Vec::with_capacity(std::mem::size_of_val(samples));
+        let mut sample_bytes = Vec::with_capacity(size_of_val(samples));
 
         for sample in samples {
             let bytes = sample.to_le_bytes();
@@ -69,20 +71,46 @@ impl UdpSource {
         self.udp.recv(buffer).await
     }
 
-    pub async fn worker(&mut self) {
+    fn decrypt<'a>(&'a mut self, buffer: &'a mut [u8]) -> Result<&'a mut [u8], Unspecified> {
+        let nonce = &buffer[0..NONCE_LEN];
+        let nonce = Nonce::try_assume_unique_for_key(nonce)?;
+
+        self.key
+            .open_in_place(nonce, Aad::empty(), &mut buffer[NONCE_LEN..])
+    }
+
+    pub async fn worker(&mut self, cancellation_token: CancellationToken) {
         info!("Starting TAP UDP Endpoint");
         let mut buffer = vec![0; 1024 * 32];
         let mut decode_data = vec![0.0; tap_interface::STEREO_FRAME_SIZE];
         loop {
-            let msg_size = match self.recv_message(&mut buffer).await {
-                Ok(msg_size) => msg_size,
-                Err(err) => {
-                    error!("Failed to recv udp packet: {err:?}");
+            let msg_size = select! {
+                res = self.recv_message(&mut buffer) => {
+
+
+                    match res {
+                        Ok(msg_size) => msg_size,
+                        Err(err) => {
+                            error!("Failed to recv udp packet: {err:?}");
+                            continue;
+                        }
+                    }
+                }
+                _ = cancellation_token.cancelled() => {
+                    info!("Closing TAP UDP Endpoint...");
+                    return;
+                }
+            };
+
+            let decrypted_msg = match self.decrypt(&mut buffer[0..msg_size]) {
+                Ok(msg) => msg,
+                Err(_err) => {
+                    trace!("Failed to decrypt message");
                     continue;
                 }
             };
 
-            let tap_message = match TapMessage::decode(&buffer[..msg_size]) {
+            let tap_message = match TapMessage::decode(&*decrypted_msg) {
                 Ok(tap_msg) => tap_msg,
                 Err(err) => {
                     error!("Failed decode proto message: {err:?}");