Improve EOF processing for potential re-reading

Author: Evgeny Khramtsov

src/reading.rs

@@ -754,32 +754,63 @@ impl<T :io::Read + io::Seek> PacketReader<T> {
 	///
 	/// Ok(None) is returned if the physical stream has ended.
 	pub fn read_packet(&mut self) -> Result<Option<Packet>, 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<Packet, OggReadError> {
+		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<Option<Packet>, 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)
 		loop {
 			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<Packet, OggReadError> {
-		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<T :io::Read + io::Seek> PacketReader<T> {
 		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));
 

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
+	));
+}