add session encryption

Cargo.toml

@@ -1,24 +1,27 @@
 [package]
 name = "denny-jack"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2024"
 
 [dependencies]
 clap = { version = "4.5.45", features = ["derive"] }
 config = "0.15.14"
-poise = "0.6.1"
+poise = "0.7.1"
 serde = "1.0.219"
-songbird = { version = "0.6.0", features = ["builtin-queue", "driver"] }
+songbird = { version = "0.6.1", features = ["builtin-queue", "driver"] }
 thiserror = "2.0.15"
 tokio = { version = "1.47.1", features = ["macros", "rt-multi-thread"] }
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
-error = "0.1.9"
 log = "0.4.27"
-audiopus = "0.3.0-rc.0"
-circular-buffer = "1.1.0"
+opus= { version="0.4.0" }
+opusic-sys = { version = "0.7.5", features = ["bundled", "no-stack-protector", "no-hardening", "no-fortify-source"] }
+circular-buffer = "2.0.1"
 tap-interface = {version = "0.1.0", registry = "ahines"}
 prost = "0.14.1"
+aws-lc-rs = "1.18.2"
+hex = "0.4.3"
+tokio-util = "0.7.20"
 
 [dependencies.symphonia]
-version = "0.5"
+version = "0.5.4"
 features = ["aac", "mp3", "isomp4", "alac"] # ...as well as any extras you need!

src/circular_buffer_source.rs

@@ -1,16 +1,16 @@
-use circular_buffer::CircularBuffer;
+use circular_buffer::FixedCircularBuffer;
 use songbird::input::core::io::MediaSource;
 use std::io::SeekFrom;
 use std::{
     io::{Read, Seek, Write},
-    sync::{Arc, Condvar, Mutex},
+    sync::{Arc, Mutex},
 };
 
 const BUFFER_SIZE: usize = 64 * 1024;
 
 #[derive(Clone, Default)]
 pub struct CircularBufferSource {
-    circular_buffer: Arc<Mutex<CircularBuffer<BUFFER_SIZE, u8>>>,
+    circular_buffer: Arc<Mutex<FixedCircularBuffer<u8, BUFFER_SIZE>>>,
 }
 
 impl Read for CircularBufferSource {

src/main.rs

@@ -7,14 +7,18 @@ use crate::circular_buffer_source::CircularBufferSource;
 use crate::dj_config::{Args, DJConfig};
 use crate::stream_task::{StreamControl, stream_task};
 use crate::udp_source::UdpSource;
+use aws_lc_rs::cipher::AES_256_KEY_LEN;
+use aws_lc_rs::rand;
 use clap::Parser;
 use log::{error, info};
+use poise::serenity_prelude::CreateMessage;
 use poise::{PrefixFrameworkOptions, serenity_prelude as serenity};
 use songbird::SerenityInit;
 use thiserror::Error;
 use tokio::net::UdpSocket;
 use tokio::sync::Mutex;
 use tokio::sync::mpsc::{Sender, channel};
+use tokio_util::sync::CancellationToken;
 use tracing_subscriber::EnvFilter;
 use tracing_subscriber::filter::LevelFilter;
 
@@ -26,12 +30,17 @@ pub enum Error {
     FailedToJoin(#[from] songbird::error::JoinError),
 }
 
+struct DennyJackControl {
+    pub udp_task: tokio::task::JoinHandle<()>,
+    pub stream_control_tx: Sender<StreamControl>,
+    pub cancellation_token: CancellationToken,
+}
+
 #[allow(dead_code)]
 struct Data {
     pub config: DJConfig,
-    pub udp_task: tokio::task::JoinHandle<()>,
     pub circular_buffer_source: CircularBufferSource,
-    pub stream_control_tx: Mutex<Option<Sender<StreamControl>>>,
+    pub dj_control: Mutex<Option<DennyJackControl>>,
 }
 type Context<'a> = poise::Context<'a, Data, Error>;
 
@@ -39,8 +48,8 @@ const COMMAND_QUEUE_SIZE: usize = 10;
 
 #[poise::command(slash_command, prefix_command, guild_only)]
 async fn join(ctx: Context<'_>) -> Result<(), Error> {
-    let mut stream_control_tx = match ctx.data().stream_control_tx.try_lock() {
-        Ok(stream_control_tx) => stream_control_tx,
+    let mut dj_control = match ctx.data().dj_control.try_lock() {
+        Ok(dj_control) => dj_control,
         Err(_err) => {
             ctx.reply("Yo yo yo, looks like someone else is streaming on the net.")
                 .await?;
@@ -67,15 +76,43 @@ async fn join(ctx: Context<'_>) -> Result<(), Error> {
         }
     };
 
+    let mut key = vec![0; AES_256_KEY_LEN];
+    rand::fill(&mut key).expect("Failed to generate key");
+
+    ctx.author()
+        .direct_message(
+            &ctx,
+            CreateMessage::new().content(format!(
+                "Yo Yo, your session key is `{}`.\nYou send your sweet tones to me at `{}`",
+                hex::encode(&key),
+                ctx.data().config.server_addr
+            )),
+        )
+        .await?;
+
+    let cancellation_token = CancellationToken::new();
+    let udp_input = UdpSocket::bind(ctx.data().config.server_addr)
+        .await
+        .unwrap();
+    let mut udp_source = UdpSource::new(udp_input, ctx.data().circular_buffer_source.clone(), &key);
+
+    let udp_task_token = cancellation_token.clone();
+    let udp_task = tokio::spawn(async move { udp_source.worker(udp_task_token).await });
+
     let manager = songbird::get(ctx.serenity_context())
         .await
-        .expect("Songbird Voice client placed in at initialisation.")
+        .expect("Songbird Voice client placed in at initialization.")
         .clone();
 
     manager.join(guild_id, connect_to).await?;
 
     let (tx, rx) = channel(COMMAND_QUEUE_SIZE);
-    *stream_control_tx = Some(tx);
+
+    *dj_control = Some(DennyJackControl {
+        udp_task,
+        stream_control_tx: tx,
+        cancellation_token,
+    });
 
     let circular_buffer_source = ctx.data().circular_buffer_source.clone();
     tokio::spawn(async move { stream_task(circular_buffer_source, manager, guild_id, rx).await });
@@ -84,16 +121,16 @@ async fn join(ctx: Context<'_>) -> Result<(), Error> {
 }
 
 async fn send_command(ctx: Context<'_>, command: StreamControl) -> Result<bool, Error> {
-    let stream_control_tx = ctx.data().stream_control_tx.lock().await;
+    let dj_control = ctx.data().dj_control.lock().await;
 
-    Ok(match &*stream_control_tx {
+    Ok(match &*dj_control {
         None => {
             ctx.reply("I can't control what doesn't exist, deep man.")
                 .await?;
             false
         }
-        Some(tx) => {
-            tx.send(command).await.unwrap();
+        Some(dj_control) => {
+            dj_control.stream_control_tx.send(command).await.unwrap();
             true
         }
     })
@@ -119,18 +156,29 @@ async fn pause(ctx: Context<'_>) -> Result<(), Error> {
 
 #[poise::command(slash_command, prefix_command, guild_only)]
 async fn disconnect(ctx: Context<'_>) -> Result<(), Error> {
-    let mut stream_control_tx = ctx.data().stream_control_tx.lock().await;
-    *stream_control_tx = None;
+    let mut dj_control = ctx.data().dj_control.lock().await;
 
-    let manager = songbird::get(ctx.serenity_context())
-        .await
-        .expect("Songbird Voice client placed in at initialisation.")
-        .clone();
+    let dj_control = dj_control.take();
 
-    manager.leave(ctx.guild_id().unwrap()).await?;
+    if let Some(dj_control) = dj_control {
+        dj_control.cancellation_token.cancel();
+        if let Err(err) = dj_control.udp_task.await {
+            error!("Failed to join UdpSource: {err}")
+        }
 
-    ctx.reply("Until next time, friend DJ Denny Jack OUT!")
-        .await?;
+        let manager = songbird::get(ctx.serenity_context())
+            .await
+            .expect("Songbird Voice client placed in at initialization.")
+            .clone();
+
+        manager.leave(ctx.guild_id().unwrap()).await?;
+
+        ctx.reply("Until next time, friend DJ Denny Jack OUT!")
+            .await?;
+    } else {
+        ctx.reply("You must be confused, I'm off the air waves for the moment")
+            .await?;
+    }
 
     Ok(())
 }
@@ -158,24 +206,19 @@ async fn main() {
     let intents =
         serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::MESSAGE_CONTENT;
 
-    let udp_input = UdpSocket::bind(cfg.server_addr).await.unwrap();
     let source = CircularBufferSource::default();
 
-    let mut udp_source = UdpSource::new(udp_input, source.clone());
-    let udp_task = tokio::spawn(async move { udp_source.worker().await });
-
     let data = Data {
         config: cfg.clone(),
-        udp_task,
         circular_buffer_source: source,
-        stream_control_tx: Mutex::new(None),
+        dj_control: Mutex::new(None),
     };
 
     let framework = poise::Framework::builder()
         .options(poise::FrameworkOptions {
             commands: vec![join(), play(), pause(), disconnect()],
             prefix_options: PrefixFrameworkOptions {
-                prefix: Some("!".to_string()),
+                prefix: Some("!".to_string().into()),
                 ..Default::default()
             },
             ..Default::default()

src/stream_task.rs

@@ -29,7 +29,7 @@ pub async fn stream_task(
     loop {
         let msg = match receiver.recv().await {
             None => {
-                info!("Stopping Stram Task...");
+                info!("Stopping Stream Task...");
                 return;
             }
             Some(msg) => msg,

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