Fix jitter restart and WAV size overflow

src/audio/recorder.rs

@@ -22,6 +22,8 @@ use std::path::{Path, PathBuf};
 const SAMPLE_RATE: u32 = 48_000;
 const BITS_PER_SAMPLE: u16 = 16;
 const CHANNELS: u16 = 1;
+const RIFF_DATA_OVERHEAD: u64 = 36;
+const MAX_RIFF_DATA_BYTES: u64 = u32::MAX as u64 - RIFF_DATA_OVERHEAD;
 
 /// Cap on buffered mic samples (~200ms). Bounds how far recording lag can drift
 /// if the capture clock runs persistently faster than playout — past this we drop
@@ -34,7 +36,7 @@ const MAX_MIC_FIFO: usize = SAMPLE_RATE as usize / 5;
 pub struct WavWriter {
     file: File,
     /// Bytes of PCM data written so far (for the size fields).
-    data_bytes: u32,
+    data_bytes: u64,
 }
 
 impl WavWriter {
@@ -42,7 +44,10 @@ impl WavWriter {
     pub fn new(path: &Path) -> io::Result<Self> {
         let mut file = File::create(path)?;
         file.write_all(&Self::header(0))?;
-        Ok(Self { file, data_bytes: 0 })
+        Ok(Self {
+            file,
+            data_bytes: 0,
+        })
     }
 
     /// The 44-byte canonical WAV/PCM header for the given data length in bytes.
@@ -68,21 +73,40 @@ impl WavWriter {
 
     /// Append PCM samples to the data chunk.
     pub fn write_samples(&mut self, samples: &[i16]) -> io::Result<()> {
+        let added_bytes = u64::try_from(samples.len())
+            .ok()
+            .and_then(|len| len.checked_mul(2))
+            .ok_or_else(|| io::Error::other("WAV sample buffer too large"))?;
+        let new_data_bytes = self
+            .data_bytes
+            .checked_add(added_bytes)
+            .ok_or_else(|| io::Error::other("WAV data size overflow"))?;
+        if new_data_bytes > MAX_RIFF_DATA_BYTES {
+            return Err(io::Error::other("WAV too large for RIFF"));
+        }
+
         let mut buf = Vec::with_capacity(samples.len() * 2);
         for &s in samples {
             buf.extend_from_slice(&s.to_le_bytes());
         }
         self.file.write_all(&buf)?;
-        self.data_bytes += (samples.len() * 2) as u32;
+        self.data_bytes = new_data_bytes;
         Ok(())
     }
 
     /// Patch the RIFF + data size fields and flush. Consumes the writer.
     pub fn finalize(mut self) -> io::Result<()> {
+        let data_bytes = u32::try_from(self.data_bytes)
+            .map_err(|_| io::Error::other("WAV too large for RIFF"))?;
+        let riff_size = self
+            .data_bytes
+            .checked_add(RIFF_DATA_OVERHEAD)
+            .and_then(|size| u32::try_from(size).ok())
+            .ok_or_else(|| io::Error::other("WAV too large for RIFF"))?;
         self.file.seek(SeekFrom::Start(4))?;
-        self.file.write_all(&(36 + self.data_bytes).to_le_bytes())?;
+        self.file.write_all(&riff_size.to_le_bytes())?;
         self.file.seek(SeekFrom::Start(40))?;
-        self.file.write_all(&self.data_bytes.to_le_bytes())?;
+        self.file.write_all(&data_bytes.to_le_bytes())?;
         self.file.flush()?;
         Ok(())
     }
@@ -209,11 +233,29 @@ mod tests {
         let _ = std::fs::remove_file(&path);
     }
 
+    #[test]
+    fn wav_writer_rejects_data_that_would_overflow_riff_header() {
+        let dir = std::env::temp_dir();
+        let path = dir.join(format!("peerspeak-overflow-{}.wav", std::process::id()));
+        let mut w = WavWriter::new(&path).unwrap();
+        w.data_bytes = MAX_RIFF_DATA_BYTES - 1;
+        let before_len = std::fs::metadata(&path).unwrap().len();
+
+        let err = w.write_samples(&[0]).unwrap_err();
+
+        assert_eq!(err.kind(), io::ErrorKind::Other);
+        assert_eq!(w.data_bytes, MAX_RIFF_DATA_BYTES - 1);
+        assert_eq!(std::fs::metadata(&path).unwrap().len(), before_len);
+        drop(w);
+        let _ = std::fs::remove_file(&path);
+    }
+
     #[test]
     fn mic_is_summed_with_mix_when_present() {
         let dir = std::env::temp_dir();
         let mut r = Recorder {
-            writer: WavWriter::new(&dir.join(format!("ps-sum-{}.wav", std::process::id()))).unwrap(),
+            writer: WavWriter::new(&dir.join(format!("ps-sum-{}.wav", std::process::id())))
+                .unwrap(),
             mic_fifo: VecDeque::new(),
             path: PathBuf::new(),
         };
@@ -223,7 +265,11 @@ mod tests {
         r.write_frame(&[10, 20]).unwrap();
         assert_eq!(r.mic_fifo.len(), 1, "two samples consumed, one mic left");
         r.write_frame(&[0, 0]).unwrap();
-        assert_eq!(r.mic_fifo.len(), 0, "remaining mic sample consumed; rest is silence");
+        assert_eq!(
+            r.mic_fifo.len(),
+            0,
+            "remaining mic sample consumed; rest is silence"
+        );
         let _ = r.finalize();
     }
 
@@ -231,7 +277,8 @@ mod tests {
     fn mic_fifo_is_capped() {
         let dir = std::env::temp_dir();
         let mut r = Recorder {
-            writer: WavWriter::new(&dir.join(format!("ps-cap-{}.wav", std::process::id()))).unwrap(),
+            writer: WavWriter::new(&dir.join(format!("ps-cap-{}.wav", std::process::id())))
+                .unwrap(),
             mic_fifo: VecDeque::new(),
             path: PathBuf::new(),
         };