Added more commands to control stream, added protobuf message support

src/udp_source.rs

@@ -1,54 +1,116 @@
-use crate::stream::Stream;
+use crate::circular_buffer_source::CircularBufferSource;
 use audiopus::coder::Decoder;
 use audiopus::{Channels, MutSignals, SampleRate};
+use log::{debug, error, info};
+use prost::{DecodeError, Message};
 use std::io::Write;
+use tap_interface::tap::message::TapMessage;
+use thiserror::Error;
 use tokio::net::UdpSocket;
 
 pub const SAMPLE_RATE: SampleRate = SampleRate::Hz48000;
-pub const SAMPLE_RATE_RAW: usize = 48_000;
-pub const AUDIO_FRAME_RATE: usize = 50;
-pub const MONO_FRAME_SIZE: usize = SAMPLE_RATE_RAW / AUDIO_FRAME_RATE;
-pub const STEREO_FRAME_SIZE: usize = 2 * MONO_FRAME_SIZE;
+
+#[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] audiopus::Error),
+    #[error("Failed to decode packet: {0}")]
+    DecodeError(#[from] DecodeError),
+}
 
 pub struct UdpSource {
     udp: UdpSocket,
-    stream: Stream,
+    source: CircularBufferSource,
     decoder: Decoder,
+    expected_seq: u32,
 }
 
 impl UdpSource {
-    pub fn new(udp_socket: UdpSocket, stream: Stream) -> Self {
+    pub fn new(udp_socket: UdpSocket, source: CircularBufferSource) -> Self {
         Self {
             udp: udp_socket,
-            stream,
+            source,
             decoder: Decoder::new(SAMPLE_RATE, Channels::Stereo).unwrap(),
+            expected_seq: 0,
         }
     }
 
+    fn decode_data(
+        &mut self,
+        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)?;
+
+        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));
+
+        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
+    }
+
     pub async fn worker(&mut self) {
+        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 len = self.udp.recv(&mut buffer).await.unwrap();
+            let msg_size = match self.recv_message(&mut buffer).await {
+                Ok(msg_size) => msg_size,
+                Err(err) => {
+                    error!("Failed to recv udp packet: {err:?}");
+                    continue;
+                }
+            };
 
-            let packet = audiopus::packet::Packet::try_from(&buffer[12..len]).unwrap();
+            let tap_message = match TapMessage::decode(&buffer[..msg_size]) {
+                Ok(tap_msg) => tap_msg,
+                Err(err) => {
+                    error!("Failed decode proto message: {err:?}");
+                    continue;
+                }
+            };
 
-            let mut samples = vec![0.0; STEREO_FRAME_SIZE];
-            let signals = MutSignals::try_from(&mut samples).unwrap();
-            let sample_size = self
-                .decoder
-                .decode_float(Some(packet), signals, false)
-                .unwrap()
-                * 2;
-
-            let mut sample_bytes = Vec::with_capacity(sample_size * std::mem::size_of::<f32>());
-
-            for sample in &samples[0..sample_size] {
-                let bytes = sample.to_le_bytes();
-
-                sample_bytes.extend_from_slice(bytes.as_slice());
+            if tap_message.seq != self.expected_seq {
+                debug!(
+                    "Mismatch sequence count found expected '{}': {tap_message:?}",
+                    tap_message.seq
+                )
             }
 
-            self.stream.write_all(&sample_bytes).unwrap();
+            (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:?}");
+            }
         }
     }
 }