615 lines
21 KiB
Rust
615 lines
21 KiB
Rust
use std::{
|
|
collections::BTreeMap,
|
|
net::UdpSocket,
|
|
sync::{
|
|
atomic::{AtomicBool, Ordering},
|
|
Arc, Mutex,
|
|
},
|
|
thread,
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use chacha20poly1305::{aead::KeyInit, ChaCha20Poly1305, Key};
|
|
use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
|
|
use ringbuf::{
|
|
storage::Heap,
|
|
traits::{Consumer, Observer, Producer, Split},
|
|
CachingCons, CachingProd, HeapRb, SharedRb,
|
|
};
|
|
use tauri::State;
|
|
|
|
use crate::{commands::config::ConfigState, crypto::SessionCipher};
|
|
|
|
const TARGET_SAMPLE_RATE: u32 = 48_000;
|
|
const PACKET_SAMPLES: usize = 960;
|
|
const MAX_PACKET_SIZE: usize = 4096;
|
|
|
|
const INITIAL_PACKET_CUSHION: usize = 3;
|
|
|
|
const VAD_THRESHOLD: f32 = 0.01;
|
|
const VAD_HANGOVER_FRAMES: usize = 10;
|
|
|
|
type AudioProducer = Arc<Mutex<CachingProd<Arc<SharedRb<Heap<f32>>>>>>;
|
|
type AudioConsumer = Arc<Mutex<CachingCons<Arc<SharedRb<Heap<f32>>>>>>;
|
|
|
|
#[derive(Clone, Default)]
|
|
pub struct VoiceState {
|
|
pub inner: Arc<VoiceStateInner>,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
pub struct VoiceStateInner {
|
|
pub session: Mutex<Option<VoiceSession>>,
|
|
}
|
|
|
|
pub struct VoiceSession {
|
|
pub input_stream: cpal::Stream,
|
|
pub output_stream: Arc<Mutex<cpal::Stream>>,
|
|
pub socket: Arc<UdpSocket>,
|
|
pub pin: u64,
|
|
pub shutdown: Arc<AtomicBool>,
|
|
|
|
pub current_input_device: Option<String>,
|
|
pub current_output_device: Option<String>,
|
|
|
|
pub producer_in: AudioProducer,
|
|
pub consumer_out: AudioConsumer,
|
|
}
|
|
|
|
// Simple Linear Resampler for real-time audio conversion
|
|
struct LinearResampler {
|
|
phase: f64,
|
|
}
|
|
|
|
impl LinearResampler {
|
|
fn new() -> Self {
|
|
Self { phase: 0.0 }
|
|
}
|
|
|
|
/// Resamples dynamic buffers from `src_rate` to `dst_rate`
|
|
fn process(&mut self, input: &[f32], src_rate: u32, dst_rate: u32, output: &mut Vec<f32>) {
|
|
if src_rate == dst_rate {
|
|
output.extend_from_slice(input);
|
|
return;
|
|
}
|
|
|
|
let ratio = src_rate as f64 / dst_rate as f64;
|
|
while self.phase < input.len() as f64 {
|
|
let idx = self.phase as usize;
|
|
let frac = (self.phase - idx as f64) as f32;
|
|
let next_idx = (idx + 1).min(input.len() - 1);
|
|
|
|
let sample = input[idx] * (1.0 - frac) + input[next_idx] * frac;
|
|
output.push(sample);
|
|
|
|
self.phase += ratio;
|
|
}
|
|
|
|
self.phase -= input.len() as f64;
|
|
}
|
|
}
|
|
|
|
impl VoiceSession {
|
|
pub fn update_input_device(
|
|
&mut self,
|
|
device_name: Option<String>,
|
|
state_inner: Arc<VoiceStateInner>,
|
|
) -> Result<(), String> {
|
|
let host = cpal::default_host();
|
|
let device = match &device_name {
|
|
Some(name) => host
|
|
.input_devices()
|
|
.map_err(|e| e.to_string())?
|
|
.find(|d| {
|
|
d.description()
|
|
.ok()
|
|
.map(|x| x.name() == *name)
|
|
.unwrap_or(false)
|
|
})
|
|
.ok_or_else(|| format!("Input device '{name}' not found"))?,
|
|
None => host
|
|
.default_input_device()
|
|
.ok_or("No default input device")?,
|
|
};
|
|
|
|
let input_config = device
|
|
.default_input_config()
|
|
.map_err(|e| format!("Failed to get default input config: {e}"))?
|
|
.config();
|
|
|
|
let producer = Arc::clone(&self.producer_in);
|
|
let new_stream = build_input_stream(&device, input_config, producer, state_inner)?;
|
|
|
|
new_stream
|
|
.play()
|
|
.map_err(|e| format!("Failed to play input stream: {e}"))?;
|
|
|
|
self.input_stream = new_stream;
|
|
self.current_input_device = device_name;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn update_output_device(
|
|
&mut self,
|
|
device_name: Option<String>,
|
|
state_inner: Arc<VoiceStateInner>,
|
|
) -> Result<(), String> {
|
|
let host = cpal::default_host();
|
|
let device = match &device_name {
|
|
Some(name) => host
|
|
.output_devices()
|
|
.map_err(|e| e.to_string())?
|
|
.find(|d| {
|
|
d.description()
|
|
.ok()
|
|
.map(|x| x.name() == *name)
|
|
.unwrap_or(false)
|
|
})
|
|
.ok_or_else(|| format!("Output device '{name}' not found"))?,
|
|
None => host
|
|
.default_output_device()
|
|
.ok_or("No default output device")?,
|
|
};
|
|
|
|
let output_config = device
|
|
.default_output_config()
|
|
.map_err(|e| format!("Failed to get default output config: {e}"))?
|
|
.config();
|
|
|
|
let consumer = Arc::clone(&self.consumer_out);
|
|
let new_stream = build_output_stream(&device, output_config, consumer, state_inner)?;
|
|
new_stream
|
|
.play()
|
|
.map_err(|e| format!("Failed to play output stream: {e}"))?;
|
|
|
|
if let Ok(mut active_stream) = self.output_stream.lock() {
|
|
*active_stream = new_stream;
|
|
}
|
|
|
|
self.current_output_device = device_name;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// Helper to build normalized Input Stream (Resampled & Downmixed to 48kHz Mono)
|
|
fn build_input_stream(
|
|
device: &cpal::Device,
|
|
config: cpal::StreamConfig,
|
|
producer: AudioProducer,
|
|
state_inner: Arc<VoiceStateInner>,
|
|
) -> Result<cpal::Stream, String> {
|
|
let native_sample_rate = config.sample_rate;
|
|
let channels = config.channels as usize;
|
|
let mut resampler = LinearResampler::new();
|
|
let mut mono_buffer = Vec::with_capacity(2048);
|
|
let mut resampled_buffer = Vec::with_capacity(2048);
|
|
|
|
let inner_input_err = Arc::clone(&state_inner);
|
|
|
|
device
|
|
.build_input_stream(
|
|
config,
|
|
move |data: &[f32], _| {
|
|
mono_buffer.clear();
|
|
resampled_buffer.clear();
|
|
|
|
// Downmix channels to mono
|
|
for chunk in data.chunks_exact(channels) {
|
|
let sum: f32 = chunk.iter().sum();
|
|
mono_buffer.push(sum / channels as f32);
|
|
}
|
|
|
|
// Resample to 48kHz standard target
|
|
resampler.process(
|
|
&mono_buffer,
|
|
native_sample_rate,
|
|
TARGET_SAMPLE_RATE,
|
|
&mut resampled_buffer,
|
|
);
|
|
|
|
if let Ok(mut prod) = producer.lock() {
|
|
let _ = prod.push_slice(&resampled_buffer);
|
|
}
|
|
},
|
|
move |err| {
|
|
eprintln!("[vc] Input error: {err}. Attempting recovery...");
|
|
if let Ok(mut lock) = inner_input_err.session.lock() {
|
|
if let Some(session) = lock.as_mut() {
|
|
let target_device = session.current_input_device.clone();
|
|
let _ = session
|
|
.update_input_device(target_device, Arc::clone(&inner_input_err));
|
|
}
|
|
}
|
|
},
|
|
None,
|
|
)
|
|
.map_err(|e| e.to_string())
|
|
}
|
|
|
|
// Helper to build normalized Output Stream (48kHz Mono -> Device Native Channels & Rate)
|
|
fn build_output_stream(
|
|
device: &cpal::Device,
|
|
config: cpal::StreamConfig,
|
|
consumer: AudioConsumer,
|
|
state_inner: Arc<VoiceStateInner>,
|
|
) -> Result<cpal::Stream, String> {
|
|
let native_sample_rate = config.sample_rate;
|
|
let channels = config.channels as usize;
|
|
let mut resampler = LinearResampler::new();
|
|
let mut raw_mono_samples = Vec::with_capacity(2048);
|
|
let mut resampled_mono = Vec::with_capacity(2048);
|
|
let mut last_sample = 0.0f32;
|
|
|
|
let inner_output_err = Arc::clone(&state_inner);
|
|
|
|
device
|
|
.build_output_stream(
|
|
config,
|
|
move |data: &mut [f32], _| {
|
|
let required_mono_samples = (data.len() / channels) * TARGET_SAMPLE_RATE as usize
|
|
/ native_sample_rate as usize;
|
|
|
|
raw_mono_samples.clear();
|
|
resampled_mono.clear();
|
|
|
|
if let Ok(mut cons) = consumer.lock() {
|
|
for _ in 0..required_mono_samples {
|
|
if let Some(s) = cons.try_pop() {
|
|
last_sample = s;
|
|
raw_mono_samples.push(s);
|
|
} else {
|
|
// Exponential decay to prevent clicking when underflowing
|
|
last_sample *= 0.92;
|
|
raw_mono_samples.push(last_sample);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Resample from 48kHz mono to target native output rate
|
|
resampler.process(
|
|
&raw_mono_samples,
|
|
TARGET_SAMPLE_RATE,
|
|
native_sample_rate,
|
|
&mut resampled_mono,
|
|
);
|
|
|
|
// Interleave mono into hardware channels
|
|
let mut res_idx = 0;
|
|
let mut out_idx = 0;
|
|
while out_idx < data.len() && res_idx < resampled_mono.len() {
|
|
let mono_val = resampled_mono[res_idx];
|
|
for ch in 0..channels {
|
|
if out_idx + ch < data.len() {
|
|
data[out_idx + ch] = mono_val;
|
|
}
|
|
}
|
|
out_idx += channels;
|
|
res_idx += 1;
|
|
}
|
|
},
|
|
move |err| {
|
|
eprintln!("[vc] Output error: {err}. Attempting recovery...");
|
|
if let Ok(mut lock) = inner_output_err.session.lock() {
|
|
if let Some(session) = lock.as_mut() {
|
|
let target_device = session.current_output_device.clone();
|
|
let _ = session
|
|
.update_output_device(target_device, Arc::clone(&inner_output_err));
|
|
}
|
|
}
|
|
},
|
|
None,
|
|
)
|
|
.map_err(|e| e.to_string())
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub fn list_input_devices() -> Result<Vec<String>, String> {
|
|
let host = cpal::default_host();
|
|
Ok(host
|
|
.input_devices()
|
|
.map_err(|e| e.to_string())?
|
|
.filter_map(|d| d.description().ok().map(|x| x.name().to_string()))
|
|
.collect())
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub fn list_output_devices() -> Result<Vec<String>, String> {
|
|
let host = cpal::default_host();
|
|
Ok(host
|
|
.output_devices()
|
|
.map_err(|e| e.to_string())?
|
|
.filter_map(|d| d.description().ok().map(|x| x.name().to_string()))
|
|
.collect())
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub fn disconnect_from_vc(voice_state: State<'_, VoiceState>) -> Result<(), String> {
|
|
let mut lock = voice_state
|
|
.inner
|
|
.session
|
|
.lock()
|
|
.map_err(|e| e.to_string())?;
|
|
|
|
if let Some(session) = lock.take() {
|
|
session.shutdown.store(true, Ordering::SeqCst);
|
|
let _ = session.input_stream.pause();
|
|
if let Ok(output) = session.output_stream.lock() {
|
|
let _ = output.pause();
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub fn connect_to_vc(
|
|
hostname: String,
|
|
pin: u64,
|
|
shared_secret: Vec<u8>, // output of js `x25519.getSharedSecret`
|
|
config_state: State<'_, ConfigState>,
|
|
voice_state: State<'_, VoiceState>,
|
|
) -> Result<(), String> {
|
|
disconnect_from_vc(voice_state.clone())?;
|
|
|
|
let key = Key::try_from(shared_secret.as_slice()).map_err(|v| v.to_string())?;
|
|
let cipher = Arc::new(Mutex::new(SessionCipher::new(ChaCha20Poly1305::new(&key))));
|
|
|
|
let config = config_state.0.lock().unwrap().clone();
|
|
let state_inner = Arc::clone(&voice_state.inner);
|
|
|
|
let socket = Arc::new(
|
|
UdpSocket::bind("0.0.0.0:0").map_err(|e| format!("Failed to bind UDP socket: {e}"))?,
|
|
);
|
|
socket
|
|
.connect(&hostname)
|
|
.map_err(|e| format!("Failed to connect UDP socket: {e}"))?;
|
|
socket
|
|
.set_read_timeout(Some(Duration::from_millis(5)))
|
|
.map_err(|e| e.to_string())?;
|
|
socket
|
|
.send(&pin.to_be_bytes())
|
|
.map_err(|e| format!("Failed to send pin: {e}"))?;
|
|
|
|
let host = cpal::default_host();
|
|
|
|
let input_device = match &config.input_device_name {
|
|
Some(name) => host
|
|
.input_devices()
|
|
.map_err(|e| e.to_string())?
|
|
.find(|d| {
|
|
d.description()
|
|
.ok()
|
|
.map(|x| x.name() == *name)
|
|
.unwrap_or(false)
|
|
})
|
|
.ok_or("Input device not found")?,
|
|
None => host
|
|
.default_input_device()
|
|
.ok_or("No default input device")?,
|
|
};
|
|
|
|
let output_device = match &config.output_device_name {
|
|
Some(name) => host
|
|
.output_devices()
|
|
.map_err(|e| e.to_string())?
|
|
.find(|d| {
|
|
d.description()
|
|
.ok()
|
|
.map(|x| x.name() == *name)
|
|
.unwrap_or(false)
|
|
})
|
|
.ok_or("Output device not found")?,
|
|
None => host
|
|
.default_output_device()
|
|
.ok_or("No default output device")?,
|
|
};
|
|
|
|
let shutdown = Arc::new(AtomicBool::new(false));
|
|
|
|
// Ring Buffer Setup
|
|
let rb_out = HeapRb::<f32>::new(19200);
|
|
let (producer_out, consumer_out) = rb_out.split();
|
|
let mut producer_out = producer_out;
|
|
let shared_consumer_out = Arc::new(Mutex::new(consumer_out));
|
|
|
|
let output_config = output_device
|
|
.default_output_config()
|
|
.map_err(|e| format!("Failed to get default output config: {e}"))?
|
|
.config();
|
|
|
|
let output_stream = build_output_stream(
|
|
&output_device,
|
|
output_config,
|
|
Arc::clone(&shared_consumer_out),
|
|
Arc::clone(&state_inner),
|
|
)?;
|
|
output_stream
|
|
.play()
|
|
.map_err(|e| format!("Failed to start output: {e}"))?;
|
|
let output_stream = Arc::new(Mutex::new(output_stream));
|
|
|
|
let rb_in = HeapRb::<f32>::new(19200);
|
|
let (producer_in, mut consumer_in) = rb_in.split();
|
|
let shared_producer_in = Arc::new(Mutex::new(producer_in));
|
|
|
|
let input_config = input_device
|
|
.default_input_config()
|
|
.map_err(|e| format!("Failed to get default input config: {e}"))?
|
|
.config();
|
|
|
|
let input_stream = build_input_stream(
|
|
&input_device,
|
|
input_config,
|
|
Arc::clone(&shared_producer_in),
|
|
Arc::clone(&state_inner),
|
|
)?;
|
|
|
|
// Sender Thread
|
|
{
|
|
let input_socket = socket.clone();
|
|
let shutdown = shutdown.clone();
|
|
let cipher = cipher.clone();
|
|
|
|
thread::spawn(move || {
|
|
let mut sequence = 0u32;
|
|
let mut frame_buf = vec![0.0f32; PACKET_SAMPLES];
|
|
let mut hangover_counter = 0;
|
|
|
|
while !shutdown.load(Ordering::Relaxed) {
|
|
if consumer_in.occupied_len() >= PACKET_SAMPLES {
|
|
let _ = consumer_in.pop_slice(&mut frame_buf);
|
|
|
|
let sum_squares: f32 = frame_buf.iter().map(|&s| s * s).sum();
|
|
let rms = (sum_squares / PACKET_SAMPLES as f32).sqrt();
|
|
|
|
let is_speaking = if rms >= VAD_THRESHOLD {
|
|
hangover_counter = VAD_HANGOVER_FRAMES;
|
|
true
|
|
} else if hangover_counter > 0 {
|
|
hangover_counter -= 1;
|
|
true
|
|
} else {
|
|
false
|
|
};
|
|
|
|
if is_speaking {
|
|
// sequence goes INSIDE the plaintext now, prefixed before the PCM
|
|
let mut plaintext = Vec::with_capacity(4 + PACKET_SAMPLES * 2);
|
|
plaintext.extend_from_slice(&sequence.to_be_bytes());
|
|
|
|
for sample in frame_buf.iter() {
|
|
let pcm = (sample.clamp(-1.0, 1.0) * 32767.0) as i16;
|
|
plaintext.extend_from_slice(&pcm.to_be_bytes());
|
|
}
|
|
|
|
let Ok(net_packet) = cipher.lock().unwrap().encrypt(&plaintext) else {
|
|
eprintln!("[vc] failed to encrypt outgoing packet");
|
|
sequence = sequence.wrapping_add(1);
|
|
continue;
|
|
};
|
|
|
|
let _ = input_socket.send(&net_packet);
|
|
sequence = sequence.wrapping_add(1);
|
|
}
|
|
} else {
|
|
thread::sleep(Duration::from_millis(2));
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
// Receiver Thread
|
|
{
|
|
let socket = socket.clone();
|
|
let shutdown = shutdown.clone();
|
|
let cipher = cipher.clone();
|
|
|
|
thread::spawn(move || {
|
|
let mut packets: BTreeMap<u32, Vec<f32>> = BTreeMap::new();
|
|
let mut expected: Option<u32> = None;
|
|
let mut is_prebuffering = true;
|
|
let mut last_good_frame = vec![0.0f32; PACKET_SAMPLES];
|
|
let mut udp_buffer = [0u8; MAX_PACKET_SIZE];
|
|
let mut next_frame_time = Instant::now();
|
|
|
|
while !shutdown.load(Ordering::Relaxed) {
|
|
if let Ok(len) = socket.recv(&mut udp_buffer) {
|
|
match cipher.lock().unwrap().decrypt(&udp_buffer[..len]) {
|
|
Ok(plaintext) => {
|
|
if plaintext.len() < 4 {
|
|
eprintln!("[vc] dropped packet: too short after decrypt");
|
|
} else {
|
|
let seq = u32::from_be_bytes(plaintext[..4].try_into().unwrap());
|
|
let pcm = &plaintext[4..];
|
|
|
|
let samples: Vec<f32> = pcm
|
|
.chunks_exact(2)
|
|
.map(|c| i16::from_be_bytes([c[0], c[1]]) as f32 / 32768.0)
|
|
.collect();
|
|
|
|
packets.insert(seq, samples);
|
|
}
|
|
}
|
|
Err(_) => {
|
|
eprintln!("[vc] dropped packet: decryption failed");
|
|
}
|
|
}
|
|
}
|
|
|
|
if is_prebuffering {
|
|
if packets.len() >= INITIAL_PACKET_CUSHION {
|
|
expected = packets.keys().next().copied();
|
|
is_prebuffering = false;
|
|
next_frame_time = Instant::now();
|
|
} else {
|
|
thread::sleep(Duration::from_millis(1));
|
|
continue;
|
|
}
|
|
}
|
|
|
|
if Instant::now() >= next_frame_time {
|
|
if let Some(seq) = expected {
|
|
if let Some(samples) = packets.remove(&seq) {
|
|
if samples.len() == PACKET_SAMPLES {
|
|
last_good_frame.copy_from_slice(&samples);
|
|
} else {
|
|
last_good_frame.clear();
|
|
last_good_frame
|
|
.extend(samples.iter().take(PACKET_SAMPLES).copied());
|
|
if last_good_frame.len() < PACKET_SAMPLES {
|
|
last_good_frame.resize(PACKET_SAMPLES, 0.0);
|
|
}
|
|
}
|
|
|
|
let _ = producer_out.push_slice(&last_good_frame);
|
|
expected = Some(seq.wrapping_add(1));
|
|
} else if packets.keys().any(|&x| x > seq) {
|
|
for sample in last_good_frame.iter_mut() {
|
|
*sample *= 0.65;
|
|
}
|
|
let _ = producer_out.push_slice(&last_good_frame);
|
|
expected = Some(seq.wrapping_add(1));
|
|
} else if packets.is_empty() {
|
|
is_prebuffering = true;
|
|
}
|
|
}
|
|
next_frame_time += Duration::from_millis(20);
|
|
}
|
|
|
|
thread::sleep(Duration::from_millis(1));
|
|
}
|
|
});
|
|
}
|
|
|
|
input_stream
|
|
.play()
|
|
.map_err(|e| format!("Failed to start input: {e}"))?;
|
|
|
|
*state_inner.session.lock().unwrap() = Some(VoiceSession {
|
|
input_stream,
|
|
output_stream,
|
|
socket,
|
|
pin,
|
|
shutdown,
|
|
current_input_device: config.input_device_name,
|
|
current_output_device: config.output_device_name,
|
|
producer_in: shared_producer_in,
|
|
consumer_out: shared_consumer_out,
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
|
|
impl Drop for VoiceSession {
|
|
fn drop(&mut self) {
|
|
self.shutdown.store(true, Ordering::SeqCst);
|
|
|
|
let _ = self.input_stream.pause();
|
|
if let Ok(output) = self.output_stream.lock() {
|
|
let _ = output.pause();
|
|
}
|
|
|
|
eprintln!("[vc] VoiceSession dropped and audio streams paused.");
|
|
}
|
|
}
|