Repository: duesee/imap-codec (language: rust)

Issue:
docs: Document `Fragmentizer`
> Should we document the `Fragmentizer` somewhere?
>
> - `README.md`?
> - `lib.rs`?

_Originally posted by @jakoschiko in https://github.com/duesee/imap-codec/issues/537#issuecomment-2223958444_
            


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 0623b420..b7faff58 100644
--- a/README.md
+++ b/README.md
@@ -54,6 +54,28 @@ $ cargo run --example=parse_command
 
 ... to parse some IMAP messages.
 
+### Robust parsing
+
+Real-world implementations need to handle invalid messages.
+Thus, `imap-codec` provides a `Fragmentizer` (see `fragmentizer*` examples) for "two-stage parsing".
+
+The idea is to detect message boundaries (consuming raw lines and literals) first, and to apply an IMAP message parser
+second.
+If the message parser errors out, we can safely discard the erroneous message w/o affecting other messages because we
+know the message boundaries.
+
+You can try this out by running ...
+
+```sh
+$ cargo run --example=fragmentizer_server
+```
+
+... and connecting via ...
+
+```sh
+$ netcat -C 127.0.0.1 12345
+```
+
 ### Tokio demo
 
 You can also start the [demo server] with ...
diff --git a/imap-codec/examples/fragmentizer.rs b/imap-codec/examples/fragmentizer.rs
new file mode 100644
index 00000000..6a91e465
--- /dev/null
+++ b/imap-codec/examples/fragmentizer.rs
@@ -0,0 +1,42 @@
+use std::io::{stdin, Read};
+
+use imap_codec::{fragmentizer::Fragmentizer, imap_types::utils::escape_byte_string};
+
+fn main() {
+    let mut fragmentizer = Fragmentizer::new(1024);
+
+    loop {
+        match fragmentizer.progress() {
+            Some(fragment_info) => {
+                println!(
+                    "[!] Fragment: {fragment_info:#?} // b\"{}\"",
+                    escape_byte_string(fragmentizer.fragment_bytes(fragment_info))
+                );
+
+                if fragmentizer.is_message_complete() {
+                    println!(
+                        "[!] Complete message: {}",
+                        escape_byte_string(fragmentizer.message_bytes())
+                    );
+                }
+            }
+            None => {
+                println!("[!] Reading stdin (ctrl+d to flush)...");
+                let mut buffer = [0; 64];
+                let count = stdin().read(&mut buffer).unwrap();
+                if count == 0 {
+                    println!("[!] Connection closed");
+                    break;
+                }
+                let chunk = &buffer[..count];
+
+                println!(
+                    "[!] Enqueueing {} byte(s) (b\"{}\")",
+                    count,
+                    escape_byte_string(chunk)
+                );
+                fragmentizer.enqueue_bytes(chunk);
+            }
+        }
+    }
+}
diff --git a/imap-codec/src/fragmentizer.rs b/imap-codec/src/fragmentizer.rs
index de7c914b..644e04c2 100644
--- a/imap-codec/src/fragmentizer.rs
+++ b/imap-codec/src/fragmentizer.rs
@@ -1,5 +1,47 @@
-//! Parser for IMAP message fragments.
-
+//! Utilities to split IMAP bytes into line and literal fragments.
+//!
+//! These utilities can be used to fragment a stream of IMAP bytes into lines (with metadata) and
+//! literals (before actually doing detailed IMAP parsing).
+//!
+//! This approach has multiple advantages: It separates literal handling from IMAP parsing and sets
+//! clear message boundaries even in the presence of malformed messages. Consequently, malformed
+//! messages can be reliably discarded. (A naive implementation of byte discardment may lead to
+//! adventurous (security) issues, such as, literal data being interpreted as command or response.)
+//! Further, this two-layered approach allows to more easily guard against excessive memory
+//! allocation by malevolant actors.
+//!
+//! # Example
+//!
+//! ```rust,ignore
+//! # use std::io::{stdin, Read};
+//! #
+//! # use imap_codec::{
+//! #     fragmentizer::Fragmentizer,
+//! #     imap_types::utils::escape_byte_string,
+//! # };
+//! #
+//! # fn read_bytes() -> &'static [u8] { b"" }
+//! #
+//! # fn main() {
+//! let mut fragmentizer = Fragmentizer::new(1024);
+//!
+//! loop {
+//!     match fragmentizer.progress() {
+//!         Some(fragment_info) => {
+//!             let fragment_bytes = fragmentizer.fragment_bytes(fragment_info);
+//!
+//!             if fragmentizer.is_message_complete() {
+//!                 let message_bytes = fragmentizer.message_bytes();
+//!             }
+//!         }
+//!         None => {
+//!             let received = read_bytes();
+//!             fragmentizer.enqueue_bytes(received);
+//!         }
+//!     }
+//! }
+//! # }
+//! ```
 use std::{collections::VecDeque, ops::Range};
 
 use imap_types::{
@@ -9,91 +51,15 @@ use imap_types::{
 
 use crate::decode::Decoder;
 
-/// The character sequence used for ending a line.
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub enum LineEnding {
-    /// The line ends with the character `\n`.
-    Lf,
-    /// The line ends with the character sequence `\r\n`.
-    CrLf,
-}
-
-/// Used by a line to announce a literal following the line.
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub struct LiteralAnnouncement {
-    /// The mode of the announced literal.
-    pub mode: LiteralMode,
-    /// The length of the announced literal in bytes.
-    pub length: u32,
-}
-
-/// Describes a fragment of the current message found by [`Fragmentizer::progress`].
-///
-/// The corresponding bytes can be retrieved via [`Fragmentizer::fragment_bytes`]
-/// until [`Fragmentizer::is_message_complete`] returns true. After that the
-/// next call of [`Fragmentizer::progress`] will start the next message.
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub enum FragmentInfo {
-    /// The fragment is a line.
-    Line {
-        /// Inclusive start index relative to the current message.
-        start: usize,
-        /// Exclusive end index relative to the current message.
-        end: usize,
-        /// Whether the next fragment will be a literal.
-        announcement: Option<LiteralAnnouncement>,
-        /// The detected ending sequence for this line.
-        ending: LineEnding,
-    },
-    /// The fragment is a literal.
-    Literal {
-        /// Inclusive start index relative to the current message.
-        start: usize,
-        /// Exclusive end index relative to the current message.
-        end: usize,
-    },
-}
-
-impl FragmentInfo {
-    /// The index range relative to the current message.
-    pub fn range(self) -> Range<usize> {
-        match self {
-            FragmentInfo::Line { start, end, .. } => start..end,
-            FragmentInfo::Literal { start, end } => start..end,
-        }
-    }
-}
-
-/// An error returned by [`Fragmentizer::decode_message`].
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub enum DecodeMessageError<'a, C: Decoder> {
-    /// The decoder failed decoding the message.
-    DecodingFailure(C::Error<'a>),
-    /// Not all bytes of the message were used when decoding the message.
-    DecodingRemainder {
-        /// The decoded message.
-        message: C::Message<'a>,
-        /// The unused bytes.
-        remainder: Secret<&'a [u8]>,
-    },
-    /// Max message size was exceeded and bytes were dropped.
-    MessageTooLong { initial: Secret<&'a [u8]> },
-}
-
-/// Parser for IMAP message fragments.
+/// Splits IMAP bytes into line and literal fragments.
 ///
-/// A fragment is either a line or a literal. A single IMAP message may consist of multiple
-/// fragments.
+/// The `Fragmentizer` prevents excessive memory allocation through a configurable maximum message size.
+/// Correct fragmentation is ensured even for messages exceeding the allowed message size.
 ///
-/// Using this fragment parser has the following benefits:
-/// - It can reliable detect when all fragments of a message were received. Only then it makes
-/// sense to decode the actual message. This allows us to decode large messages efficiently.
-/// - Even if the message is malformed it allows us to drop the right number of bytes. This
-/// prevents us from treating untrusted bytes (e.g. literal bytes) as IMAP messages.
-/// - It has support for max message length. Any exceeding bytes will be dropped. This
-/// prevents an attacker from allocating indefinite amount of memory.
-/// - It allows us to handle literals before the message is complete. This is important for IMAP
-/// servers because they need to accept or reject literals from the client.
+/// If the message size is exceeded,
+/// [`Fragmentizer::decode_message`] will fail and
+/// [`Fragmentizer::message_bytes`] will emit truncated message bytes.
+/// However, fragmentation will seamlessly continue with the following message.
 #[derive(Clone, Debug)]
 pub struct Fragmentizer {
     /// Enqueued bytes that are not parsed by [`Fragmentizer::progress`] yet.
@@ -111,14 +77,9 @@ pub struct Fragmentizer {
 }
 
 impl Fragmentizer {
-    /// Creates a `Fragmentizer` with maximum message size.
+    /// Creates a `Fragmentizer` with a maximum message size.
     ///
-    /// The maximum message size is bounded by `max_message_size` preventing excessive memory allocation.
-    ///
-    /// Correct fragmentation is ensured even for messages exceeding the allowed maximum message size.
-    /// If the message size is exceeded,
-    /// [`Fragmentizer::decode_message`] will fail and
-    /// [`Fragmentizer::message_bytes`] will contain truncated message bytes.
+    /// The maximum message size is bounded by `max_message_size` to prevent excessive memory allocation.
     pub fn new(max_message_size: u32) -> Self {
         Self {
             unparsed_buffer: VecDeque::new(),
@@ -129,7 +90,7 @@ impl Fragmentizer {
         }
     }
 
-    /// Creates a `Fragmentizer` without maximum message size.
+    /// Creates a `Fragmentizer` without a maximum message size.
     ///
     /// <div class="warning">
     /// This is dangerous because it allows an attacker to allocate an excessive amount of memory
@@ -145,65 +106,6 @@ impl Fragmentizer {
         }
     }
 
-    /// Enqueues more byte that can be parsed by [`Fragmentizer::progress`].
-    ///
-    /// Note that the message size limit is not enforced on the enqueued bytes. You can control
-    /// the size of the enqueued bytes by only calling this function if more bytes are necessary.
-    /// More bytes are necessary if [`Fragmentizer::progress`] returns `None`.
-    pub fn enqueue_bytes(&mut self, bytes: &[u8]) {
-        self.unparsed_buffer.extend(bytes);
-    }
-
-    /// Returns whether the size limit is exceeded for the current message.
-    pub fn is_max_message_size_exceeded(&self) -> bool {
-        self.max_message_size_exceeded
-    }
-
-    /// Returns whether the current message was fully parsed.
-    ///
-    /// If it returns true then it makes sense to call [`Fragmentizer::decode_message`]
-    /// to decode the message. Alternatively, you can access all bytes of the message via
-    /// [`Fragmentizer::message_bytes`].
-    pub fn is_message_complete(&self) -> bool {
-        self.parser.is_none()
-    }
-
-    /// Returns the bytes of the current message.
-    ///
-    /// Note that the bytes might be incomplete:
-    /// - The message might not be fully parsed yet and [`Fragmentizer::progress`] need to be
-    ///   called. You can check whether the message is complete via
-    ///   [`Fragmentizer::is_message_complete`].
-    /// - The size limit might be exceeded and bytes might be dropped. You can check this
-    ///   via [`Fragmentizer::is_max_message_size_exceeded`]
-    pub fn message_bytes(&self) -> &[u8] {
-        &self.message_buffer
-    }
-
-    /// Returns the bytes for a fragment of the current message.
-    pub fn fragment_bytes(&self, fragment_info: FragmentInfo) -> &[u8] {
-        let (start, end) = match fragment_info {
-            FragmentInfo::Line { start, end, .. } => (start, end),
-            FragmentInfo::Literal { start, end } => (start, end),
-        };
-        let start = start.min(self.message_buffer.len());
-        let end = end.min(self.message_buffer.len());
-        &self.message_buffer[start..end]
-    }
-
-    /// Skips the current message and starts the next message immediately.
-    ///
-    /// Warning: Using this method might be dangerous. If client and server don't
-    /// agree at which point a message is skipped, then the client or server might
-    /// treat untrusted bytes (e.g. literal bytes) as IMAP messages. Currently the
-    /// only valid use-case is a server that rejects synchronizing literals from the
-    /// client.
-    pub fn skip_message(&mut self) {
-        self.max_message_size_exceeded = false;
-        self.message_buffer.clear();
-        self.parser = Some(Parser::Line(LineParser::new(0)));
-    }
-
     /// Continue parsing the current message until the next fragment is detected.
     ///
     /// Returns `None` if more bytes need to be enqueued via [`Fragmentizer::enqueue_bytes`].
@@ -251,6 +153,65 @@ impl Fragmentizer {
         fragment
     }
 
+    /// Enqueues more byte that can be parsed by [`Fragmentizer::progress`].
+    ///
+    /// Note that the message size limit is not enforced on the enqueued bytes. You can control
+    /// the size of the enqueued bytes by only calling this function if more bytes are necessary.
+    /// More bytes are necessary if [`Fragmentizer::progress`] returns `None`.
+    pub fn enqueue_bytes(&mut self, bytes: &[u8]) {
+        self.unparsed_buffer.extend(bytes);
+    }
+
+    /// Returns the bytes for a fragment of the current message.
+    pub fn fragment_bytes(&self, fragment_info: FragmentInfo) -> &[u8] {
+        let (start, end) = match fragment_info {
+            FragmentInfo::Line { start, end, .. } => (start, end),
+            FragmentInfo::Literal { start, end } => (start, end),
+        };
+        let start = start.min(self.message_buffer.len());
+        let end = end.min(self.message_buffer.len());
+        &self.message_buffer[start..end]
+    }
+
+    /// Returns whether the current message was fully parsed.
+    ///
+    /// If it returns true then it makes sense to call [`Fragmentizer::decode_message`]
+    /// to decode the message. Alternatively, you can access all bytes of the message via
+    /// [`Fragmentizer::message_bytes`].
+    pub fn is_message_complete(&self) -> bool {
+        self.parser.is_none()
+    }
+
+    /// Returns the bytes of the current message.
+    ///
+    /// Note that the bytes might be incomplete:
+    /// - The message might not be fully parsed yet and [`Fragmentizer::progress`] need to be
+    ///   called. You can check whether the message is complete via
+    ///   [`Fragmentizer::is_message_complete`].
+    /// - The size limit might be exceeded and bytes might be dropped. You can check this
+    ///   via [`Fragmentizer::is_max_message_size_exceeded`]
+    pub fn message_bytes(&self) -> &[u8] {
+        &self.message_buffer
+    }
+
+    /// Returns whether the size limit is exceeded for the current message.
+    pub fn is_max_message_size_exceeded(&self) -> bool {
+        self.max_message_size_exceeded
+    }
+
+    /// Skips the current message and starts the next message immediately.
+    ///
+    /// Warning: Using this method might be dangerous. If client and server don't
+    /// agree at which point a message is skipped, then the client or server might
+    /// treat untrusted bytes (e.g. literal bytes) as IMAP messages. Currently the
+    /// only valid use-case is a server that rejects synchronizing literals from the
+    /// client.
+    pub fn skip_message(&mut self) {
+        self.max_message_size_exceeded = false;
+        self.message_buffer.clear();
+        self.parser = Some(Parser::Line(LineParser::new(0)));
+    }
+
     /// Tries to decode the [`Tag`] for the current message.
     ///
     /// Note that decoding the [`Tag`] is on best effort basis. Some message types don't have
@@ -533,6 +494,77 @@ impl LiteralParser {
     }
 }
 
+/// Describes a fragment of the current message found by [`Fragmentizer::progress`].
+///
+/// The corresponding bytes can be retrieved via [`Fragmentizer::fragment_bytes`]
+/// until [`Fragmentizer::is_message_complete`] returns true. After that the
+/// next call of [`Fragmentizer::progress`] will start the next message.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum FragmentInfo {
+    /// The fragment is a line.
+    Line {
+        /// Inclusive start index relative to the current message.
+        start: usize,
+        /// Exclusive end index relative to the current message.
+        end: usize,
+        /// Whether the next fragment will be a literal.
+        announcement: Option<LiteralAnnouncement>,
+        /// The detected ending sequence for this line.
+        ending: LineEnding,
+    },
+    /// The fragment is a literal.
+    Literal {
+        /// Inclusive start index relative to the current message.
+        start: usize,
+        /// Exclusive end index relative to the current message.
+        end: usize,
+    },
+}
+
+impl FragmentInfo {
+    /// The index range relative to the current message.
+    pub fn range(self) -> Range<usize> {
+        match self {
+            FragmentInfo::Line { start, end, .. } => start..end,
+            FragmentInfo::Literal { start, end } => start..end,
+        }
+    }
+}
+
+/// Used by a line to announce a literal following the line.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct LiteralAnnouncement {
+    /// The mode of the announced literal.
+    pub mode: LiteralMode,
+    /// The length of the announced literal in bytes.
+    pub length: u32,
+}
+
+/// The character sequence used for ending a line.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum LineEnding {
+    /// The line ends with the character `\n`.
+    Lf,
+    /// The line ends with the character sequence `\r\n`.
+    CrLf,
+}
+
+/// An error returned by [`Fragmentizer::decode_message`].
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum DecodeMessageError<'a, C: Decoder> {
+    /// The decoder failed decoding the message.
+    DecodingFailure(C::Error<'a>),
+    /// Not all bytes of the message were used when decoding the message.
+    DecodingRemainder {
+        /// The decoded message.
+        message: C::Message<'a>,
+        /// The unused bytes.
+        remainder: Secret<&'a [u8]>,
+    },
+    /// Max message size was exceeded and bytes were dropped.
+    MessageTooLong { initial: Secret<&'a [u8]> },
+}
+
 fn parse_tag(message_bytes: &[u8]) -> Option<Tag> {
     let mut bytes = message_bytes.iter().enumerate();
     let sp = loop {
@@ -1221,27 +1253,17 @@ mod tests {
     #[test]
     fn parse_line_examples() {
         assert_not_line(b"");
-
         assert_not_line(b"foo");
 
         assert_is_line(b"\n", 1, None, LineEnding::Lf);
-
         assert_is_line(b"\r\n", 2, None, LineEnding::CrLf);
-
         assert_is_line(b"\n\r", 1, None, LineEnding::Lf);
-
         assert_is_line(b"foo\n", 4, None, LineEnding::Lf);
-
         assert_is_line(b"foo\r\n", 5, None, LineEnding::CrLf);
-
         assert_is_line(b"foo\n\r", 4, None, LineEnding::Lf);
-
         assert_is_line(b"foo\nbar\n", 4, None, LineEnding::Lf);
-
         assert_is_line(b"foo\r\nbar\r\n", 5, None, LineEnding::CrLf);
-
         assert_is_line(b"\r\nfoo\r\n", 2, None, LineEnding::CrLf);
-
         assert_is_line(
             b"{1}\r\n",
             5,
@@ -1251,7 +1273,6 @@ mod tests {
             }),
             LineEnding::CrLf,
         );
-
         assert_is_line(
             b"{1}\n",
             4,
@@ -1261,7 +1282,6 @@ mod tests {
             }),
             LineEnding::Lf,
         );
-
         assert_is_line(
             b"foo {1}\r\n",
             9,
@@ -1271,7 +1291,6 @@ mod tests {
             }),
             LineEnding::CrLf,
         );
-
         assert_is_line(
             b"foo {2} {1}\r\n",
             13,
@@ -1281,19 +1300,12 @@ mod tests {
             }),
             LineEnding::CrLf,
         );
-
         assert_is_line(b"foo {1} \r\n", 10, None, LineEnding::CrLf);
-
         assert_is_line(b"foo \n {1}\r\n", 5, None, LineEnding::Lf);
-
         assert_is_line(b"foo {1} foo\r\n", 13, None, LineEnding::CrLf);
-
         assert_is_line(b"foo {1\r\n", 8, None, LineEnding::CrLf);
-
         assert_is_line(b"foo 1}\r\n", 8, None, LineEnding::CrLf);
-
         assert_is_line(b"foo { 1}\r\n", 10, None, LineEnding::CrLf);
-
         assert_is_line(
             b"foo {{1}\r\n",
             10,
@@ -1303,7 +1315,6 @@ mod tests {
             }),
             LineEnding::CrLf,
         );
-
         assert_is_line(
             b"foo {42}\r\n",
             10,
@@ -1313,7 +1324,6 @@ mod tests {
             }),
             LineEnding::CrLf,
         );
-
         assert_is_line(
             b"foo {42+}\r\n",
             11,
@@ -1323,7 +1333,6 @@ mod tests {
             }),
             LineEnding::CrLf,
         );
-
         assert_is_line(
             b"foo +{42}\r\n",
             11,
@@ -1333,21 +1342,13 @@ mod tests {
             }),
             LineEnding::CrLf,
         );
-
         assert_is_line(b"foo {+}\r\n", 9, None, LineEnding::CrLf);
-
         assert_is_line(b"foo {42++}\r\n", 12, None, LineEnding::CrLf);
-
         assert_is_line(b"foo {+42+}\r\n", 12, None, LineEnding::CrLf);
-
         assert_is_line(b"foo {+42}\r\n", 11, None, LineEnding::CrLf);
-
         assert_is_line(b"foo {42}+\r\n", 11, None, LineEnding::CrLf);
-
         assert_is_line(b"foo {-42}\r\n", 11, None, LineEnding::CrLf);
-
         assert_is_line(b"foo {42-}\r\n", 11, None, LineEnding::CrLf);
-
         assert_is_line(
             b"foo {4294967295}\r\n",
             18,
@@ -1357,7 +1358,6 @@ mod tests {
             }),
             LineEnding::CrLf,
         );
-
         assert_is_line(b"foo {4294967296}\r\n", 18, None, LineEnding::CrLf);
     }
 
diff --git a/imap-codec/src/lib.rs b/imap-codec/src/lib.rs
index 45f9f892..b9e594d7 100644
--- a/imap-codec/src/lib.rs
+++ b/imap-codec/src/lib.rs
@@ -22,11 +22,13 @@
 //! ### Example
 //!
 //! ```rust
-//! // Run 'cargo add imap-types' in your project's root directory.
-//! use imap_codec::{decode::Decoder, GreetingCodec};
-//! use imap_types::{
-//!     core::Text,
-//!     response::{Code, Greeting, GreetingKind},
+//! use imap_codec::{
+//!     decode::Decoder,
+//!     imap_types::{
+//!         core::Text,
+//!         response::{Code, Greeting, GreetingKind},
+//!     },
+//!     GreetingCodec,
 //! };
 //!
 //! let (remaining, greeting) = GreetingCodec::default()
@@ -54,11 +56,13 @@
 //! ### Example
 //!
 //! ```rust
-//! // Run 'cargo add imap-types' in your project's root directory.
-//! use imap_codec::{encode::Encoder, GreetingCodec};
-//! use imap_types::{
-//!     core::Text,
-//!     response::{Code, Greeting, GreetingKind},
+//! use imap_codec::{
+//!     encode::Encoder,
+//!     imap_types::{
+//!         core::Text,
+//!         response::{Code, Greeting, GreetingKind},
+//!     },
+//!     GreetingCodec,
 //! };
 //!
 //! let greeting = Greeting {


Test patch: