From 490297e61812ac66277e6f8c72ec9eb05bb9ad6d Mon Sep 17 00:00:00 2001
From: Ahmed Charles <acharles@outlook.com>
Date: Mon, 19 Feb 2024 09:01:28 -0800
Subject: [PATCH] fix: indefinite strings can't contain indefinite strings

Signed-off-by: Ahmed Charles <acharles@outlook.com>
---
 ciborium-ll/src/seg.rs      | 56 +++++++++++++++++++++++++++++++------
 ciborium/tests/error.rs     |  4 +--
 ciborium/tests/recursion.rs |  4 +--
 3 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/ciborium-ll/src/seg.rs b/ciborium-ll/src/seg.rs
index 95eaac9..5377e07 100644
--- a/ciborium-ll/src/seg.rs
+++ b/ciborium-ll/src/seg.rs
@@ -158,14 +158,20 @@ impl<'r, R: Read, P: Parser> Segment<'r, R, P> {
     }
 }
 
+#[derive(Eq, PartialEq)]
+enum State {
+    Initial,
+    Continue,
+    Finished,
+}
+
 /// A sequence of CBOR segments
 ///
 /// CBOR allows for bytes or text items to be segmented. This type represents
 /// the state of that segmented input stream.
 pub struct Segments<'r, R: Read, P: Parser> {
     reader: &'r mut Decoder<R>,
-    finish: bool,
-    nested: usize,
+    state: State,
     parser: PhantomData<P>,
     unwrap: fn(Header) -> Result<Option<usize>, ()>,
 }
@@ -178,8 +184,7 @@ impl<'r, R: Read, P: Parser> Segments<'r, R, P> {
     ) -> Self {
         Self {
             reader: decoder,
-            finish: false,
-            nested: 0,
+            state: State::Initial,
             parser: PhantomData,
             unwrap,
         }
@@ -190,16 +195,26 @@ impl<'r, R: Read, P: Parser> Segments<'r, R, P> {
     /// Returns `Ok(None)` at the conclusion of the stream.
     #[inline]
     pub fn pull(&mut self) -> Result<Option<Segment<R, P>>, Error<R::Error>> {
-        while !self.finish {
+        while self.state != State::Finished {
             let offset = self.reader.offset();
             match self.reader.pull()? {
-                Header::Break if self.nested == 1 => return Ok(None),
-                Header::Break if self.nested > 1 => self.nested -= 1,
+                Header::Break => {
+                    self.state = State::Finished;
+                    return Ok(None);
+                }
                 header => match (self.unwrap)(header) {
                     Err(..) => return Err(Error::Syntax(offset)),
-                    Ok(None) => self.nested += 1,
+                    Ok(None) => {
+                        if self.state == State::Initial {
+                            self.state = State::Continue;
+                        } else {
+                            return Err(Error::Syntax(offset));
+                        }
+                    }
                     Ok(Some(len)) => {
-                        self.finish = self.nested == 0;
+                        if self.state == State::Initial {
+                            self.state = State::Finished;
+                        }
                         return Ok(Some(Segment {
                             reader: self.reader,
                             unread: len,
@@ -214,3 +229,26 @@ impl<'r, R: Read, P: Parser> Segments<'r, R, P> {
         Ok(None)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn segments() {
+        fn t(data: &[u8], len: usize) {
+            let mut dec = Decoder::from(data);
+            let mut segs = Segments::<_, Bytes>::new(&mut dec, |header| match header {
+                Header::Bytes(len) => Ok(len),
+                _ => Err(()),
+            });
+            while let Some(mut seg) = segs.pull().unwrap() {
+                let mut b = [0; 1];
+                assert_eq!(Some(&b"0"[..]), seg.pull(&mut b).unwrap());
+            }
+            assert_eq!(len, dec.offset());
+        }
+        t(b"\x410\x00", 2);
+        t(b"\x5f\x410\xff\x00", 4);
+    }
+}
diff --git a/ciborium/tests/error.rs b/ciborium/tests/error.rs
index 589ab94..2bc42d9 100644
--- a/ciborium/tests/error.rs
+++ b/ciborium/tests/error.rs
@@ -131,8 +131,8 @@ fn eof() -> Error<std::io::Error> {
     case("7f4100ff", Error::Syntax(1)),
 
     // Indefinite-length string chunks not definite length:
-    //case("5f5f4100ffff", Error::Syntax(0)), These should fail, but do not currently.
-    //case("7f7f6100ffff", Error::Syntax(0)),
+    case("5f5f4100ffff", Error::Syntax(1)),
+    case("7f7f6100ffff", Error::Syntax(1)),
 
     // Break occurring on its own outside of an indefinite-length item:
     case("ff", Error::Semantic(None, "invalid type: break, expected non-break".into())),
diff --git a/ciborium/tests/recursion.rs b/ciborium/tests/recursion.rs
index cda1ce2..90a9692 100644
--- a/ciborium/tests/recursion.rs
+++ b/ciborium/tests/recursion.rs
@@ -33,7 +33,7 @@ fn map() {
 fn bytes() {
     let bytes = [0x5f; 128 * 1024];
     match from_reader::<Value, _>(&bytes[..]).unwrap_err() {
-        Error::Io(..) => (),
+        Error::Syntax(..) => (),
         e => panic!("incorrect error: {:?}", e),
     }
 }
@@ -42,7 +42,7 @@ fn bytes() {
 fn text() {
     let bytes = [0x7f; 128 * 1024];
     match from_reader::<Value, _>(&bytes[..]).unwrap_err() {
-        Error::Io(..) => (),
+        Error::Syntax(..) => (),
         e => panic!("incorrect error: {:?}", e),
     }
 }