use std::collections::HashMap; use std::convert::Infallible; use std::fs::{self, File}; use std::path::{Path as FilePath, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use axum::extract::Path; use axum::http::StatusCode; use axum::response::sse::{Event, KeepAlive, Sse}; use axum::response::IntoResponse; use axum::Json; use file_format::{FileFormat, Kind}; use futures::Stream; use once_cell::sync::Lazy; use ropus::{Application, Bitrate, Channels as OpusChannels, DecodeMode, Decoder as OpusDecoder, Encoder as OpusEncoder}; use rubato::audioadapter::Adapter; use rubato::audioadapter_buffers::direct::SequentialSliceOfVecs; use rubato::{Fft, FixedSync, Resampler}; use serde::{Deserialize, Serialize}; use symphonia::core::audio::sample::Sample; use symphonia::core::codecs::audio::{well_known::CODEC_ID_OPUS, AudioDecoder, AudioDecoderOptions}; use symphonia::core::codecs::CodecParameters; use symphonia::core::errors::Error as SymphoniaError; use symphonia::core::formats::{FormatOptions, Track, TrackFlags, TrackType}; use symphonia::core::io::MediaSourceStream; use symphonia::core::meta::MetadataOptions; use symphonia::core::formats::probe::Hint; use symphonia::core::units::Timestamp; use tokio::sync::broadcast; use tokio_stream::wrappers::BroadcastStream; use tokio_stream::StreamExt; use webm_iterable::matroska_spec::{Master, MatroskaSpec, SimpleBlock}; use webm_iterable::{WebmIterator, WebmWriter, WriteOptions}; use crate::api_token::APIToken; const OUTPUT_SAMPLE_RATE: u32 = 48_000; const OPUS_FRAME_SAMPLES: usize = 960; const OPUS_BITRATE: u32 = 32_000; const CLUSTER_DURATION_MS: u64 = 30_000; const OPUS_PRE_SKIP: u16 = 312; const DEFAULT_MAX_PASS_THROUGH_BYTES: u64 = 25 * 1024 * 1024; static JOBS: Lazy>>> = Lazy::new(|| RwLock::new(HashMap::new())); #[derive(Debug, Deserialize)] pub struct CreateMediaJobRequest { pub input_path: String, pub output_path: Option, pub max_pass_through_bytes: Option, } #[derive(Debug, Serialize)] pub struct CreateMediaJobResponse { pub job_id: String, } #[derive(Clone, Debug, Serialize)] #[serde(rename_all = "snake_case")] pub enum MediaJobPhase { Probing, Transcoding, Completed, Failed, Cancelled, } #[derive(Clone, Debug, Serialize)] pub struct MediaJobEvent { pub phase: MediaJobPhase, pub progress: Option, pub result: Option, pub error: Option, } #[derive(Clone, Debug, Serialize)] pub struct MediaJobResult { pub output_path: String, pub detected_format: String, pub detected_codec: String, pub duration_ms: u64, pub pass_through: bool, } #[derive(Clone, Debug, Serialize)] pub struct MediaError { pub code: String, pub message: String, } impl MediaError { fn new(code: &str, message: impl Into) -> Self { Self { code: code.to_string(), message: message.into() } } } struct MediaJob { cancelled: AtomicBool, current: Mutex, events: broadcast::Sender, } impl MediaJob { fn new() -> Self { let initial = MediaJobEvent { phase: MediaJobPhase::Probing, progress: Some(0.0), result: None, error: None, }; let (events, _) = broadcast::channel(32); Self { cancelled: AtomicBool::new(false), current: Mutex::new(initial), events } } fn publish(&self, event: MediaJobEvent) { *self.current.lock().unwrap() = event.clone(); let _ = self.events.send(event); } fn is_cancelled(&self) -> bool { self.cancelled.load(Ordering::Relaxed) } } pub async fn create_job( _token: APIToken, Json(request): Json, ) -> Result, (StatusCode, Json)> { let input_path = PathBuf::from(&request.input_path); if !input_path.is_file() { return Err((StatusCode::BAD_REQUEST, Json(MediaError::new("file_not_found", "The selected media file does not exist.")))); } let output_path = request.output_path.map(PathBuf::from).unwrap_or_else(|| { let parent = input_path.parent().unwrap_or_else(|| FilePath::new(".")); let stem = input_path.file_stem().and_then(|value| value.to_str()).unwrap_or("media"); parent.join(format!("{stem}-normalized.webm")) }); let job_id = format!("{}-{}", std::process::id(), rand::random::()); let job = Arc::new(MediaJob::new()); JOBS.write().unwrap().insert(job_id.clone(), Arc::clone(&job)); let completed_job_id = job_id.clone(); tauri::async_runtime::spawn(async move { let max_pass_through_bytes = request.max_pass_through_bytes.unwrap_or(DEFAULT_MAX_PASS_THROUGH_BYTES); let task_job = Arc::clone(&job); let result = tokio::task::spawn_blocking(move || normalize_media(&input_path, &output_path, max_pass_through_bytes, &task_job)).await; match result { Ok(Ok(result)) => job.publish(MediaJobEvent { phase: MediaJobPhase::Completed, progress: Some(1.0), result: Some(result), error: None, }), Ok(Err(error)) if error.code == "cancelled" => job.publish(MediaJobEvent { phase: MediaJobPhase::Cancelled, progress: None, result: None, error: None, }), Ok(Err(error)) => job.publish(MediaJobEvent { phase: MediaJobPhase::Failed, progress: None, result: None, error: Some(error), }), Err(error) => job.publish(MediaJobEvent { phase: MediaJobPhase::Failed, progress: None, result: None, error: Some(MediaError::new("internal_error", format!("The media worker failed: {error}"))), }), } tokio::time::sleep(std::time::Duration::from_secs(600)).await; JOBS.write().unwrap().remove(&completed_job_id); }); Ok(Json(CreateMediaJobResponse { job_id })) } pub async fn get_job_events( _token: APIToken, Path(job_id): Path, ) -> Result>>, StatusCode> { let job = JOBS.read().unwrap().get(&job_id).cloned().ok_or(StatusCode::NOT_FOUND)?; let current = job.current.lock().unwrap().clone(); let initial = tokio_stream::once(current); let updates = BroadcastStream::new(job.events.subscribe()).filter_map(|event| event.ok()); let stream = initial.chain(updates).map(|event| { let data = serde_json::to_string(&event).unwrap_or_else(|_| "{}".to_string()); Ok(Event::default().event(phase_name(&event.phase)).data(data)) }); Ok(Sse::new(stream).keep_alive(KeepAlive::default())) } pub async fn cancel_job(_token: APIToken, Path(job_id): Path) -> impl IntoResponse { match JOBS.read().unwrap().get(&job_id) { Some(job) => { job.cancelled.store(true, Ordering::Relaxed); StatusCode::NO_CONTENT } None => StatusCode::NOT_FOUND, } } fn phase_name(phase: &MediaJobPhase) -> &'static str { match phase { MediaJobPhase::Probing => "probing", MediaJobPhase::Transcoding => "transcoding", MediaJobPhase::Completed => "completed", MediaJobPhase::Failed => "failed", MediaJobPhase::Cancelled => "cancelled", } } fn normalize_media(input_path: &FilePath, output_path: &FilePath, max_pass_through_bytes: u64, job: &MediaJob) -> Result { check_cancelled(job)?; let detected = FileFormat::from_file(input_path) .map_err(|error| MediaError::new("unknown_format", format!("The file type could not be identified: {error}")))?; if detected.kind() == Kind::Executable { return Err(MediaError::new("unsafe_file", "The selected file contains executable data and cannot be processed as media.")); } if !matches!(detected.kind(), Kind::Audio | Kind::Video) && !matches!(detected, FileFormat::ExtensibleBinaryMetaLanguage) { return Err(MediaError::new("not_media", format!("The selected file is not supported media (detected as {detected:?})."))); } let file = File::open(input_path).map_err(|error| MediaError::new("file_open_failed", error.to_string()))?; let mss = MediaSourceStream::new(Box::new(file), Default::default()); let mut hint = Hint::new(); if let Some(extension) = input_path.extension().and_then(|value| value.to_str()) { hint.with_extension(extension); } let mut format = symphonia::default::get_probe() .probe(&hint, mss, FormatOptions::default(), MetadataOptions::default()) .map_err(map_probe_error)?; let detected_format = format!("{detected:?} / {}", format.format_info().long_name); let tracks = format.tracks(); let mut audio_tracks: Vec<&Track> = tracks.iter().filter(|track| { matches!(track.codec_params, Some(CodecParameters::Audio(_))) }).collect(); if audio_tracks.is_empty() { return Err(MediaError::new("no_audio_track", "The selected media file does not contain an audio track.")); } audio_tracks.sort_by_key(|track| !track.flags.contains(TrackFlags::DEFAULT)); let selected = audio_tracks.into_iter().find(|track| is_decodable(track)) .ok_or_else(|| MediaError::new("unsupported_codec", "None of the audio tracks uses a supported codec."))?; let track_id = selected.id; let params = selected.codec_params.as_ref().and_then(CodecParameters::audio).unwrap().clone(); let detected_codec = if params.codec == CODEC_ID_OPUS { "opus".to_string() } else { format!("{}", params.codec) }; let duration_ms = selected.time_base.zip(selected.duration).and_then(|(time_base, duration)| { let timestamp = Timestamp::new(i64::try_from(duration.get()).ok()?); let time = time_base.calc_time(timestamp)?; Some((time.as_secs_f64() * 1000.0).max(0.0).round() as u64) }).unwrap_or(0); let audio_only = tracks.iter().all(|track| track.track_type() != Some(TrackType::Video)); let channels = params.channels.as_ref().map(|value| value.count()).unwrap_or(0); let pass_through = is_webm_container(input_path) && audio_only && tracks.iter().filter(|track| track.track_type() == Some(TrackType::Audio)).count() == 1 && params.codec == CODEC_ID_OPUS && params.sample_rate == Some(OUTPUT_SAMPLE_RATE) && channels == 1 && input_path.metadata().map(|metadata| metadata.len() <= max_pass_through_bytes).unwrap_or(false); let partial_path = partial_path(output_path); if let Some(parent) = partial_path.parent() { fs::create_dir_all(parent).map_err(|error| MediaError::new("output_create_failed", error.to_string()))?; } let result = if pass_through { check_cancelled(job)?; fs::copy(input_path, &partial_path).map_err(|error| MediaError::new("output_write_failed", error.to_string()))?; check_cancelled(job)?; Ok(MediaJobResult { output_path: output_path.to_string_lossy().into_owned(), detected_format: detected_format.clone(), detected_codec, duration_ms, pass_through: true, }) } else { job.publish(MediaJobEvent { phase: MediaJobPhase::Transcoding, progress: Some(0.0), result: None, error: None, }); transcode(&mut *format, track_id, ¶ms, &partial_path, output_path, detected_format, detected_codec, duration_ms, job) }; match result { Ok(result) => { if let Err(error) = fs::rename(&partial_path, output_path) { let _ = fs::remove_file(&partial_path); return Err(MediaError::new("output_commit_failed", error.to_string())); } Ok(result) } Err(error) => { let _ = fs::remove_file(&partial_path); Err(error) } } } fn is_decodable(track: &Track) -> bool { let Some(params) = track.codec_params.as_ref().and_then(CodecParameters::audio) else { return false; }; params.codec == CODEC_ID_OPUS || symphonia::default::get_codecs().make_audio_decoder(params, &AudioDecoderOptions::default()).is_ok() } fn transcode( format: &mut dyn symphonia::core::formats::FormatReader, track_id: u32, params: &symphonia::core::codecs::audio::AudioCodecParameters, partial_path: &FilePath, output_path: &FilePath, detected_format: String, detected_codec: String, expected_duration_ms: u64, job: &MediaJob, ) -> Result { let input_rate = params.sample_rate.unwrap_or(OUTPUT_SAMPLE_RATE); let input_channels = params.channels.as_ref().map(|value| value.count()).unwrap_or(1).max(1); let mut symphonia_decoder: Option> = if params.codec == CODEC_ID_OPUS { None } else { Some(symphonia::default::get_codecs().make_audio_decoder(params, &AudioDecoderOptions::default()) .map_err(|_| MediaError::new("unsupported_codec", "The selected audio codec is not supported."))?) }; let mut opus_decoder = if params.codec == CODEC_ID_OPUS { let channels = if input_channels == 1 { OpusChannels::Mono } else { OpusChannels::Stereo }; Some(OpusDecoder::new(OUTPUT_SAMPLE_RATE, channels) .map_err(|error| MediaError::new("decoder_init_failed", error.to_string()))?) } else { None }; let mut opus_encoder = OpusEncoder::builder(OUTPUT_SAMPLE_RATE, OpusChannels::Mono, Application::Audio) .bitrate(Bitrate::Bits(OPUS_BITRATE)) .vbr(true) .build() .map_err(|error| MediaError::new("encoder_init_failed", error.to_string()))?; let file = File::create(partial_path).map_err(|error| MediaError::new("output_create_failed", error.to_string()))?; let mut writer = WebmOpusWriter::new(file)?; let mut pending = Vec::::with_capacity(OPUS_FRAME_SAMPLES * 3); let mut interleaved = Vec::::new(); let mut opus_pcm = vec![0i16; 5_760 * input_channels]; let mut encoded = [0u8; 4_000]; let mut produced_samples = 0u64; let mut input_samples = 0u64; let mut resampled_samples = 0u64; loop { check_cancelled(job)?; let packet = match format.next_packet() { Ok(Some(packet)) => packet, Ok(None) => break, Err(SymphoniaError::ResetRequired) => return Err(MediaError::new("stream_reset", "The media stream changed unexpectedly.")), Err(SymphoniaError::IoError(error)) if error.kind() == std::io::ErrorKind::UnexpectedEof => break, Err(error) => return Err(MediaError::new("damaged_container", format!("The media container is damaged: {error}"))), }; if packet.track_id != track_id { continue; } let mono = if let Some(decoder) = symphonia_decoder.as_mut() { let decoded = match decoder.decode(&packet) { Ok(decoded) => decoded, Err(SymphoniaError::DecodeError(_)) => continue, Err(error) => return Err(MediaError::new("decode_failed", format!("Audio decoding failed: {error}"))), }; interleaved.resize(decoded.samples_interleaved(), f32::MID); decoded.copy_to_slice_interleaved(&mut interleaved); downmix_to_mono(&interleaved, decoded.num_planes()) } else { let decoder = opus_decoder.as_mut().unwrap(); let frames = decoder.decode(&packet.data, &mut opus_pcm, DecodeMode::Normal) .map_err(|error| MediaError::new("decode_failed", error.to_string()))?; let samples = &opus_pcm[..frames * input_channels]; if input_channels == 1 { samples.iter().map(|sample| f32::from(*sample) / 32768.0).collect() } else { samples.chunks_exact(input_channels).map(|frame| { frame.iter().map(|sample| f32::from(*sample) / 32768.0).sum::() / input_channels as f32 }).collect() } }; input_samples += mono.len() as u64; let expected_resampled_samples = input_samples.saturating_mul(u64::from(OUTPUT_SAMPLE_RATE)) / u64::from(input_rate); let expected_chunk_samples = expected_resampled_samples.saturating_sub(resampled_samples) as usize; let mut resampled = resample_chunk(&mono, input_rate)?; resampled.resize(expected_chunk_samples, 0.0); resampled.truncate(expected_chunk_samples); resampled_samples += resampled.len() as u64; pending.extend_from_slice(&resampled); let mut consumed = 0usize; while pending.len() - consumed >= OPUS_FRAME_SAMPLES { check_cancelled(job)?; let frame = &pending[consumed..consumed + OPUS_FRAME_SAMPLES]; let length = opus_encoder.encode_float(frame, &mut encoded) .map_err(|error| MediaError::new("encode_failed", error.to_string()))?; writer.write_packet(&encoded[..length], produced_samples)?; produced_samples += OPUS_FRAME_SAMPLES as u64; consumed += OPUS_FRAME_SAMPLES; } if consumed > 0 { pending.drain(..consumed); } if expected_duration_ms > 0 { let current_ms = produced_samples.saturating_mul(1000) / u64::from(OUTPUT_SAMPLE_RATE); job.publish(MediaJobEvent { phase: MediaJobPhase::Transcoding, progress: Some((current_ms as f64 / expected_duration_ms as f64).clamp(0.0, 0.99)), result: None, error: None, }); } } if !pending.is_empty() { pending.resize(OPUS_FRAME_SAMPLES, 0.0); let length = opus_encoder.encode_float(&pending, &mut encoded) .map_err(|error| MediaError::new("encode_failed", error.to_string()))?; writer.write_packet(&encoded[..length], produced_samples)?; produced_samples += OPUS_FRAME_SAMPLES as u64; } writer.finish()?; check_cancelled(job)?; Ok(MediaJobResult { output_path: output_path.to_string_lossy().into_owned(), detected_format, detected_codec, duration_ms: produced_samples.saturating_mul(1000) / u64::from(OUTPUT_SAMPLE_RATE), pass_through: false, }) } fn downmix_to_mono(samples: &[f32], channels: usize) -> Vec { if channels <= 1 { return samples.to_vec(); } samples.chunks_exact(channels).map(|frame| frame.iter().copied().sum::() / channels as f32).collect() } fn resample_chunk(samples: &[f32], input_rate: u32) -> Result, MediaError> { if input_rate == OUTPUT_SAMPLE_RATE || samples.is_empty() { return Ok(samples.to_vec()); } let input_data = vec![samples.to_vec()]; let input = SequentialSliceOfVecs::new(&input_data, 1, samples.len()) .map_err(|error| MediaError::new("resample_failed", error.to_string()))?; let mut resampler = Fft::::new(input_rate as usize, OUTPUT_SAMPLE_RATE as usize, samples.len().max(256), 1, FixedSync::Input) .map_err(|error| MediaError::new("resample_failed", error.to_string()))?; let output = resampler.process_all(&input, samples.len(), None) .map_err(|error| MediaError::new("resample_failed", error.to_string()))?; let mut result = Vec::with_capacity(output.frames()); for frame in 0..output.frames() { result.push(output.read_sample(0, frame).unwrap_or(0.0)); } Ok(result) } struct WebmOpusWriter { writer: WebmWriter, cluster_start_ms: Option, } impl WebmOpusWriter { fn new(file: File) -> Result { let mut writer = WebmWriter::new(file); write_tags(&mut writer, &[ MatroskaSpec::Ebml(Master::Start), MatroskaSpec::EbmlVersion(1), MatroskaSpec::EbmlReadVersion(1), MatroskaSpec::EbmlMaxIdLength(4), MatroskaSpec::EbmlMaxSizeLength(8), MatroskaSpec::DocType("webm".to_string()), MatroskaSpec::DocTypeVersion(4), MatroskaSpec::DocTypeReadVersion(2), MatroskaSpec::Ebml(Master::End), ])?; writer.write_advanced(&MatroskaSpec::Segment(Master::Start), WriteOptions::is_unknown_sized_element()).map_err(webm_error)?; write_tags(&mut writer, &[ MatroskaSpec::Info(Master::Start), MatroskaSpec::TimestampScale(1_000_000), MatroskaSpec::MuxingApp("MindWork AI Studio".to_string()), MatroskaSpec::WritingApp("MindWork AI Studio".to_string()), MatroskaSpec::Info(Master::End), MatroskaSpec::Tracks(Master::Start), MatroskaSpec::TrackEntry(Master::Start), MatroskaSpec::TrackNumber(1), MatroskaSpec::TrackUID(1), MatroskaSpec::TrackType(2), MatroskaSpec::FlagDefault(1), MatroskaSpec::CodecID("A_OPUS".to_string()), MatroskaSpec::CodecPrivate(opus_head()), MatroskaSpec::CodecDelay(u64::from(OPUS_PRE_SKIP) * 1_000_000_000 / u64::from(OUTPUT_SAMPLE_RATE)), MatroskaSpec::SeekPreRoll(80_000_000), MatroskaSpec::Audio(Master::Start), MatroskaSpec::SamplingFrequency(f64::from(OUTPUT_SAMPLE_RATE)), MatroskaSpec::Channels(1), MatroskaSpec::Audio(Master::End), MatroskaSpec::TrackEntry(Master::End), MatroskaSpec::Tracks(Master::End), ])?; Ok(Self { writer, cluster_start_ms: None }) } fn write_packet(&mut self, packet: &[u8], sample_position: u64) -> Result<(), MediaError> { let timestamp_ms = sample_position.saturating_mul(1000) / u64::from(OUTPUT_SAMPLE_RATE); let rotate = self.cluster_start_ms.map(|start| timestamp_ms.saturating_sub(start) >= CLUSTER_DURATION_MS).unwrap_or(true); if rotate { if self.cluster_start_ms.is_some() { self.writer.write(&MatroskaSpec::Cluster(Master::End)).map_err(webm_error)?; } self.writer.write(&MatroskaSpec::Cluster(Master::Start)).map_err(webm_error)?; self.writer.write(&MatroskaSpec::Timestamp(timestamp_ms)).map_err(webm_error)?; self.cluster_start_ms = Some(timestamp_ms); } let relative = timestamp_ms.saturating_sub(self.cluster_start_ms.unwrap_or(timestamp_ms)); if relative > i16::MAX as u64 { return Err(MediaError::new("webm_timestamp_overflow", "The WebM cluster timestamp exceeded its safe range.")); } let block: MatroskaSpec = SimpleBlock::new_uncheked(packet, 1, relative as i16, false, None, false, true).into(); self.writer.write(&block).map_err(webm_error) } fn finish(mut self) -> Result<(), MediaError> { if self.cluster_start_ms.is_some() { self.writer.write(&MatroskaSpec::Cluster(Master::End)).map_err(webm_error)?; } self.writer.write(&MatroskaSpec::Segment(Master::End)).map_err(webm_error)?; self.writer.into_inner().map_err(webm_error)?; Ok(()) } } fn write_tags(writer: &mut WebmWriter, tags: &[MatroskaSpec]) -> Result<(), MediaError> { for tag in tags { writer.write(tag).map_err(webm_error)?; } Ok(()) } fn opus_head() -> Vec { let mut data = b"OpusHead".to_vec(); data.push(1); data.push(1); data.extend_from_slice(&OPUS_PRE_SKIP.to_le_bytes()); data.extend_from_slice(&OUTPUT_SAMPLE_RATE.to_le_bytes()); data.extend_from_slice(&0i16.to_le_bytes()); data.push(0); data } fn partial_path(output_path: &FilePath) -> PathBuf { let mut name = output_path.file_name().unwrap_or_default().to_os_string(); name.push(".partial"); output_path.with_file_name(name) } fn is_webm_container(path: &FilePath) -> bool { let Ok(file) = File::open(path) else { return false; }; WebmIterator::new(file, &[]).take(16).filter_map(Result::ok).any(|tag| { matches!(tag, MatroskaSpec::DocType(doc_type) if doc_type.eq_ignore_ascii_case("webm")) }) } fn check_cancelled(job: &MediaJob) -> Result<(), MediaError> { if job.is_cancelled() { Err(MediaError::new("cancelled", "The media job was cancelled.")) } else { Ok(()) } } fn map_probe_error(error: SymphoniaError) -> MediaError { match error { SymphoniaError::Unsupported(_) => MediaError::new("unsupported_container", "This media container or codec is not supported."), _ => MediaError::new("damaged_container", format!("The media container could not be read: {error}")), } } fn webm_error(error: impl std::fmt::Display) -> MediaError { MediaError::new("webm_write_failed", format!("The WebM output could not be written: {error}")) } #[cfg(test)] mod tests { use super::*; use std::io::Cursor; #[test] fn opus_head_describes_48_khz_mono() { let head = opus_head(); assert_eq!(&head[..8], b"OpusHead"); assert_eq!(head[9], 1); assert_eq!(u32::from_le_bytes(head[12..16].try_into().unwrap()), 48_000); } #[test] fn webm_writer_rotates_clusters_before_relative_timestamp_overflow() { let path = std::env::temp_dir().join(format!("ai-studio-media-writer-{}.webm", rand::random::())); let file = File::create(&path).unwrap(); let mut writer = WebmOpusWriter::new(file).unwrap(); writer.write_packet(&[0xf8, 0xff, 0xfe], 0).unwrap(); writer.write_packet(&[0xf8, 0xff, 0xfe], 31 * 48_000).unwrap(); writer.finish().unwrap(); let bytes = fs::read(&path).unwrap(); let clusters = WebmIterator::new(Cursor::new(bytes), &[]) .filter_map(Result::ok) .filter(|tag| matches!(tag, MatroskaSpec::Cluster(Master::Start))) .count(); let _ = fs::remove_file(path); assert_eq!(clusters, 2); } #[test] fn downmix_is_bounded_and_balanced() { assert_eq!(downmix_to_mono(&[1.0, -1.0, 0.5, 0.5], 2), vec![0.0, 0.5]); } #[test] fn wav_is_normalized_to_one_mono_opus_track_with_frame_bounded_duration() { let directory = std::env::temp_dir().join(format!("ai-studio-media-test-{}", rand::random::())); fs::create_dir_all(&directory).unwrap(); let input = directory.join("input.wav"); let output = directory.join("output.webm"); fs::write(&input, wav_silence(44_100, 4_410)).unwrap(); let job = MediaJob::new(); let result = normalize_media(&input, &output, DEFAULT_MAX_PASS_THROUGH_BYTES, &job).unwrap(); assert!(!result.pass_through); assert!(result.duration_ms.abs_diff(100) <= 20); let file = File::open(&output).unwrap(); let tags: Vec<_> = WebmIterator::new(file, &[]) .filter_map(Result::ok) .collect(); assert_eq!(tags.iter().filter(|tag| matches!(tag, MatroskaSpec::TrackEntry(Master::Start))).count(), 1); assert!(tags.iter().any(|tag| matches!(tag, MatroskaSpec::CodecID(codec) if codec == "A_OPUS"))); assert!(tags.iter().any(|tag| matches!(tag, MatroskaSpec::Channels(1)))); assert!(tags.iter().any(|tag| matches!(tag, MatroskaSpec::SamplingFrequency(rate) if *rate == 48_000.0))); let _ = fs::remove_dir_all(directory); } #[test] fn cancellation_does_not_leave_an_output_file() { let directory = std::env::temp_dir().join(format!("ai-studio-media-cancel-{}", rand::random::())); fs::create_dir_all(&directory).unwrap(); let input = directory.join("input.wav"); let output = directory.join("output.webm"); fs::write(&input, wav_silence(48_000, 960)).unwrap(); let job = MediaJob::new(); job.cancelled.store(true, Ordering::Relaxed); let error = normalize_media(&input, &output, DEFAULT_MAX_PASS_THROUGH_BYTES, &job).unwrap_err(); assert_eq!(error.code, "cancelled"); assert!(!output.exists()); assert!(!partial_path(&output).exists()); let _ = fs::remove_dir_all(directory); } #[test] fn suitable_audio_only_webm_opus_is_passed_through() { let directory = std::env::temp_dir().join(format!("ai-studio-media-pass-through-{}", rand::random::())); fs::create_dir_all(&directory).unwrap(); let input = directory.join("input.webm"); let output = directory.join("output.webm"); let mut writer = WebmOpusWriter::new(File::create(&input).unwrap()).unwrap(); writer.write_packet(&[0xf8, 0xff, 0xfe], 0).unwrap(); writer.finish().unwrap(); let job = MediaJob::new(); let result = normalize_media(&input, &output, DEFAULT_MAX_PASS_THROUGH_BYTES, &job).unwrap(); assert!(result.pass_through); assert_eq!(fs::read(input).unwrap(), fs::read(output).unwrap()); let _ = fs::remove_dir_all(directory); } fn wav_silence(sample_rate: u32, samples: u32) -> Vec { let data_size = samples * 2; let mut wav = Vec::with_capacity(44 + data_size as usize); wav.extend_from_slice(b"RIFF"); wav.extend_from_slice(&(36 + data_size).to_le_bytes()); wav.extend_from_slice(b"WAVEfmt "); wav.extend_from_slice(&16u32.to_le_bytes()); wav.extend_from_slice(&1u16.to_le_bytes()); wav.extend_from_slice(&1u16.to_le_bytes()); wav.extend_from_slice(&sample_rate.to_le_bytes()); wav.extend_from_slice(&(sample_rate * 2).to_le_bytes()); wav.extend_from_slice(&2u16.to_le_bytes()); wav.extend_from_slice(&16u16.to_le_bytes()); wav.extend_from_slice(b"data"); wav.extend_from_slice(&data_size.to_le_bytes()); wav.resize(44 + data_size as usize, 0); wav } }