Fix LIST/UIDL streaming parsing and reject zero message numbers

.gitignore

@@ -1,4 +1,6 @@
 /.env
 /target
 /.direnv
-/.pre-commit-config.yaml
+/.pre-commit-config.yaml
+/.DS_Store
+/docs/superpowers/

Cargo.toml

@@ -23,6 +23,7 @@ tokio = { version = "1.38.2", features = [
 	"time",
 	"rt",
 	"macros",
+	"io-util",
 ], optional = true }
 
 [dev-dependencies]

src/lib.rs

@@ -365,7 +365,10 @@ impl<S: Read + Write + Unpin + Send> Client<S> {
         let response = self.send_request(request).await?;
 
         match response {
-            Response::Message(resp) => Ok(resp),
+            Response::Message(resp) => {
+                self.marked_as_del.push(msg_number);
+                Ok(resp)
+            }
             _ => err!(
                 ErrorKind::UnexpectedResponse,
                 "Did not received the expected dele response"

src/request.rs

@@ -57,4 +57,8 @@ impl Request {
     pub fn command(&self) -> &Command {
         &self.command
     }
+
+    pub(crate) fn arg_count(&self) -> usize {
+        self.args.len()
+    }
 }

src/response/mod.rs

@@ -8,7 +8,7 @@ pub mod uidl;
 use bytes::Bytes;
 use nom::IResult;
 
-use crate::command::Command;
+use crate::{command::Command, request::Request};
 
 use self::{
     capability::Capability, list::List, stat::Stat, types::message::Text, uidl::UidlResponse,
@@ -29,6 +29,49 @@ impl Status {
     }
 }
 
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub(crate) enum ResponseShape {
+    Message,
+    Greeting,
+    Stat,
+    ListSingle,
+    ListMulti,
+    UidlSingle,
+    UidlMulti,
+    Bytes,
+    Capability,
+    #[cfg(feature = "sasl")]
+    Auth,
+}
+
+impl From<&Request> for ResponseShape {
+    fn from(request: &Request) -> Self {
+        match request.command() {
+            Command::Greet => Self::Greeting,
+            Command::Stat => Self::Stat,
+            Command::List => {
+                if request.arg_count() > 0 {
+                    Self::ListSingle
+                } else {
+                    Self::ListMulti
+                }
+            }
+            Command::Uidl => {
+                if request.arg_count() > 0 {
+                    Self::UidlSingle
+                } else {
+                    Self::UidlMulti
+                }
+            }
+            Command::Retr | Command::Top => Self::Bytes,
+            Command::Capa => Self::Capability,
+            #[cfg(feature = "sasl")]
+            Command::Auth | Command::Base64(_) => Self::Auth,
+            _ => Self::Message,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub enum Response {
     Stat(Stat),
@@ -46,4 +89,11 @@ impl Response {
     pub fn from_bytes<'a>(input: &'a [u8], command: &Command) -> IResult<&'a [u8], Self> {
         parser::parse(input, command)
     }
+
+    pub(crate) fn from_shape<'a>(
+        input: &'a [u8],
+        shape: &ResponseShape,
+    ) -> IResult<&'a [u8], Self> {
+        parser::parse_shape(input, shape)
+    }
 }

src/response/parser/mod.rs

@@ -10,13 +10,13 @@ use crate::command::Command;
 
 use self::{
     rfc1939::{
-        error_response, list_response, rfc822_response, stat_response, status, string_response,
-        uidl_list_response, uidl_response,
+        error_response, list_response, list_single_response, rfc822_response, stat_response,
+        status, string_response, uidl_list_response, uidl_response,
     },
     rfc2449::capability_response,
 };
 
-use super::Response;
+use super::{Response, ResponseShape};
 
 pub(crate) fn parse<'a>(input: &'a [u8], request: &Command) -> IResult<&'a [u8], Response> {
     if input.is_empty() {
@@ -43,7 +43,7 @@ pub(crate) fn parse<'a>(input: &'a [u8], request: &Command) -> IResult<&'a [u8],
         match request {
             Command::Stat => stat_response(input),
             Command::Uidl => alt((uidl_response, uidl_list_response))(input),
-            Command::List => alt((stat_response, list_response))(input),
+            Command::List => alt((list_single_response, list_response))(input),
             Command::Retr | Command::Top => rfc822_response(input),
             Command::Capa => capability_response(input),
             _ => string_response(input),
@@ -53,6 +53,47 @@ pub(crate) fn parse<'a>(input: &'a [u8], request: &Command) -> IResult<&'a [u8],
     }
 }
 
+pub(crate) fn parse_shape<'a>(
+    input: &'a [u8],
+    shape: &ResponseShape,
+) -> IResult<&'a [u8], Response> {
+    if input.is_empty() {
+        return Err(nom::Err::Incomplete(nom::Needed::Unknown));
+    }
+
+    #[cfg(feature = "sasl")]
+    if matches!(shape, ResponseShape::Auth) {
+        match rfc1734::auth(input) {
+            Ok((input, base64_challenge)) => {
+                if let Ok(challenge) = crate::base64::decode(base64_challenge) {
+                    return Ok((input, Response::Challenge(challenge.into())));
+                }
+            }
+            Err(nom::Err::Incomplete(needed)) => return Err(nom::Err::Incomplete(needed)),
+            Err(_) => {}
+        }
+    }
+
+    let (input, status) = status(input)?;
+
+    if status.success() {
+        match shape {
+            ResponseShape::Message | ResponseShape::Greeting => string_response(input),
+            ResponseShape::Stat => stat_response(input),
+            ResponseShape::ListSingle => list_single_response(input),
+            ResponseShape::ListMulti => list_response(input),
+            ResponseShape::UidlSingle => uidl_response(input),
+            ResponseShape::UidlMulti => uidl_list_response(input),
+            ResponseShape::Bytes => rfc822_response(input),
+            ResponseShape::Capability => capability_response(input),
+            #[cfg(feature = "sasl")]
+            ResponseShape::Auth => string_response(input),
+        }
+    } else {
+        error_response(input)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use crate::response::{types::DataType, uidl::UidlResponse};
@@ -117,7 +158,22 @@ mod test {
 
         let result = parse(data, &Command::List);
 
-        assert!(result.is_err())
+        assert!(result.is_err());
+
+        let data = b"+OK 1 120\r\n2 200\r\n.\r\n+OK goodbye\r\n";
+
+        let (output, response) = parse(data, &Command::List).unwrap();
+
+        assert_eq!(output, b"+OK goodbye\r\n");
+
+        match response {
+            Response::List(list) => {
+                assert_eq!(list.items().len(), 2);
+            }
+            _ => {
+                panic!("LIST multiline response was parsed as a single-line response");
+            }
+        }
     }
 
     #[test]
@@ -156,6 +212,63 @@ mod test {
                 unreachable!()
             }
         }
+
+        let data =
+            b"+OK 1 whqtswO00WBw418f9t5JxYwZ\r\n2 QhdPYR:00WBw1Ph7x7\r\n.\r\n+OK goodbye\r\n";
+
+        let (output, response) = parse(data, &Command::Uidl).unwrap();
+
+        assert_eq!(output, b"+OK goodbye\r\n");
+
+        match response {
+            Response::Uidl(uidl) => match uidl {
+                UidlResponse::Multiple(list) => {
+                    assert_eq!(list.items().len(), 2);
+                }
+                _ => {
+                    panic!("UIDL multiline response was parsed as a single-line response");
+                }
+            },
+            _ => {
+                unreachable!()
+            }
+        }
+
+        let data = b"+OK 1 whqtswO00WBw418f9t5JxYwZ\r\n.\r\n+OK goodbye\r\n";
+
+        let (output, response) = parse(data, &Command::Uidl).unwrap();
+
+        assert_eq!(output, b"+OK goodbye\r\n");
+
+        match response {
+            Response::Uidl(uidl) => {
+                match uidl {
+                    UidlResponse::Multiple(list) => {
+                        assert_eq!(list.items().len(), 1);
+                    }
+                    _ => {
+                        panic!("Single-item UIDL multiline response was parsed as a single-line response");
+                    }
+                }
+            }
+            _ => {
+                unreachable!()
+            }
+        }
+    }
+
+    #[test]
+    fn test_list_rejects_zero_message_number() {
+        let data = b"+OK 1 messages (120 bytes)\r\n0 120\r\n.\r\n";
+
+        assert!(parse(data, &Command::List).is_err());
+    }
+
+    #[test]
+    fn test_uidl_rejects_zero_message_number() {
+        let data = b"+OK unique-id listing follows\r\n0 baduid\r\n.\r\n";
+
+        assert!(parse(data, &Command::Uidl).is_err());
     }
 
     #[test]

src/response/parser/rfc1939.rs

@@ -1,12 +1,13 @@
 use bytes::Bytes;
+use nom::error::{Error as NomError, ErrorKind};
 use nom::{
     branch::alt,
     bytes::streaming::{tag, take_until, take_while, take_while_m_n},
     character::{
         is_alphanumeric,
         streaming::{char, digit1, not_line_ending, space0, space1},
     },
-    combinator::{map, opt, value},
+    combinator::{map, map_res, opt, value},
     multi::many_till,
     sequence::{delimited, preceded, terminated, tuple},
     IResult,
@@ -43,6 +44,29 @@ fn stat(input: &[u8]) -> IResult<&[u8], Stat> {
     Ok((input, Stat::new(count, size)))
 }
 
+fn message_number(input: &[u8]) -> IResult<&[u8], &[u8]> {
+    map_res(digit1, |bytes: &[u8]| {
+        let number = std::str::from_utf8(bytes)
+            .ok()
+            .and_then(|text| text.parse::<usize>().ok());
+
+        match number {
+            Some(value) if value > 0 => Ok(bytes),
+            _ => Err(()),
+        }
+    })(input)
+}
+
+fn scan_listing(input: &[u8]) -> IResult<&[u8], Stat> {
+    let (input, count) = message_number(input)?;
+    let (input, _) = space1(input)?;
+    let (input, size) = digit1(input)?;
+    let (input, _) = opt(not_line_ending)(input)?;
+    let (input, _) = eol(input)?;
+
+    Ok((input, Stat::new(count, size)))
+}
+
 pub(crate) fn stat_response(input: &[u8]) -> IResult<&[u8], Response> {
     let (input, stats) = stat(input)?;
 
@@ -68,9 +92,18 @@ fn list_stats(input: &[u8]) -> IResult<&[u8], Stat> {
 }
 
 pub(crate) fn list_response(input: &[u8]) -> IResult<&[u8], Response> {
-    let (input, stats) = alt((map(list_stats, Some), map(message_parser, |_| None)))(input)?;
+    let (input, (stats, first_item)) = alt((
+        map(list_stats, |stats| (Some(stats), None)),
+        map(scan_listing, |item| (None, Some(item))),
+        map(message_parser, |_| (None, None)),
+    ))(input)?;
+
+    let (input, (mut items, _end)) =
+        many_till(preceded(opt(tag(".")), scan_listing), end_of_multiline)(input)?;
 
-    let (input, (items, _end)) = many_till(preceded(opt(tag(".")), stat), end_of_multiline)(input)?;
+    if let Some(first_item) = first_item {
+        items.insert(0, first_item);
+    }
 
     let list = List::new(stats, items);
 
@@ -90,7 +123,7 @@ impl UniqueIdParser {
 }
 
 fn uidl(input: &[u8]) -> IResult<&[u8], UniqueId> {
-    let (input, index) = digit1(input)?;
+    let (input, index) = message_number(input)?;
     let (input, _) = space1(input)?;
     let (input, id) = UniqueIdParser::parse(input)?;
     let (input, _) = eol(input)?;
@@ -99,9 +132,17 @@ fn uidl(input: &[u8]) -> IResult<&[u8], UniqueId> {
 }
 
 pub(crate) fn uidl_list_response(input: &[u8]) -> IResult<&[u8], Response> {
-    let (input, message) = message_parser(input)?;
+    let (input, (message, first_item)) = alt((
+        map(uidl, |item| (None, Some(item))),
+        map(message_parser, |message| (message, None)),
+    ))(input)?;
+
+    let (input, (mut list, _end)) =
+        many_till(preceded(opt(tag(".")), uidl), end_of_multiline)(input)?;
 
-    let (input, (list, _end)) = many_till(preceded(opt(tag(".")), uidl), end_of_multiline)(input)?;
+    if let Some(first_item) = first_item {
+        list.insert(0, first_item);
+    }
 
     let list = Uidl::new(message, list);
 
@@ -111,9 +152,27 @@ pub(crate) fn uidl_list_response(input: &[u8]) -> IResult<&[u8], Response> {
 pub(crate) fn uidl_response(input: &[u8]) -> IResult<&[u8], Response> {
     let (input, unique_id) = uidl(input)?;
 
+    if looks_like_multiline_continuation(input) {
+        return Err(nom::Err::Error(NomError::new(input, ErrorKind::Verify)));
+    }
+
     Ok((input, Response::Uidl(unique_id.into())))
 }
 
+pub(crate) fn list_single_response(input: &[u8]) -> IResult<&[u8], Response> {
+    let (input, stats) = scan_listing(input)?;
+
+    if looks_like_multiline_continuation(input) {
+        return Err(nom::Err::Error(NomError::new(input, ErrorKind::Verify)));
+    }
+
+    Ok((input, Response::Stat(stats)))
+}
+
+fn looks_like_multiline_continuation(input: &[u8]) -> bool {
+    matches!(input.first(), Some(b'0'..=b'9') | Some(b'.'))
+}
+
 pub(crate) fn rfc822_response(input: &[u8]) -> IResult<&[u8], Response> {
     let (input, _message) = message_parser(input)?;