Fixed IMAP encoding on non-UTF-8 messages + improved Received headers in SMTP

crates/imap-proto/src/protocol/fetch.rs

@@ -112,7 +112,7 @@ pub enum DataItem<'x> {
     BodySection {
         sections: Vec<Section>,
         origin_octet: Option<u32>,
-        contents: Cow<'x, str>,
+        contents: Cow<'x, [u8]>,
     },
     Envelope {
         envelope: Envelope<'x>,
@@ -127,19 +127,19 @@ pub enum DataItem<'x> {
         uid: u32,
     },
     Rfc822 {
-        contents: Cow<'x, str>,
+        contents: Cow<'x, [u8]>,
     },
     Rfc822Header {
-        contents: Cow<'x, str>,
+        contents: Cow<'x, [u8]>,
     },
     Rfc822Size {
         size: usize,
     },
     Rfc822Text {
-        contents: Cow<'x, str>,
+        contents: Cow<'x, [u8]>,
     },
     Preview {
-        contents: Option<Cow<'x, str>>,
+        contents: Option<Cow<'x, [u8]>>,
     },
     ModSeq {
         modseq: u64,
@@ -733,7 +733,7 @@ impl<'x> DataItem<'x> {
                 }
                 match contents {
                     BodyContents::Text(text) => {
-                        literal_string(buf, text);
+                        literal_string(buf, text.as_bytes());
                     }
                     BodyContents::Bytes(bytes) => {
                         buf.extend_from_slice(b"~{");
@@ -1322,7 +1322,7 @@ mod tests {
                         Section::Mime,
                     ],
                     origin_octet: 11.into(),
-                    contents: "howdy".into(),
+                    contents: b"howdy"[..].into(),
                 },
                 "BODY[1.2.MIME]<11> {5}\r\nhowdy",
             ),
@@ -1333,7 +1333,7 @@ mod tests {
                         fields: vec!["Subject".into(), "x-special".into()],
                     }],
                     origin_octet: None,
-                    contents: "howdy".into(),
+                    contents: b"howdy"[..].into(),
                 },
                 "BODY[HEADER.FIELDS.NOT (SUBJECT X-SPECIAL)] {5}\r\nhowdy",
             ),
@@ -1344,7 +1344,7 @@ mod tests {
                         fields: vec!["From".into(), "List-Archive".into()],
                     }],
                     origin_octet: None,
-                    contents: "howdy".into(),
+                    contents: b"howdy"[..].into(),
                 },
                 "BODY[HEADER.FIELDS (FROM LIST-ARCHIVE)] {5}\r\nhowdy",
             ),
@@ -1382,10 +1382,10 @@ mod tests {
                             super::DataItem::Uid { uid: 983 },
                             super::DataItem::Rfc822Size { size: 443 },
                             super::DataItem::Rfc822Text {
-                                contents: "hi".into()
+                                contents: b"hi"[..].into()
                             },
                             super::DataItem::Rfc822Header {
-                                contents: "header".into()
+                                contents: b"header"[..].into()
                             },
                         ],
                     }],

crates/imap-proto/src/protocol/mod.rs

@@ -187,11 +187,11 @@ pub fn quoted_string_or_nil(buf: &mut Vec<u8>, text: Option<&str>) {
     }
 }
 
-pub fn literal_string(buf: &mut Vec<u8>, text: &str) {
+pub fn literal_string(buf: &mut Vec<u8>, text: &[u8]) {
     buf.push(b'{');
     buf.extend_from_slice(text.len().to_string().as_bytes());
     buf.extend_from_slice(b"}\r\n");
-    buf.extend_from_slice(text.as_bytes());
+    buf.extend_from_slice(text);
 }
 
 pub fn quoted_timestamp(buf: &mut Vec<u8>, timestamp: i64) {

crates/imap/src/op/fetch.rs

@@ -379,7 +379,10 @@ impl SessionData {
                     }
                     Attribute::Preview { .. } => {
                         items.push(DataItem::Preview {
-                            contents: email.get(&Property::Preview).as_string().map(|p| p.into()),
+                            contents: email
+                                .get(&Property::Preview)
+                                .as_string()
+                                .map(|p| p.as_bytes().into()),
                         });
                     }
                     Attribute::Rfc822Size => {
@@ -392,7 +395,7 @@ impl SessionData {
                     }
                     Attribute::Rfc822 => {
                         items.push(DataItem::Rfc822 {
-                            contents: String::from_utf8_lossy(raw_message.as_ref().unwrap()),
+                            contents: raw_message.as_ref().unwrap().into(),
                         });
                     }
                     Attribute::Rfc822Header => {
@@ -403,7 +406,7 @@ impl SessionData {
                             .get(message.offset_header..message.offset_body)
                         {
                             items.push(DataItem::Rfc822Header {
-                                contents: String::from_utf8_lossy(header),
+                                contents: header.into(),
                             });
                         }
                     }
@@ -415,7 +418,7 @@ impl SessionData {
                             .get(message.offset_body..message.offset_end)
                         {
                             items.push(DataItem::Rfc822Text {
-                                contents: String::from_utf8_lossy(text),
+                                contents: text.into(),
                             });
                         }
                     }
@@ -593,7 +596,7 @@ pub trait AsImapDataItem<'x> {
         &'z self,
         sections: &[Section],
         partial: Option<(u32, u32)>,
-    ) -> Option<Cow<'x, str>>;
+    ) -> Option<Cow<'x, [u8]>>;
     fn binary(
         &self,
         sections: &[u32],
@@ -805,14 +808,16 @@ impl<'x> AsImapDataItem<'x> for Message<'x> {
         &'z self,
         sections: &[Section],
         partial: Option<(u32, u32)>,
-    ) -> Option<Cow<'x, str>> {
+    ) -> Option<Cow<'x, [u8]>> {
         let mut part = self.root_part();
         if sections.is_empty() {
-            return String::from_utf8_lossy(get_partial_bytes(
+            return Some(
-                self.raw_message.get(part.offset_header..part.offset_end)?,
+                get_partial_bytes(
-                partial,
+                    self.raw_message.get(part.offset_header..part.offset_end)?,
-            ))
+                    partial,
-            .into();
+                )
+                .into(),
+            );
         }
 
         let mut message = self;
@@ -848,13 +853,15 @@ impl<'x> AsImapDataItem<'x> for Message<'x> {
                     }
                 }
                 Section::Header => {
-                    return String::from_utf8_lossy(get_partial_bytes(
+                    return Some(
-                        message
+                        get_partial_bytes(
-                            .raw_message
+                            message
-                            .get(part.offset_header..part.offset_body)?,
+                                .raw_message
-                        partial,
+                                .get(part.offset_header..part.offset_body)?,
-                    ))
+                            partial,
-                    .into();
+                        )
+                        .into(),
+                    );
                 }
                 Section::HeaderFields { not, fields } => {
                     let mut headers =
@@ -876,22 +883,19 @@ impl<'x> AsImapDataItem<'x> for Message<'x> {
                     headers.extend_from_slice(b"\r\n");
 
                     return Some(if partial.is_none() {
-                        String::from_utf8(headers).map_or_else(
+                        headers.into()
-                            |err| String::from_utf8_lossy(err.as_bytes()).into_owned().into(),
-                            |s| s.into(),
-                        )
                     } else {
-                        String::from_utf8_lossy(get_partial_bytes(&headers, partial))
+                        get_partial_bytes(&headers, partial).to_vec().into()
-                            .into_owned()
-                            .into()
                     });
                 }
                 Section::Text => {
-                    return String::from_utf8_lossy(get_partial_bytes(
+                    return Some(
-                        message.raw_message.get(part.offset_body..part.offset_end)?,
+                        get_partial_bytes(
-                        partial,
+                            message.raw_message.get(part.offset_body..part.offset_end)?,
-                    ))
+                            partial,
-                    .into();
+                        )
+                        .into(),
+                    );
                 }
                 Section::Mime => {
                     let mut headers =
@@ -912,14 +916,9 @@ impl<'x> AsImapDataItem<'x> for Message<'x> {
                     }
                     headers.extend_from_slice(b"\r\n");
                     return Some(if partial.is_none() {
-                        String::from_utf8(headers).map_or_else(
+                        headers.into()
-                            |err| String::from_utf8_lossy(err.as_bytes()).into_owned().into(),
-                            |s| s.into(),
-                        )
                     } else {
-                        String::from_utf8_lossy(get_partial_bytes(&headers, partial))
+                        get_partial_bytes(&headers, partial).to_vec().into()
-                            .into_owned()
-                            .into()
                     });
                 }
             }
@@ -928,17 +927,13 @@ impl<'x> AsImapDataItem<'x> for Message<'x> {
         // BODY[x] should return both headers and body, but most clients
         // expect BODY[x] to return only the body, just like BOXY[x.TEXT] does.
 
-        String::from_utf8_lossy(get_partial_bytes(
+        Some(
-            message.raw_message.get(part.offset_body..part.offset_end)?,
+            get_partial_bytes(
-            partial,
+                message.raw_message.get(part.offset_body..part.offset_end)?,
-        ))
+                partial,
-        .into()
+            )
-
+            .into(),
-        /*String::from_utf8_lossy(get_partial_bytes(
+        )
-            raw_message.get(part.offset_header..part.offset_end)?,
-            partial,
-        ))
-        .into()*/
     }
 
     fn binary(

crates/smtp/src/inbound/data.rs

@@ -415,7 +415,7 @@ impl<T: AsyncWrite + AsyncRead + IsTls + Unpin> Session<T> {
             let params = self
                 .build_script_parameters()
                 .with_message(edited_message.as_ref().unwrap_or(&raw_message).clone())
-                .set_variable("from", auth_message.from().to_string())
+                .set_variable("dmarc_from", auth_message.from().to_string())
                 .set_variable(
                     "arc_result",
                     arc_output
@@ -735,7 +735,7 @@ impl<T: AsyncWrite + AsyncRead + IsTls + Unpin> Session<T> {
                 .iprev
                 .as_ref()
                 .and_then(|ir| ir.ptr.as_ref())
-                .and_then(|ptr| ptr.first().map(|s| s.as_str()))
+                .and_then(|ptr| ptr.first().map(|s| s.strip_suffix('.').unwrap_or(s)))
                 .unwrap_or("unknown")
                 .as_bytes(),
         );
@@ -747,12 +747,12 @@ impl<T: AsyncWrite + AsyncRead + IsTls + Unpin> Session<T> {
         headers.extend_from_slice(self.instance.hostname.as_bytes());
         headers.extend_from_slice(b" (Stalwart SMTP) with ");
         headers.extend_from_slice(
-            if self.stream.is_tls() {
+            match (self.stream.is_tls(), self.data.authenticated_as.is_empty()) {
-                "ESMTPS"
+                (true, true) => b"ESMTPS",
-            } else {
+                (true, false) => b"ESMTPSA",
-                "ESMTP"
+                (false, true) => b"ESMTP",
-            }
+                (false, false) => b"ESMTPA",
-            .as_bytes(),
+            },
         );
         headers.extend_from_slice(b" id ");
         headers.extend_from_slice(format!("{id:X}").as_bytes());

crates/smtp/src/scripts/functions.rs

@@ -28,6 +28,7 @@ pub fn register_functions() -> FunctionMap {
     FunctionMap::new()
         .with_function("trim", |v| v.to_cow().trim().to_string().into())
         .with_function("len", |v| v.to_cow().len().into())
+        .with_function("is_empty", |v| v.to_cow().as_ref().is_empty().into())
         .with_function("to_lowercase", |v| {
             v.to_cow().to_lowercase().to_string().into()
         })
@@ -40,12 +41,24 @@ pub fn register_functions() -> FunctionMap {
                 .unwrap_or("unknown")
                 .into()
         })
-        .with_function("domain", |v| {
+        .with_function("is_email", |v| is_email_valid(v.to_cow().as_ref()).into())
+        .with_function("domain_part", |v| {
             v.to_cow()
                 .rsplit_once('@')
                 .map_or(Variable::default(), |(_, d)| d.trim().to_string().into())
         })
-        .with_function("base_domain", |v| {
+        .with_function("local_part", |v| {
+            v.to_cow()
+                .rsplit_once('@')
+                .map_or(Variable::default(), |(u, _)| u.trim().to_string().into())
+        })
+        .with_function("domain_name_part", |v| {
+            v.to_cow()
+                .rsplit_once('@')
+                .and_then(|(_, d)| d.trim().split('.').rev().nth(1).map(|s| s.to_string()))
+                .map_or(Variable::default(), Variable::from)
+        })
+        .with_function("subdomain_part", |v| {
             v.to_cow()
                 .rsplit_once('@')
                 .map_or(Variable::default(), |(_, d)| {
@@ -81,7 +94,101 @@ pub fn register_functions() -> FunctionMap {
                 .all(|c| c.is_lowercase())
                 .into()
         })
-        .with_function("word_count", |v| {
+        .with_function("count_words", |v| {
             v.to_cow().as_ref().split_whitespace().count().into()
         })
+        .with_function("count_chars", |v| {
+            v.to_cow().as_ref().chars().count().into()
+        })
+        .with_function("count_control_chars", |v| {
+            v.to_cow()
+                .as_ref()
+                .chars()
+                .filter(|c| {
+                    matches!(c, '\u{0000}'..='\u{0008}'
+                    | '\u{000B}'
+                    | '\u{000C}'
+                    | '\u{000E}'..='\u{001F}'
+                    | '\u{007F}')
+                })
+                .count()
+                .into()
+        })
+        .with_function("count", |v| {
+            if let Variable::Array(l) = v {
+                l.len().into()
+            } else {
+                1.into()
+            }
+        })
+}
+
+fn is_email_valid(email: &str) -> bool {
+    let mut last_ch = 0;
+    let mut in_quote = false;
+    let mut at_count = 0;
+    let mut dot_count = 0;
+    let mut lp_len = 0;
+    let mut value = 0;
+
+    for ch in email.bytes() {
+        match ch {
+            b'0'..=b'9'
+            | b'a'..=b'z'
+            | b'A'..=b'Z'
+            | b'!'
+            | b'#'
+            | b'$'
+            | b'%'
+            | b'&'
+            | b'\''
+            | b'*'
+            | b'+'
+            | b'-'
+            | b'/'
+            | b'='
+            | b'?'
+            | b'^'
+            | b'_'
+            | b'`'
+            | b'{'
+            | b'|'
+            | b'}'
+            | b'~'
+            | 0x7f..=u8::MAX => {
+                value += 1;
+            }
+            b'.' if !in_quote => {
+                if last_ch != b'.' && last_ch != b'@' && value != 0 {
+                    value += 1;
+                    if at_count == 1 {
+                        dot_count += 1;
+                    }
+                } else {
+                    return false;
+                }
+            }
+            b'@' if !in_quote => {
+                at_count += 1;
+                lp_len = value;
+                value = 0;
+            }
+            b'>' | b':' | b',' | b' ' if in_quote => {
+                value += 1;
+            }
+            b'\"' if !in_quote || last_ch != b'\\' => {
+                in_quote = !in_quote;
+            }
+            b'\\' if in_quote && last_ch != b'\\' => (),
+            _ => {
+                if !in_quote {
+                    return false;
+                }
+            }
+        }
+
+        last_ch = ch;
+    }
+
+    at_count == 1 && dot_count > 0 && lp_len > 0 && value > 0
 }