From 6c0a38742ab295470521984619b369c9ac4582a7 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 14 Jul 2026 18:54:14 +0000 Subject: [PATCH] Fix WAV chunk ordering and padding --- src/web/audio_utils.rs | 182 ++++++++++++++++++++++++++++++++++++++--- 1 file changed, 172 insertions(+), 10 deletions(-) diff --git a/src/web/audio_utils.rs b/src/web/audio_utils.rs index 21065b58..464a84d5 100644 --- a/src/web/audio_utils.rs +++ b/src/web/audio_utils.rs @@ -115,7 +115,8 @@ impl AudioSplitter { return Err("Not a valid WAV file".to_string()); } - // Find both fmt and data chunks - they can be in any order + // Find both mandatory chunks. RIFF/WAVE requires fmt to precede data, + // while optional and unknown chunks may appear around them. let mut pos = 12; let mut data_start = 0; let mut data_size = 0; @@ -135,6 +136,9 @@ impl AudioSplitter { fmt_pos = pos + 8; fmt_size = chunk_size; } else if chunk_id == b"data" { + if fmt_pos == 0 { + return Err("Invalid WAV file: fmt chunk must precede data chunk".to_string()); + } data_start = pos + 8; data_size = chunk_size; } @@ -189,12 +193,18 @@ impl AudioSplitter { if fmt_size < 16 { return Err("WAV fmt chunk too small".to_string()); } - if fmt_pos + fmt_size > audio_data.len() { + let fmt_end = fmt_pos + .checked_add(fmt_size) + .ok_or_else(|| "Invalid WAV file: fmt chunk size would cause overflow".to_string())?; + if fmt_end > audio_data.len() { return Err("WAV fmt chunk exceeds file length".to_string()); } // Validate that the data chunk size doesn't exceed the file length - if data_start + data_size > audio_data.len() { + let data_end = data_start + .checked_add(data_size) + .ok_or_else(|| "Invalid WAV file: data chunk size would cause overflow".to_string())?; + if data_end > audio_data.len() { return Err("WAV file data chunk exceeds file length".to_string()); } @@ -204,10 +214,10 @@ impl AudioSplitter { } let data_header_pos = data_start - 8; let header = &audio_data[0..data_header_pos]; - let audio_samples = &audio_data[data_start..data_start + data_size]; + let audio_samples = &audio_data[data_start..data_end]; // Parse format info from fmt chunk - let fmt_data = &audio_data[fmt_pos..fmt_pos + fmt_size]; + let fmt_data = &audio_data[fmt_pos..fmt_end]; let channels = u16::from_le_bytes([fmt_data[2], fmt_data[3]]); let sample_rate = u32::from_le_bytes([fmt_data[4], fmt_data[5], fmt_data[6], fmt_data[7]]); let bytes_per_second = @@ -228,7 +238,10 @@ impl AudioSplitter { ); // Calculate chunk sizes with checked math and alignment - let header_overhead = header.len() + 8; // Header and data chunk header + let header_overhead = header + .len() + .checked_add(8) // Header and data chunk header + .ok_or_else(|| "WAV file header size would cause overflow".to_string())?; let available = self .max_chunk_size .checked_sub(header_overhead) @@ -240,7 +253,20 @@ impl AudioSplitter { })?; // Align to sample boundaries and ensure at least one frame - let chunk_duration_bytes = (available / block_align) * block_align; + let mut chunk_duration_bytes = (available / block_align) + .checked_mul(block_align) + .ok_or_else(|| "WAV sample-aligned chunk size would cause overflow".to_string())?; + + // RIFF chunks are word-aligned. Reserve a padding byte when a full + // sample window would otherwise end at the size limit with odd length. + let padded_chunk_duration_bytes = chunk_duration_bytes + .checked_add(chunk_duration_bytes % 2) + .ok_or_else(|| "WAV padded chunk size would cause overflow".to_string())?; + if padded_chunk_duration_bytes > available { + chunk_duration_bytes = chunk_duration_bytes + .checked_sub(block_align) + .ok_or_else(|| "WAV chunk size is smaller than one audio frame".to_string())?; + } if chunk_duration_bytes < block_align { return Err( "Computed chunk size is smaller than one audio frame; increase max_chunk_size" @@ -312,20 +338,43 @@ impl AudioSplitter { original_header: &[u8], samples: &[u8], ) -> Result, String> { - let mut wav = Vec::new(); + let padding_size = samples.len() % 2; + let output_size = original_header + .len() + .checked_add(8) + .and_then(|size| size.checked_add(samples.len())) + .and_then(|size| size.checked_add(padding_size)) + .ok_or_else(|| "WAV output size would cause overflow".to_string())?; + if output_size > self.max_chunk_size { + return Err(format!( + "WAV output exceeds maximum chunk size: {output_size} bytes (max: {} bytes)", + self.max_chunk_size + )); + } + + let data_size = u32::try_from(samples.len()) + .map_err(|_| "WAV data chunk exceeds the RIFF size limit".to_string())?; + let riff_size = output_size + .checked_sub(8) + .and_then(|size| u32::try_from(size).ok()) + .ok_or_else(|| "WAV output exceeds the RIFF size limit".to_string())?; + + let mut wav = Vec::with_capacity(output_size); // Copy the original header up to the data chunk wav.extend_from_slice(original_header); // Update the data chunk size wav.extend_from_slice(b"data"); - wav.extend_from_slice(&(samples.len() as u32).to_le_bytes()); + wav.extend_from_slice(&data_size.to_le_bytes()); // Add the actual audio samples wav.extend_from_slice(samples); + if padding_size == 1 { + wav.push(0); + } // Update the RIFF chunk size (file size - 8) - let riff_size = (wav.len() - 8) as u32; wav[4..8].copy_from_slice(&riff_size.to_le_bytes()); Ok(wav) @@ -398,6 +447,74 @@ pub fn merge_transcriptions( mod tests { use super::*; + const PCM_U8_MONO_FMT: [u8; 16] = [ + 1, 0, // PCM + 1, 0, // mono + 0x40, 0x1f, 0, 0, // 8 kHz + 0x40, 0x1f, 0, 0, // 8,000 bytes/second + 1, 0, // one-byte frames + 8, 0, // eight bits/sample + ]; + + fn wav_with_chunks(chunks: &[(&[u8; 4], &[u8])]) -> Vec { + let mut wav = Vec::from(&b"RIFF\0\0\0\0WAVE"[..]); + + for (id, data) in chunks { + wav.extend_from_slice(*id); + wav.extend_from_slice(&(data.len() as u32).to_le_bytes()); + wav.extend_from_slice(data); + if data.len() % 2 == 1 { + wav.push(0); + } + } + + let riff_size = (wav.len() - 8) as u32; + wav[4..8].copy_from_slice(&riff_size.to_le_bytes()); + wav + } + + fn parse_wav_chunks(wav: &[u8]) -> Vec<([u8; 4], Vec)> { + assert!(wav.len() >= 12); + assert_eq!(&wav[0..4], b"RIFF"); + assert_eq!(&wav[8..12], b"WAVE"); + + let riff_size = u32::from_le_bytes(wav[4..8].try_into().unwrap()) as usize; + assert_eq!(riff_size.checked_add(8), Some(wav.len())); + + let mut chunks = Vec::new(); + let mut pos = 12; + let mut saw_fmt = false; + let mut saw_data = false; + + while pos < wav.len() { + let header_end = pos.checked_add(8).unwrap(); + assert!(header_end <= wav.len()); + + let id: [u8; 4] = wav[pos..pos + 4].try_into().unwrap(); + let size = u32::from_le_bytes(wav[pos + 4..header_end].try_into().unwrap()) as usize; + let data_end = header_end.checked_add(size).unwrap(); + assert!(data_end <= wav.len()); + + if &id == b"fmt " { + saw_fmt = true; + } else if &id == b"data" { + assert!(saw_fmt, "fmt chunk must precede data chunk"); + saw_data = true; + } + + chunks.push((id, wav[header_end..data_end].to_vec())); + pos = data_end.checked_add(size % 2).unwrap(); + assert!(pos <= wav.len()); + if size % 2 == 1 { + assert_eq!(wav[pos - 1], 0); + } + } + + assert!(saw_fmt); + assert!(saw_data); + chunks + } + #[test] fn test_should_split() { let splitter = AudioSplitter::new(); @@ -461,6 +578,51 @@ mod tests { assert!(result.unwrap_err().contains("Not a valid WAV")); } + #[test] + fn test_split_wav_rejects_data_before_fmt() { + let splitter = AudioSplitter::new(); + let wav = wav_with_chunks(&[(b"data", &[1, 2, 3, 4]), (b"fmt ", &PCM_U8_MONO_FMT)]); + + assert_eq!( + splitter.split_wav(&wav).unwrap_err(), + "Invalid WAV file: fmt chunk must precede data chunk" + ); + } + + #[test] + fn test_split_wav_emits_valid_word_aligned_chunks_with_optional_chunks() { + let splitter = AudioSplitter { + max_chunk_size: 71, + overlap_seconds: 0.0, + }; + let samples = [1, 2, 3, 4, 5, 6, 7, 8, 9]; + let wav = wav_with_chunks(&[ + (b"JUNK", &[10, 11, 12]), + (b"fmt ", &PCM_U8_MONO_FMT), + (b"LIST", &[13]), + (b"data", &samples), + ]); + + let chunks = splitter.split_wav(&wav).unwrap(); + assert!(chunks.len() > 1); + + let mut emitted_samples = Vec::new(); + for chunk in chunks { + assert!(chunk.data.len() <= splitter.max_chunk_size); + let parsed = parse_wav_chunks(&chunk.data); + assert_eq!( + parsed.iter().map(|(id, _)| *id).collect::>(), + vec![*b"JUNK", *b"fmt ", *b"LIST", *b"data"] + ); + assert_eq!(parsed[0].1, [10, 11, 12]); + assert_eq!(parsed[1].1, PCM_U8_MONO_FMT); + assert_eq!(parsed[2].1, [13]); + emitted_samples.extend_from_slice(&parsed[3].1); + } + + assert_eq!(emitted_samples, samples); + } + #[test] fn test_merge_transcriptions() { let results = vec![