From 7a6dbd3ee3ab836f1d9a0d5bfff5035f6e8097b4 Mon Sep 17 00:00:00 2001 From: Evgeny Khramtsov Date: Sat, 29 Jun 2024 14:36:18 +0200 Subject: [PATCH] Improve EOF processing for potential re-reading --- src/reading.rs | 60 ++++++++++++++++++++++++++++++++++++++------------ src/test.rs | 47 +++++++++++++++++++++++++++++++++++---- 2 files changed, 89 insertions(+), 18 deletions(-) diff --git a/src/reading.rs b/src/reading.rs index 7bdf789..9e55213 100644 --- a/src/reading.rs +++ b/src/reading.rs @@ -754,6 +754,31 @@ impl PacketReader { /// /// Ok(None) is returned if the physical stream has ended. pub fn read_packet(&mut self) -> Result, OggReadError> { + self.do_read_packet(false) + } + + /// Reads a packet and returns it on success. + /// + /// Unlike the `read_packet` function, this function returns an `Err(_)` + /// if the physical stream has ended, with an error kind of + /// [`UnexpectedEof`](std::io::ErrorKind::UnexpectedEof). + /// + /// This function might be useful in the following scenarios: + /// 1. When you expect a new packet to arrive. + /// 2. When the underlying reader's buffer is incomplete, allowing you to add more data + /// and re-read upon encountering `UnexpectedEof`. + /// + /// Note: This function incurs negligible overhead by reading the current + /// stream position before reading an Ogg page. + pub fn read_packet_expected(&mut self) -> Result { + match tri!(self.do_read_packet(true)) { + Some(p) => Ok(p), + None => tri!(Err(Error::new(ErrorKind::UnexpectedEof, + "Expected ogg packet but found end of physical stream"))), + } + } + + fn do_read_packet(&mut self, restore_position: bool) -> Result, OggReadError> { // Read pages until we got a valid entire packet // (packets may span multiple pages, so reading one page // doesn't always suffice to give us a valid packet) @@ -761,25 +786,31 @@ impl PacketReader { if let Some(pck) = self.base_pck_rdr.read_packet() { return Ok(Some(pck)); } - let page = tri!(self.read_ogg_page()); + let page = if restore_position { + // Save the current position to rewind if EOF is reached. + let pos = tri!(self.rdr.stream_position()); + tri!(match self.read_ogg_page() { + Err(OggReadError::ReadError(e)) if e.kind() == ErrorKind::UnexpectedEof => { + // Rewind to the saved position to allow for potential re-reading. + self.rdr.seek(SeekFrom::Start(pos))?; + Err(OggReadError::ReadError(e)) + } + Ok(None) => { + // Rewind to the saved position to allow for potential re-reading. + self.rdr.seek(SeekFrom::Start(pos))?; + Ok(None) + } + ret => ret, + }) + } else { + tri!(self.read_ogg_page()) + }; match page { Some(page) => tri!(self.base_pck_rdr.push_page(page)), None => return Ok(None), } } } - /// Reads a packet, and returns it on success. - /// - /// The difference to the `read_packet` function is that this function - /// returns an Err(_) if the physical stream has ended. - /// This function is useful if you expect a new packet to come. - pub fn read_packet_expected(&mut self) -> Result { - match tri!(self.read_packet()) { - Some(p) => Ok(p), - None => tri!(Err(Error::new(ErrorKind::UnexpectedEof, - "Expected ogg packet but found end of physical stream"))), - } - } /// Reads until the new page header, and then returns the page header array. /// @@ -818,7 +849,8 @@ impl PacketReader { let header_buf :[u8; 27] = match tri!(self.read_until_pg_header()) { Some(s) => s, None if self.read_some_pg => return Ok(None), - None => return Err(OggReadError::NoCapturePatternFound) + None => tri!(Err(Error::new(ErrorKind::UnexpectedEof, + "Expected ogg packet but found end of physical stream"))), }; let (mut pg_prs, page_segments) = tri!(PageParser::new(header_buf)); diff --git a/src/test.rs b/src/test.rs index ff55ae7..2b37d9b 100644 --- a/src/test.rs +++ b/src/test.rs @@ -542,6 +542,8 @@ fn test_issue_14() { // data is treated as invalid. #[test] fn test_issue_7() { + use std::io::ErrorKind::UnexpectedEof; + let mut c = Cursor::new(Vec::new()); let test_arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; let test_arr_2 = [2, 4, 8, 16, 32, 64, 128, 127, 126, 125, 124]; @@ -564,17 +566,54 @@ fn test_issue_7() { assert!(r.read_packet().unwrap().is_none()); } - // Non-Ogg data should return an error. + // Truncated data should return the UnexpectedEof error. let c = Cursor::new(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); { let mut r = PacketReader::new(c); - assert!(matches!(r.read_packet(), Err(OggReadError::NoCapturePatternFound))); + assert!(matches!(r.read_packet(), + Err(OggReadError::ReadError(e)) if e.kind() == UnexpectedEof)); } - // Empty data is considered non-Ogg data. + // Empty data should return the UnexpectedEof error. let c = Cursor::new(&[]); { let mut r = PacketReader::new(c); - assert!(matches!(r.read_packet(), Err(OggReadError::NoCapturePatternFound))); + assert!(matches!(r.read_packet(), + Err(OggReadError::ReadError(e)) if e.kind() == UnexpectedEof)); } } + +#[test] +fn test_read_with_append() { + use std::io::ErrorKind::UnexpectedEof; + + let first_chunk = + [79, 103, 103, 83, 0, 2, 0, 0, 0, + 0, 0, 0, 0, 0, 14, 247, 95, 68, 0]; + let second_chunk = + [0, 0, 0, 139, 130, 226, 240, 1, 30, 1, 118, 111, 114, + 98, 105, 115, 0, 0, 0, 0, 1, 128, 187, 0, 0, 0, 0, 0, + 0, 128, 56, 1, 0, 0, 0, 0, 0, 184, 1]; + + let c = Cursor::new(Vec::new()); + let mut r = PacketReader::new(c); + + // Add the first chunk + let c = r.get_mut(); + c.get_mut().extend(first_chunk); + // Attempt to read an incomplete page + assert!(matches!(r.read_packet_expected(), + Err(OggReadError::ReadError(ref e)) if e.kind() == UnexpectedEof + )); + + // Add the second chunk + let c = r.get_mut(); + c.get_mut().extend(second_chunk); + // Attempt to read a complete page + assert!(matches!(r.read_packet_expected(), Ok(_))); + + // No data left in the buffer + assert!(matches!(r.read_packet_expected(), + Err(OggReadError::ReadError(ref e)) if e.kind() == UnexpectedEof + )); +}