From edcc9eedd72967b6c0dbc63c150f8b58455daecb Mon Sep 17 00:00:00 2001
From: Mohammad AlSaleh <CE.Mohammad.AlSaleh@gmail.com>
Date: Fri, 15 Nov 2024 16:10:09 +0300
Subject: [PATCH] Add the ability to ignore length in the reader

 This is the reader counterpart to `into_header_for_infinite_file()`.

Signed-off-by: Mohammad AlSaleh <CE.Mohammad.AlSaleh@gmail.com>
---
 src/lib.rs  | 77 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/read.rs | 25 +++++++++++++++--
 2 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 244d4ea..109cef9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1026,3 +1026,80 @@ fmt \x28\x00\x00\x00\xfe\xff\x0a\x00\x80\x3e\x00\x00\x00\xe2\x04\x00\
 \x00\x00\x10\x00\x80\x00\x00\xaa\x00\x38\x9b\x71\
 data\xFF\xFF\xFF\xFF"[..]);
 }
+
+#[cfg(test)]
+#[test]
+fn test_ignore_length() {
+    macro_rules! test {
+        (no_ignore, $ty:ty, $buf:expr) => {{
+            let mut out_reader = WavReader::new(std::io::Cursor::new($buf)).unwrap();
+            let out_samples_res = out_reader
+                .samples()
+                .collect::<crate::Result<Vec<$ty>>>();
+            assert!{{
+                if let Err(crate::Error::IoError(e)) = out_samples_res {
+                    e.kind() == std::io::ErrorKind::UnexpectedEof
+                } else {
+                    false
+                }
+            }};
+        }};
+        (ignore, $buf:expr, $in_samples:expr) => {{
+            let mut out_reader = WavReader::new(std::io::Cursor::new($buf)).unwrap();
+            out_reader.set_ignore_length(true);
+            let out_samples = out_reader
+                .samples()
+                .collect::<crate::Result<Vec<_>>>()
+                .unwrap();
+            assert_eq!($in_samples, out_samples);
+        }};
+        ($ty:ty, $files:expr) => {
+            for fname in $files {
+                let mut reader = WavReader::open(fname).unwrap();
+                let spec = reader.spec();
+                let samples = reader
+                    .samples()
+                    .collect::<crate::Result<Vec<$ty>>>()
+                    .unwrap();
+                let mut buf = spec.into_header_for_infinite_file();
+                samples.iter().for_each(|sample| {
+                    sample.write(&mut buf, spec.bits_per_sample).unwrap();
+                });
+                // expect EOF errors
+                test!(no_ignore, $ty, buf.clone());
+                // expect success and getting back exactly the samples written
+                test!(ignore, buf, samples);
+            };
+        }
+    }
+    let files_s16 = &[
+        "testsamples/pcmwaveformat-16bit-44100Hz-mono.wav",
+        "testsamples/pcmwaveformat-16bit-44100Hz-mono-extra.wav",
+        "testsamples/waveformatex-16bit-44100Hz-mono.wav",
+        "testsamples/waveformatex-16bit-44100Hz-mono-extra.wav",
+        "testsamples/waveformatex-16bit-44100Hz-stereo.wav",
+    ];
+
+    test!{i16, files_s16};
+
+    let files_u8 = &[
+        "testsamples/pcmwaveformat-8bit-44100Hz-mono.wav",
+        "testsamples/waveformatex-8bit-11025Hz-mono.wav",
+    ];
+
+    test!{i16, files_u8};
+
+    let files_s24 = &[
+        "testsamples/pcmwaveformat-24bit-4byte-48kHz-stereo.wav",
+        "testsamples/waveformatextensible-24bit-192kHz-mono.wav",
+        "testsamples/waveformatextensible-24bit-4byte-48kHz-stereo.wav",
+    ];
+
+    test!{i32, files_s24};
+
+    let files_s32 = &[
+        "testsamples/waveformatextensible-32bit-48kHz-stereo.wav",
+    ];
+
+    test!{i32, files_s32};
+}
diff --git a/src/read.rs b/src/read.rs
index 7024947..58ba7f3 100644
--- a/src/read.rs
+++ b/src/read.rs
@@ -83,7 +83,7 @@ impl<R> ReadExt for R
             if progress > 0 {
                 n += progress;
             } else {
-                return Err(io::Error::new(io::ErrorKind::Other, "Failed to read enough bytes."));
+                return Err(io::ErrorKind::UnexpectedEof)?;
             }
         }
         Ok(())
@@ -285,6 +285,8 @@ pub enum Chunk<'r, R: 'r + io::Read> {
 pub struct ChunksReader<R: io::Read> {
     /// the underlying reader
     reader: R,
+    /// Don't err if EOF is reached before length
+    ignore_length: bool,
     /// the Wave format specification, if it has been read already
     pub spec_ex: Option<WavSpecEx>,
     /// when inside the main data state, keeps track of decoding and chunk
@@ -311,6 +313,7 @@ impl<R: io::Read> ChunksReader<R> {
         try!(read_wave_header(&mut reader));
         Ok(ChunksReader {
             reader: reader,
+            ignore_length: false,
             spec_ex: None,
             data_state: None,
         })
@@ -801,6 +804,21 @@ impl<R> WavReader<R>
         })
     }
 
+    /// Set whether to ignore header length.
+    ///
+    /// Length can be set to a maximum value that represents infinite length in
+    /// non-seekable sinks (like stdout).
+    ///
+    /// Setting this to `true` allows ignoring that length.
+    pub fn set_ignore_length(&mut self, ignore_length: bool) {
+        self.reader.ignore_length = ignore_length;
+    }
+
+    /// Get whether header length is ignored or not.
+    pub fn ignore_length(&self) -> bool {
+        self.reader.ignore_length
+    }
+
     /// Returns information about the WAVE file.
     pub fn spec(&self) -> WavSpec {
         self.reader.spec_ex
@@ -900,7 +918,10 @@ fn iter_next<R, S>(reader: &mut ChunksReader<R>) -> Option<Result<S>>
                                   data.spec_ex.spec.sample_format,
                                   data.spec_ex.bytes_per_sample,
                                   data.spec_ex.spec.bits_per_sample);
-        Some(sample.map_err(Error::from))
+        match sample {
+            Err(Error::IoError(e)) if e.kind() == io::ErrorKind::UnexpectedEof && reader.ignore_length => None,
+            _ => Some(sample.map_err(Error::from)),
+        }
     } else {
         None
     }