Skip to content

Fix WAV chunk ordering and RIFF padding #232

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AnthonyRonning wants to merge 1 commit into
base: master
Choose a base branch
from
+172 −10
Open

Fix WAV chunk ordering and RIFF padding #232

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
182 changes: 172 additions & 10 deletions src/web/audio_utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
}
Expand Down Expand Up @@ -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());
}

Expand All @@ -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 =
Expand All @@ -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)
Expand All @@ -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"
Expand Down Expand Up @@ -312,20 +338,43 @@ impl AudioSplitter {
original_header: &[u8],
samples: &[u8],
) -> Result<Vec<u8>, 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)
Expand Down Expand Up @@ -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<u8> {
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<u8>)> {
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();
Expand Down Expand Up @@ -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<_>>(),
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![
Expand Down
Loading