144 lines
4.2 KiB
Rust
144 lines
4.2 KiB
Rust
use bytes::{BufMut, BytesMut};
|
|
use futures::StreamExt;
|
|
use livekit::prelude::*;
|
|
use livekit::webrtc::audio_stream::native::NativeAudioStream;
|
|
use std::env;
|
|
use tokio::fs::File;
|
|
use tokio::io::{AsyncWriteExt, BufWriter};
|
|
|
|
const WAV_HEADER_SIZE: usize = 44;
|
|
const FILE_PATH: &str = "record.wav";
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct WavHeader {
|
|
pub sample_rate: u32,
|
|
pub bit_depth: u16,
|
|
pub num_channels: u16,
|
|
}
|
|
|
|
pub struct WavWriter {
|
|
header: WavHeader,
|
|
data: BytesMut,
|
|
writer: BufWriter<File>,
|
|
}
|
|
|
|
impl WavWriter {
|
|
pub async fn create<P: AsRef<std::path::Path>>(
|
|
path: P,
|
|
header: WavHeader,
|
|
) -> Result<WavWriter, std::io::Error> {
|
|
let file = File::create(path).await?;
|
|
let writer = BufWriter::new(file);
|
|
|
|
let mut wav_writer = WavWriter {
|
|
header,
|
|
data: BytesMut::new(),
|
|
writer,
|
|
};
|
|
|
|
wav_writer.write_header()?;
|
|
Ok(wav_writer)
|
|
}
|
|
|
|
fn write_header(&mut self) -> Result<(), std::io::Error> {
|
|
let byte_rate = (self.header.sample_rate
|
|
* self.header.bit_depth as u32
|
|
* self.header.num_channels as u32);
|
|
|
|
let block_align = byte_rate / self.header.sample_rate as u32;
|
|
|
|
self.data.put_slice(b"RIFF");
|
|
self.data.put_u32_le(0); // Placeholder for file size
|
|
self.data.put_slice(b"WAVE");
|
|
self.data.put_slice(b"fmt ");
|
|
self.data.put_u32_le(16); // Subchunk1Size (16 for PCM)
|
|
self.data.put_u16_le(1); // AudioFormat (1 for PCM)
|
|
self.data.put_u16_le(self.header.num_channels);
|
|
self.data.put_u32_le(self.header.sample_rate);
|
|
self.data.put_u32_le(byte_rate);
|
|
self.data.put_u16_le(32);
|
|
self.data.put_u16_le(self.header.bit_depth);
|
|
self.data.put_slice(b"data");
|
|
self.data.put_u32_le(0); // Placeholder for data size
|
|
|
|
assert_eq!(self.data.len(), WAV_HEADER_SIZE);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn write_sample(&mut self, sample: i16) -> Result<(), std::io::Error> {
|
|
self.data.put_i16_le(sample);
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn finalize(mut self) -> Result<(), std::io::Error> {
|
|
let data_size = self.data.len() as u32 - WAV_HEADER_SIZE as u32;
|
|
let file_size = data_size + WAV_HEADER_SIZE as u32 - 8;
|
|
self.data.as_mut()[4..8].copy_from_slice(&file_size.to_le_bytes());
|
|
self.data.as_mut()[40..44].copy_from_slice(&data_size.to_le_bytes());
|
|
|
|
self.writer.write_all(&self.data).await?;
|
|
self.writer.flush().await?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let url = env::var("LIVEKIT_URL").expect("LIVEKIT_URL is not set");
|
|
let token = env::var("LIVEKIT_TOKEN").expect("LIVEKIT_TOKEN is not set");
|
|
|
|
let (room, mut rx) = Room::connect(&url, &token).await.unwrap();
|
|
let session = room.session();
|
|
println!("Connected to room: {} - {}", session.name(), session.sid());
|
|
|
|
while let Some(msg) = rx.recv().await {
|
|
match msg {
|
|
RoomEvent::TrackSubscribed {
|
|
track,
|
|
publication: _,
|
|
participant: _,
|
|
} => {
|
|
if let RemoteTrack::Audio(audio_track) = track {
|
|
record_track(audio_track).await.unwrap();
|
|
break;
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
println!("Done");
|
|
}
|
|
|
|
async fn record_track(audio_track: RemoteAudioTrack) -> Result<(), std::io::Error> {
|
|
println!("Recording track {:?}", audio_track.sid());
|
|
let rtc_track = audio_track.rtc_track();
|
|
|
|
// TODO(theomonnom): Remove hardcoded values
|
|
let header = WavHeader {
|
|
sample_rate: 48000,
|
|
bit_depth: 16,
|
|
num_channels: 1,
|
|
};
|
|
|
|
let mut wav_writer = WavWriter::create(FILE_PATH, header).await?;
|
|
let mut audio_stream = NativeAudioStream::new(rtc_track);
|
|
|
|
let max_record = 5 * header.sample_rate * header.num_channels as u32;
|
|
let mut sample_count = 0;
|
|
'recv_loop: while let Some(frame) = audio_stream.next().await {
|
|
for sample in frame.data {
|
|
wav_writer.write_sample(sample).await.unwrap();
|
|
sample_count += 1;
|
|
|
|
if sample_count >= max_record {
|
|
break 'recv_loop;
|
|
}
|
|
}
|
|
}
|
|
|
|
wav_writer.finalize().await?;
|
|
Ok(())
|
|
}
|