145 lines
4.5 KiB
Rust
145 lines
4.5 KiB
Rust
use crate::circular_buffer_source::CircularBufferSource;
|
|
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: u32 = 48_000;
|
|
|
|
#[derive(Debug, Error)]
|
|
#[allow(clippy::enum_variant_names)]
|
|
pub enum Error {
|
|
#[error("IO Error {0}")]
|
|
IOError(#[from] std::io::Error),
|
|
#[error("Failed to decode audio in packet: {0}")]
|
|
OpusError(#[from] opus::Error),
|
|
#[error("Failed to decode packet: {0}")]
|
|
DecodeError(#[from] DecodeError),
|
|
}
|
|
|
|
pub struct UdpSource {
|
|
udp: UdpSocket,
|
|
source: CircularBufferSource,
|
|
decoder: opus::Decoder,
|
|
expected_seq: u32,
|
|
key: RandomizedNonceKey,
|
|
}
|
|
|
|
impl UdpSource {
|
|
pub fn new(udp_socket: UdpSocket, source: CircularBufferSource, key: &[u8]) -> Self {
|
|
Self {
|
|
udp: udp_socket,
|
|
source,
|
|
decoder: opus::Decoder::new(SAMPLE_RATE, opus::Channels::Stereo).unwrap(),
|
|
expected_seq: 0,
|
|
key: RandomizedNonceKey::new(&AES_256_GCM, key).unwrap(),
|
|
}
|
|
}
|
|
|
|
fn decode_data(
|
|
&mut self,
|
|
input_buffer: &[u8],
|
|
output_buffer: &mut [f32],
|
|
) -> Result<usize, Error> {
|
|
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(size_of_val(samples));
|
|
|
|
for sample in samples {
|
|
let bytes = sample.to_le_bytes();
|
|
|
|
sample_bytes.extend_from_slice(bytes.as_slice());
|
|
}
|
|
|
|
self.source.write_all(&sample_bytes)?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn recv_message(&mut self, buffer: &mut [u8]) -> Result<usize, std::io::Error> {
|
|
self.udp.recv(buffer).await
|
|
}
|
|
|
|
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 = 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(&*decrypted_msg) {
|
|
Ok(tap_msg) => tap_msg,
|
|
Err(err) => {
|
|
error!("Failed decode proto message: {err:?}");
|
|
continue;
|
|
}
|
|
};
|
|
|
|
if tap_message.seq != self.expected_seq {
|
|
debug!(
|
|
"Mismatch sequence count found expected '{}': {tap_message:?}",
|
|
tap_message.seq
|
|
)
|
|
}
|
|
|
|
(self.expected_seq, _) = tap_message.seq.overflowing_add(1);
|
|
|
|
let decoded_samples = match self.decode_data(&tap_message.audio_data, &mut decode_data)
|
|
{
|
|
Ok(decoded_samples) => decoded_samples,
|
|
Err(err) => {
|
|
error!("Failed to decode audio data: {err:?}");
|
|
continue;
|
|
}
|
|
};
|
|
|
|
if let Err(err) = self.write_audio_data_to_source(&decode_data[..decoded_samples]) {
|
|
error!("Failed to write data to source: {err:?}");
|
|
}
|
|
}
|
|
}
|
|
}
|