test: more

torrust · Oct 16, 2024 · 7b2a4df · 7b2a4df
1 parent f7065d3
commit 7b2a4df
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/src/parsers/error.rs b/src/parsers/error.rs
@@ -23,6 +23,9 @@ pub enum Error {
     #[error("R/W error: {0}")]
     Rw(#[from] rw::error::Error),
 
+    #[error("Read byte after peeking does match peeked byte; {0}")]
+    ReadByteAfterPeekingDoesMatchPeekedByte(ReadContext),
+
     #[error("Unrecognized first byte for new bencoded value; {0}")]
     UnrecognizedFirstBencodeValueByte(ReadContext),
 

diff --git a/src/parsers/mod.rs b/src/parsers/mod.rs
@@ -112,24 +112,24 @@ impl<R: Read> BencodeParser<R> {
                     string::parse(&mut self.byte_reader, writer)?;
                 }
                 BENCODE_BEGIN_LIST => {
-                    let _byte = Self::consume_byte(&mut self.byte_reader)?;
+                    let _byte = Self::read_peeked_byte(peeked_byte, &mut self.byte_reader)?;
                     self.begin_bencoded_value(BencodeType::List, writer)?;
                     writer.write_byte(Self::JSON_ARRAY_BEGIN)?;
                     self.stack.push(State::ExpectingFirstListItemOrEnd);
                 }
                 BENCODE_BEGIN_DICT => {
-                    let _byte = Self::consume_byte(&mut self.byte_reader)?;
+                    let _byte = Self::read_peeked_byte(peeked_byte, &mut self.byte_reader)?;
                     self.begin_bencoded_value(BencodeType::Dict, writer)?;
                     writer.write_byte(Self::JSON_OBJ_BEGIN)?;
                     self.stack.push(State::ExpectingFirstDictFieldOrEnd);
                 }
                 BENCODE_END_LIST_OR_DICT => {
-                    let _byte = Self::consume_byte(&mut self.byte_reader)?;
+                    let _byte = Self::read_peeked_byte(peeked_byte, &mut self.byte_reader)?;
                     self.end_list_or_dict(writer)?;
                 }
                 b'\n' => {
                     // Ignore line breaks at the beginning, the end, or between values
-                    let _byte = Self::consume_byte(&mut self.byte_reader)?;
+                    let _byte = Self::read_peeked_byte(peeked_byte, &mut self.byte_reader)?;
                 }
                 _ => {
                     return Err(error::Error::UnrecognizedFirstBencodeValueByte(
@@ -202,10 +202,27 @@ impl<R: Read> BencodeParser<R> {
     ///
     /// # Errors
     ///
-    /// Will return and errors if it can't read from the input.
-    fn consume_byte(reader: &mut ByteReader<R>) -> Result<Option<u8>, error::Error> {
+    /// Will return and errors if:
+    ///
+    /// - It can't read from the input.
+    /// - The byte read is not the expected one (the previously peeked byte).
+    fn read_peeked_byte(
+        peeked_byte: u8,
+        reader: &mut ByteReader<R>,
+    ) -> Result<Option<u8>, error::Error> {
         match reader.read_byte() {
-            Ok(byte) => Ok(Some(byte)),
+            Ok(byte) => {
+                if byte == peeked_byte {
+                    return Ok(Some(byte));
+                }
+                Err(error::Error::ReadByteAfterPeekingDoesMatchPeekedByte(
+                    ReadContext {
+                        byte: Some(byte),
+                        pos: reader.input_byte_counter(),
+                        latest_bytes: reader.captured_bytes(),
+                    },
+                ))
+            }
             Err(err) => {
                 if err.kind() == io::ErrorKind::UnexpectedEof {
                     return Ok(None);