Merge pull request #537 from saadkhantareen/fix/464-poison-stream-encoders

fix(audio-io): poison stream encoders after a failed write so later calls cannot claim success (#464)

crates/audio-io/src/aiff.rs

@@ -185,6 +185,9 @@ pub(crate) struct Writer<W> {
     quantizers: Vec<Quantizer>,
     frames: u64,
     bytes: Vec<u8>,
+    /// Set when a block write has failed: the sink may hold a prefix of that block while
+    /// `frames` does not count it, so the writer refuses every later call.
+    failed: bool,
 }
 
 impl<W: Write + Seek> Writer<W> {
@@ -193,7 +196,7 @@ impl<W: Write + Seek> Writer<W> {
         let start = out.stream_position()?;
         out.write_all(&header(channels, sample_rate, depth, 0)?)?;
         let quantizers = Quantizer::per_channel(u32::from(sample_format(depth).0), dither, channels);
-        Ok(Writer { out, start, channels, sample_rate, depth, quantizers, frames: 0, bytes: Vec::new() })
+        Ok(Writer { out, start, channels, sample_rate, depth, quantizers, frames: 0, bytes: Vec::new(), failed: false })
     }
 
     /// The frame count after `frames` more, or an error when AIFF can't hold that many (4 GiB).
@@ -206,13 +209,22 @@ impl<W: Write + Seek> Writer<W> {
     /// Appends `frames` frames of planar audio (missing samples are silence), unless they don't
     /// fit (see [`Writer::check_room`]).
     pub(crate) fn write(&mut self, block: &[Vec<f32>], frames: usize) -> Result<()> {
+        if self.failed {
+            return Err(crate::failed_write_error());
+        }
         let total = self.check_room(frames as u64)?;
-        write_interleaved(&mut self.out, &mut self.bytes, block, frames, &mut self.quantizers, sample_format(self.depth), true)?;
+        if let Err(e) = write_interleaved(&mut self.out, &mut self.bytes, block, frames, &mut self.quantizers, sample_format(self.depth), true) {
+            self.failed = true;
+            return Err(AudioError::from(e));
+        }
         self.frames = total;
         Ok(())
     }
 
     pub(crate) fn finish(mut self) -> Result<W> {
+        if self.failed {
+            return Err(crate::failed_write_error());
+        }
         let data_len = u64::from(sample_format(self.depth).0 / 8).saturating_mul(self.channels as u64).saturating_mul(self.frames);
         if data_len % 2 == 1 {
             self.out.write_all(&[0])?;

crates/audio-io/src/flac.rs

@@ -318,6 +318,10 @@ pub(crate) struct Writer<W> {
     pending: Vec<Vec<i64>>,
     frames: u64,
     frame_no: u64,
+    /// Set when a frame write has failed: `write_frame` clears the pending samples before
+    /// writing, so a failed frame is lost while `frames` still counts it; the writer refuses
+    /// every later call.
+    failed: bool,
 }
 
 impl<W: Write + Seek> Writer<W> {
@@ -336,6 +340,7 @@ impl<W: Write + Seek> Writer<W> {
             pending: (0..channels).map(|_| Vec::with_capacity(BLOCK_SIZE)).collect(),
             frames: 0,
             frame_no: 0,
+            failed: false,
         })
     }
 
@@ -351,6 +356,9 @@ impl<W: Write + Seek> Writer<W> {
     /// Appends `frames` frames of planar audio (missing samples are silence), unless they don't
     /// fit (see [`Writer::check_room`]).
     pub(crate) fn write(&mut self, block: &[Vec<f32>], frames: usize) -> Result<()> {
+        if self.failed {
+            return Err(crate::failed_write_error());
+        }
         let total = self.check_room(frames as u64)?;
         let mut done = 0;
         while done < frames {
@@ -361,8 +369,11 @@ impl<W: Write + Seek> Writer<W> {
                 p.extend((0..n).map(|i| i64::from(q.quantize(src.get(i).copied().unwrap_or(0.0)))));
             }
             done += n;
-            if held + n == BLOCK_SIZE {
-                self.write_frame()?;
+            if held + n == BLOCK_SIZE
+                && let Err(e) = self.write_frame()
+            {
+                self.failed = true;
+                return Err(e);
             }
         }
         self.frames = total;
@@ -408,6 +419,9 @@ impl<W: Write + Seek> Writer<W> {
     }
 
     pub(crate) fn finish(mut self) -> Result<W> {
+        if self.failed {
+            return Err(crate::failed_write_error());
+        }
         self.write_frame()?;
         let info = stream_info(self.channels, self.sample_rate, self.bps, self.frames);
         self.out.seek(SeekFrom::Start(self.start.saturating_add(4)))?;

crates/audio-io/src/lib.rs

@@ -294,6 +294,13 @@ pub fn encode_with_channel_mask(buf: &AudioBuffer, opts: &EncodeOptions, mask: u
     }
 }
 
+/// The error a stream encoder returns for every call after one of its writes has failed: the
+/// sink may already hold a prefix of the failed block, so the encoder's counters can no longer
+/// describe the bytes on the sink, and later writes or `finish` must not claim success.
+pub(crate) fn failed_write_error() -> AudioError {
+    AudioError::Encode("the stream encoder already failed on an earlier write; its output is incomplete".into())
+}
+
 /// Writes WAV/BWF (RF64 above 4 GiB), AIFF / AIFF-C `fl32` or FLAC block by block, producing
 /// the same bytes [`encode_with_channel_mask`] would for the whole buffer.
 pub struct FileEncoder<W: Write + Seek>(Writer<W>);
@@ -329,7 +336,9 @@ impl<W: Write + Seek> FileEncoder<W> {
     }
 
     /// Appends the first `frames` frames of `block`, one `Vec` per channel (missing samples are
-    /// silence), unless they don't fit (see [`FileEncoder::check_room`]).
+    /// silence), unless they don't fit (see [`FileEncoder::check_room`]). Once a write has
+    /// failed, the encoder is unusable: the sink may hold a prefix of the failed block, so
+    /// every later `write` and `finish` fails instead of completing an inconsistent file.
     pub fn write(&mut self, block: &[Vec<f32>], frames: usize) -> Result<()> {
         match &mut self.0 {
             Writer::Wav(w) => w.write(block, frames),

crates/audio-io/src/wav.rs

@@ -321,6 +321,9 @@ pub(crate) struct Writer<W> {
     quantizers: Vec<Quantizer>,
     frames: u64,
     bytes: Vec<u8>,
+    /// Set when a block write has failed: the sink may hold a prefix of that block while
+    /// `frames` does not count it, so the writer refuses every later call.
+    failed: bool,
 }
 
 impl<W: Write + Seek> Writer<W> {
@@ -339,7 +342,20 @@ impl<W: Write + Seek> Writer<W> {
         let start = out.stream_position()?;
         out.write_all(&header(channels, sample_rate, format, bwf, mask, 0, force_rf64)?)?;
         let quantizers = Quantizer::per_channel(u32::from(format.0), dither, channels);
-        Ok(Writer { out, start, channels, sample_rate, format, bwf: bwf.cloned(), mask, force_rf64, quantizers, frames: 0, bytes: Vec::new() })
+        Ok(Writer {
+            out,
+            start,
+            channels,
+            sample_rate,
+            format,
+            bwf: bwf.cloned(),
+            mask,
+            force_rf64,
+            quantizers,
+            frames: 0,
+            bytes: Vec::new(),
+            failed: false,
+        })
     }
 
     fn header(&self, frames: u64) -> Result<Vec<u8>> {
@@ -355,13 +371,22 @@ impl<W: Write + Seek> Writer<W> {
 
     /// Appends `frames` frames of planar audio (missing samples are silence).
     pub(crate) fn write(&mut self, block: &[Vec<f32>], frames: usize) -> Result<()> {
+        if self.failed {
+            return Err(crate::failed_write_error());
+        }
         let total = self.check_room(frames as u64)?;
-        write_interleaved(&mut self.out, &mut self.bytes, block, frames, &mut self.quantizers, self.format, false)?;
+        if let Err(e) = write_interleaved(&mut self.out, &mut self.bytes, block, frames, &mut self.quantizers, self.format, false) {
+            self.failed = true;
+            return Err(AudioError::from(e));
+        }
         self.frames = total;
         Ok(())
     }
 
     pub(crate) fn finish(mut self) -> Result<W> {
+        if self.failed {
+            return Err(crate::failed_write_error());
+        }
         let data_len = u64::from(self.format.0 / 8).saturating_mul(self.channels as u64).saturating_mul(self.frames);
         if data_len % 2 == 1 {
             self.out.write_all(&[0])?;

crates/audio-io/tests/stream.rs

@@ -1,5 +1,5 @@
 use soundcraft_audio_io::{AudioBuffer, AudioError, BitDepth, BwfInfo, EncodeOptions, FileEncoder, FileFormat, encode, encode_with_channel_mask};
-use std::io::Cursor;
+use std::io::{Cursor, Seek, SeekFrom, Write};
 
 fn tone(channels: usize, frames: usize) -> AudioBuffer {
     let channels = (0..channels).map(|c| (0..frames).map(|i| ((i * (c + 3)) as f32 * 0.013).sin() * 1.2).collect()).collect();
@@ -49,6 +49,50 @@ fn streamed_files_match_whole_buffer_encodes() {
     }
 }
 
+/// A `Write + Seek` sink that accepts `budget` bytes in total and then fails every write, so
+/// a `write_all` on it appends a prefix of the block and only then reports the error.
+struct FailAfter {
+    inner: Cursor<Vec<u8>>,
+    budget: usize,
+}
+
+impl Write for FailAfter {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        if self.budget == 0 {
+            return Err(std::io::Error::other("sink is full"));
+        }
+        let n = buf.len().min(self.budget);
+        self.budget -= n;
+        self.inner.write(&buf[..n])
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.inner.flush()
+    }
+}
+
+impl Seek for FailAfter {
+    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
+        self.inner.seek(pos)
+    }
+}
+
+#[test]
+fn a_failed_write_poisons_the_encoder() {
+    let buf = tone(2, 20_001);
+    for format in [FileFormat::Wav, FileFormat::Aiff, FileFormat::Flac] {
+        let opts = EncodeOptions { format, bit_depth: BitDepth::Int16, dither: false, bwf: None };
+        let sink = FailAfter { inner: Cursor::new(Vec::new()), budget: 256 };
+        let mut enc = FileEncoder::new(sink, &opts, 2, buf.sample_rate, None).unwrap();
+        // The sink accepts a prefix of the first block, then fails.
+        assert!(enc.write(&buf.channels, buf.frames()).is_err(), "{format:?}");
+        // Later calls must not claim success: the bytes on the sink no longer match the
+        // encoder's counters (and for FLAC the failed frame's samples are already gone).
+        assert!(matches!(enc.write(&buf.channels, 16), Err(AudioError::Encode(_))), "{format:?}");
+        assert!(matches!(enc.finish(), Err(AudioError::Encode(_))), "{format:?}");
+    }
+}
+
 #[test]
 fn blocks_past_a_format_limit_are_refused_before_anything_is_written() {
     let buf = tone(2, 100);