use mailparse to parse headers instead of reimplementing

src/header/headers.rs

@@ -1,5 +1,7 @@
-use super::{Header, parse_header};
-use crate::CowMailError;
+use mailparse::MailHeader;
+
+use super::{Header, build_header};
+use crate::{CowMailError, extend_over_newline};
 
 pub trait HeaderMap<'a, H: Header<'a> + 'a> {
     fn iter_headers(&'a self, key: &str) -> impl Iterator<Item = &'a H>;
@@ -93,29 +95,29 @@ impl<'a, H: Header<'a> + 'a> HeaderMap<'a, H> for Vec<H> {
     }
 }
 
+/// Rebuild the header lines mailparse has already scanned, so that a part's
+/// header block is only walked once. `parsed` has to be the headers of
+/// `raw_data`, so that their key and value slices point into it. Returns the
+/// headers and the index the body starts at.
+pub(crate) fn build_headers<'a, H: Header<'a>>(
+    raw_data: &'a [u8],
+    parsed: &[MailHeader<'a>],
+) -> Result<(Vec<H>, usize), CowMailError> {
+    let mut headers = Vec::with_capacity(parsed.len());
+    let mut line_end = 0;
+
+    for (ix, header) in parsed.iter().enumerate() {
+        let (header, end) = build_header(raw_data, header, parsed.get(ix + 1))?;
+        headers.push(header);
+        line_end = end;
+    }
+
+    // Whatever terminates the block: the blank line, or nothing at all when the
+    // message ends with its headers.
+    Ok((headers, extend_over_newline(raw_data, line_end)))
+}
 
 pub fn parse_headers<'a, H: Header<'a>>(raw_data: &'a [u8]) -> Result<(Vec<H>, usize), CowMailError> {
-    let mut headers: Vec<H> = Vec::new();
-    let mut ix = 0;
-    loop {
-        if ix >= raw_data.len() {
-            break;
-        } else if raw_data[ix] == b'\n' {
-            ix += 1;
-            break;
-        } else if raw_data[ix] == b'\r' {
-            if ix + 1 < raw_data.len() && raw_data[ix + 1] == b'\n' {
-                ix += 2;
-                break;
-            } else {
-                return Err(CowMailError::Generic(
-                    "Headers were followed by an unexpected lone CR character!".to_string(),
-                ));
-            }
-        }
-        let (header, ix_next) = parse_header(&raw_data[ix..])?;
-        headers.push(header);
-        ix += ix_next;
-    }
-    Ok((headers, ix))
+    let (parsed, _) = mailparse::parse_headers(raw_data)?;
+    build_headers(raw_data, &parsed)
 }

src/header/mod.rs

@@ -3,15 +3,16 @@ mod headers;
 
 pub use encoding::encode_header_value;
 pub use headers::{HeaderMap, parse_headers};
+pub(crate) use headers::build_headers;
 
 use std::borrow::Cow;
 use std::collections::BTreeMap;
 use std::fmt::{self, Display};
 
 use charset::decode_latin1;
-use mailparse::parse_content_type;
+use mailparse::{MailHeader, parse_content_type};
 
-use crate::CowMailError;
+use crate::{CowMailError, offset_in};
 
 pub trait Header<'a>: Clone + Display {
     fn new<K: ToString, V: ToString>(key: K, value: V) -> Self;
@@ -222,96 +223,65 @@ impl Display for CowHeader<'_> {
     }
 }
 
-enum HeaderParseState {
-    Initial,
-    Key,
-    PreValue,
-    Value,
-    ValueNewline,
+
+/// End of the — possibly folded — header line whose value ends at `value_end`:
+/// past the first newline that is not followed by continuation whitespace. Only
+/// `\r`, `\n`, spaces and tabs can appear between the end of a value and the end
+/// of its line; anything else would have moved `value_end` along.
+fn header_line_end(raw_data: &[u8], value_end: usize) -> usize {
+    let mut ix = value_end;
+    loop {
+        let Some(newline_ix) = raw_data[ix..].iter().position(|&byte| byte == b'\n') else {
+            return raw_data.len();
+        };
+        ix += newline_ix + 1;
+        if !matches!(raw_data.get(ix), Some(b' ' | b'\t')) {
+            return ix;
+        }
+    }
+}
+
+/// Rebuild the header line mailparse scanned into `parsed`, which has to be a
+/// header of `raw_data`, so that its key and value slices point into it.
+///
+/// A `MailHeader` keeps the key and the value, but not the line they came from,
+/// which is what we serialize. The key always starts at the first byte of its
+/// line, and a line ends where the next one's key begins — hence `next`; the
+/// last header of a block has none and needs a look ahead instead. Returns the
+/// header and the index just past its line.
+fn build_header<'a, H: Header<'a>>(
+    raw_data: &'a [u8],
+    parsed: &MailHeader<'a>,
+    next: Option<&MailHeader<'a>>,
+) -> Result<(H, usize), CowMailError> {
+    let offset_of = |bytes: &[u8]| {
+        offset_in(raw_data, bytes)
+            .ok_or_else(|| CowMailError::Generic("header is not a slice of its part".to_string()))
+    };
+
+    let key = parsed.get_key_raw();
+    let value = parsed.get_value_raw();
+
+    let line_start = offset_of(key)?;
+    let value_start = offset_of(value)?;
+    let value_end = value_start + value.len();
+
+    let line_end = match next {
+        Some(next) => offset_of(next.get_key_raw())?,
+        None => header_line_end(raw_data, value_end),
+    };
+
+    let header = H::from(
+        &raw_data[line_start..line_end],
+        key.len(),
+        value_start - line_start,
+        value_end - line_start,
+    );
+
+    Ok((header, line_end))
 }
 
 pub fn parse_header<'a, H: Header<'a>>(raw_data: &'a [u8]) -> Result<(H, usize), CowMailError> {
-    let mut it = raw_data.iter();
-    let mut ix = 0;
-    let mut c = match it.next() {
-        None => return Err(CowMailError::Generic("Empty string provided".to_string())),
-        Some(v) => *v,
-    };
-
-    let mut ix_key_end = None;
-    let mut ix_value_start = 0;
-    let mut ix_value_end = 0;
-
-    let mut state = HeaderParseState::Initial;
-    loop {
-        match state {
-            HeaderParseState::Initial => {
-                if c == b' ' {
-                    return Err(CowMailError::Generic(
-                        "Header cannot start with a space; it is likely an overhanging line from a previous header"
-                            .to_string(),
-                    ));
-                };
-                state = HeaderParseState::Key;
-                continue;
-            }
-            HeaderParseState::Key => {
-                if c == b':' {
-                    ix_key_end = Some(ix);
-                    state = HeaderParseState::PreValue;
-                } else if c == b'\n' {
-                    // Technically this is invalid. We'll handle it gracefully
-                    // since it does appear to happen in the wild and other
-                    // MTAs deal with it. Our handling is to just treat everything
-                    // encountered so far on this line as the header key, and
-                    // leave the value empty.
-                    ix_key_end = Some(ix);
-                    ix_value_start = ix;
-                    ix_value_end = ix;
-                    ix += 1;
-                    break;
-                }
-            }
-            HeaderParseState::PreValue => {
-                if c != b' ' {
-                    ix_value_start = ix;
-                    ix_value_end = ix;
-                    state = HeaderParseState::Value;
-                    continue;
-                }
-            }
-            HeaderParseState::Value => {
-                if c == b'\n' {
-                    state = HeaderParseState::ValueNewline;
-                } else if c != b'\r' {
-                    ix_value_end = ix + 1;
-                }
-            }
-            HeaderParseState::ValueNewline => {
-                if c == b' ' || c == b'\t' {
-                    state = HeaderParseState::Value;
-                    continue;
-                } else {
-                    break;
-                }
-            }
-        }
-        ix += 1;
-        c = match it.next() {
-            None => break,
-            Some(v) => *v,
-        };
-    }
-    match ix_key_end {
-        Some(v) => Ok((H::from(&raw_data[0..ix], v, ix_value_start, ix_value_end), ix)),
-
-        None => Ok((
-            // Technically this is invalid. We'll handle it gracefully
-            // since we handle the analogous situation above. Our handling
-            // is to just treat everything encountered on this line as
-            // the header key, and leave the value empty.
-            H::from(&raw_data[0..ix], ix, ix, ix),
-            ix,
-        )),
-    }
-}
+    let (parsed, _) = mailparse::parse_header(raw_data)?;
+    build_header(raw_data, &parsed, None)
+}

src/lib.rs

@@ -7,7 +7,7 @@ pub mod header;
 mod tests;
 
 pub use body::{BodyContent, ContentTransferEncoding};
-use header::CowHeader;
+use header::{CowHeader, build_headers};
 pub use header::{Header, HeaderMap, parse_headers};
 
 use mailparse::{
@@ -928,15 +928,16 @@ fn line_ending(bytes: &[u8]) -> &'static [u8] {
 /// buffer, which holds for every `ParsedMail::raw_bytes` relative to its
 /// parent's. Returns `None` instead of panicking, so that a mailparse that
 /// stops handing back subslices surfaces as an error, not as a bad index.
-fn offset_in(outer: &[u8], inner: &[u8]) -> Option<usize> {
+pub(crate) fn offset_in(outer: &[u8], inner: &[u8]) -> Option<usize> {
     let ix = (inner.as_ptr() as usize).checked_sub(outer.as_ptr() as usize)?;
     (ix + inner.len() <= outer.len()).then_some(ix)
 }
 
-/// Undo mailparse's `strip_trailing_crlf`: widen `end` over the newline that
-/// separates a part from the delimiter line following it, so that the slice
-/// covers everything the part occupied in the message.
-fn extend_over_newline(raw_data: &[u8], end: usize) -> usize {
+/// Widen `end` over the newline that follows it, if there is one. Used to undo
+/// mailparse's `strip_trailing_crlf`, so that a part's slice covers everything
+/// it occupied in the message, and to step over the blank line that ends a
+/// header block.
+pub(crate) fn extend_over_newline(raw_data: &[u8], end: usize) -> usize {
     if raw_data[end..].starts_with(b"\r\n") {
         end + 2
     } else if raw_data[end..].starts_with(b"\n") {
@@ -957,7 +958,7 @@ fn build_mail<'a, H: Header<'a> + 'a>(
     parsed: &ParsedMail<'a>,
     in_multipart_digest: bool,
 ) -> Result<CowMail<'a, H>, CowMailError> {
-    let (headers, body_idx) = parse_headers::<H>(raw_data)?;
+    let (headers, body_idx) = build_headers::<H>(raw_data, &parsed.headers)?;
 
     let mut body_end_idx = raw_data.len();
     let mut end_bytes: &[u8] = &[];
@@ -1042,4 +1043,4 @@ fn build_mail<'a, H: Header<'a> + 'a>(
     };
 
     Ok(mail)
-}
+}

src/tests.rs

@@ -961,3 +961,142 @@
         assert_eq!(mail.subparts[0].end_bytes, b"\r\n");
         assert_eq!(mail.end_bytes, b"");
     }
+
+    // The header lines below are rebuilt from the key/value slices mailparse
+    // hands back, so each of these pins a case where the line extends past the
+    // value: folding, a stray CR, an empty value, or the end of the buffer.
+    #[test]
+    fn test_roundtrip_folded_header() {
+        let raw = concat!(
+            "Subject: a\r\n",
+            " continued\r\n",
+            "To: john@example.com\r\n",
+            "\r\n",
+            "body\r\n"
+        );
+        assert_roundtrip(raw);
+
+        let mail = CowMail::try_from(raw.as_bytes()).unwrap();
+        assert_eq!(mail.headers.len(), 2);
+        assert_eq!(mail.headers[0].get_key_raw(), b"Subject");
+        assert_eq!(mail.headers[0].get_value_raw(), b"a\r\n continued");
+        assert_eq!(mail.headers[0].get_value(), "a continued");
+        assert_eq!(mail.headers[1].get_key_raw(), b"To");
+    }
+
+    // A fold whose last line holds nothing but whitespace: the value ends before
+    // the line does, for the last header as well as for one in the middle.
+    #[test]
+    fn test_roundtrip_header_folded_onto_blank_line() {
+        let raw = concat!("Subject: x\r\n", " \r\n", "To: y\r\n", " \r\n", "\r\n", "body\r\n");
+        assert_roundtrip(raw);
+
+        let mail = CowMail::try_from(raw.as_bytes()).unwrap();
+        assert_eq!(mail.headers.len(), 2);
+        assert_eq!(mail.headers[0].get_key_raw(), b"Subject");
+        assert_eq!(mail.headers[1].get_key_raw(), b"To");
+        assert_eq!(mail.raw_body(), b"body");
+    }
+
+    // Invalid, but it happens in the wild: mailparse treats the whole line as the
+    // key and leaves the value empty. Note that the key keeps the trailing CR.
+    #[test]
+    fn test_roundtrip_header_without_colon() {
+        let raw = concat!("Subject\r\n", "To: y\r\n", "\r\n", "body\r\n");
+        assert_roundtrip(raw);
+
+        let mail = CowMail::try_from(raw.as_bytes()).unwrap();
+        assert_eq!(mail.headers.len(), 2);
+        assert_eq!(mail.headers[0].get_key_raw(), b"Subject\r");
+        assert_eq!(mail.headers[0].get_value_raw(), b"");
+        assert_eq!(mail.headers[1].get_value_raw(), b"y");
+    }
+
+    #[test]
+    fn test_roundtrip_empty_header_value() {
+        let raw = concat!("X-Empty:\r\n", "Subject: x\r\n", "\r\n", "body\r\n");
+        assert_roundtrip(raw);
+
+        let mail = CowMail::try_from(raw.as_bytes()).unwrap();
+        assert_eq!(mail.headers[0].get_key_raw(), b"X-Empty");
+        assert_eq!(mail.headers[0].get_value_raw(), b"");
+        assert_eq!(mail.headers[1].get_value_raw(), b"x");
+    }
+
+    #[test]
+    fn test_roundtrip_no_headers() {
+        assert_roundtrip("\r\nbody\r\n");
+        assert_roundtrip("\nbody\n");
+
+        let mail = CowMail::try_from("\r\nbody\r\n".as_bytes()).unwrap();
+        assert!(mail.headers.is_empty());
+        assert_eq!(mail.raw_body(), b"body");
+    }
+
+    #[test]
+    fn test_roundtrip_lf_headers() {
+        let raw = concat!("Subject: a\n", " continued\n", "To: y\n", "\n", "body\n");
+        assert_roundtrip(raw);
+
+        let mail = CowMail::try_from(raw.as_bytes()).unwrap();
+        assert_eq!(mail.headers[0].get_value_raw(), b"a\n continued");
+        assert_eq!(mail.headers[1].get_key_raw(), b"To");
+    }
+
+    // A message that ends with its header block, with and without the final
+    // newline. The blank line separating headers from body is always written, so
+    // these are the one case that does not round-trip byte-exact.
+    #[test]
+    fn test_header_block_without_blank_line() {
+        let mail = CowMail::try_from("Subject: x\r\nTo: y\r\n".as_bytes()).unwrap();
+        assert_eq!(mail.headers.len(), 2);
+        assert_eq!(mail.headers[1].get_value_raw(), b"y");
+        assert_eq!(mail.raw_body(), b"");
+        assert_eq!(
+            String::from_utf8_lossy(&serialize(&mail)),
+            "Subject: x\r\nTo: y\r\n\r\n"
+        );
+
+        let mail = CowMail::try_from("Subject: x".as_bytes()).unwrap();
+        assert_eq!(mail.headers.len(), 1);
+        assert_eq!(mail.headers[0].get_value_raw(), b"x");
+        assert_eq!(mail.raw_body(), b"");
+        assert_eq!(String::from_utf8_lossy(&serialize(&mail)), "Subject: x\r\n");
+    }
+
+    // The header line is rebuilt from pointer offsets, not by searching for the
+    // key and value bytes, so a value identical to its key still lands at the
+    // right place.
+    #[test]
+    fn test_roundtrip_header_value_equal_to_key() {
+        for raw in [
+            "Test: Test\r\n\r\nbody\r\n",
+            "Test:Test\r\n\r\nbody\r\n",
+            "Test: Test\r\n Test\r\n\r\nbody\r\n",
+            "T: T\r\nT: T\r\n\r\nbody\r\n",
+        ] {
+            assert_roundtrip(raw);
+        }
+
+        let mail = CowMail::try_from("Test: Test\r\nTest:Test\r\n\r\nbody\r\n".as_bytes()).unwrap();
+        assert_eq!(mail.headers.len(), 2);
+        for header in &mail.headers {
+            assert_eq!(header.get_key_raw(), b"Test");
+            assert_eq!(header.get_value_raw(), b"Test");
+        }
+        assert_eq!(mail.headers[0].get_raw().as_ref(), b"Test: Test\r\n");
+        assert_eq!(mail.headers[1].get_raw().as_ref(), b"Test:Test\r\n");
+    }
+
+    // A header line that is nothing but a colon: both the key and the value are
+    // empty, and they sit at different offsets.
+    #[test]
+    fn test_roundtrip_header_colon_only() {
+        let raw = ":\r\nTo: y\r\n\r\nbody\r\n";
+        assert_roundtrip(raw);
+
+        let mail = CowMail::try_from(raw.as_bytes()).unwrap();
+        assert_eq!(mail.headers[0].get_key_raw(), b"");
+        assert_eq!(mail.headers[0].get_value_raw(), b"");
+        assert_eq!(mail.headers[0].get_raw().as_ref(), b":\r\n");
+    }