husk-milter/src/types/cheader.rs

//! CHeader struct
//!
//! This struct is for collecting header lines from the milter.

use std::ffi::CString;
use std::fmt;

use anyhow::Result;

use crate::types::recipient::Recipient;

#[derive(Debug, PartialEq, Clone)]
pub struct CHeader {
    lines: Vec<(CString, CString)>
}

impl CHeader {
    pub fn new() -> Self {
        CHeader {
            lines: Vec::new()
        }
    }

    pub fn add(mut self, key: CString, value: CString) -> Self {
        self.lines.push((key, value));
        self
    }

    pub fn append(&mut self, key: CString, value: CString) {
        self.lines.push((key, value));
    }

    pub fn find(&self, key: &str) -> Option<(&str, &str)> {
        let lower_key = key.to_lowercase();
        for (line_key, line_value) in self.lines.iter() {
            if let Ok(line_compare) = line_key.to_str() {
                if line_compare.to_ascii_lowercase() == lower_key {
                    if let Ok(line_value_as_str) = line_value.to_str() {
                        return Some(
                            (line_compare, line_value_as_str)
                        );
                    }
                }
            }
        }
        None
    }

    pub fn get_value_or<'ch>(&'ch self, key: &str, default: &'ch str) -> &'ch str {
        if let Some((_, value)) = self.find(key) {
            value
        } else {
            default
        }
    }

    pub fn iter_lines(&self) -> std::slice::Iter<'_, (CString, CString)> {
        self.lines.iter()
    }

    /// Remove a recipient from a value taken from a to, cc or bcc header
    /// line. Only the email address is matched.
    pub fn remove_recipient(header_value: &CString, to_remove: &Vec<Recipient>) -> Result<CString> {
        let content = String::from_utf8_lossy(header_value.as_bytes());
        let parts = content.split(",")
            .filter(|part| {
                let mut keep = true;
                let p = part.to_ascii_lowercase();
                for rcpt in to_remove {
                    let rcpt = rcpt.email.to_ascii_lowercase();
                    if p.contains(&rcpt) { keep = false; break; }
                }
                keep
            })
            .collect::<Vec<&str>>();
        CString::new(parts.join(","))
            .map_err(|e| e.into())
    }

    /// Returns a new header with `to_remove` removed from to, cc and bcc.
    pub fn without_recipients(&self, to_remove: &Vec<Recipient>) -> CHeader {
        let mut result = CHeader::new();

        for (key, value) in self.iter_lines() {
            let new_value = match key.as_bytes().to_ascii_lowercase().as_slice() {
                b"to"
                | b"cc"
                | b"bcc" => Self::remove_recipient(value, to_remove),
                _ => Ok(value.clone())
            };
            if let Ok(value) = new_value {
                if ! value.is_empty() {
                    result.append(key.clone(), value);
                }
            }
        }

        result
    }

}

impl fmt::Display for CHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", [
            self.iter_lines()
                .filter_map(|(k, v)| {
                    match (k.to_str(), v.to_str()) {
                        (Ok(k), Ok(v)) => Some([k, v].join(": ")),
                        _ => None,
                    }
                })
                .collect::<Vec<String>>()
                .join("\r\n"),
            "\r\n".to_string()]
            .join("")
        )
    }
}

#[cfg(test)]
pub mod tests {
    use std::ffi::CString;
    use crate::types::recipient::Recipient;

    use super::CHeader;


    #[test]
    pub fn basic_functionality() {
        let cheader = CHeader::new()
            .add(CString::new("Content-Type").unwrap(), CString::new("text/plain").unwrap())
            .add(CString::new("Content-Transfer-Encoding").unwrap(), CString::new("7bit").unwrap());

        assert_eq!("Content-Type: text/plain\r\nContent-Transfer-Encoding: 7bit\r\n", cheader.to_string());
    }

    #[test]
    pub fn find() {
        let cheader = CHeader::new()
            .add(CString::new("Content-Type").unwrap(), CString::new("text/plain").unwrap())
            .add(CString::new("Content-Transfer-Encoding").unwrap(), CString::new("7bit").unwrap());

        let line = cheader.find("content-type");

        assert!(line.is_some());

        let (key, value) = line.unwrap();

        assert_eq!("Content-Type", key);
        assert_eq!("text/plain", value);
    }

    #[test]
    pub fn find_none() {
        let cheader = CHeader::new()
            .add(CString::new("Content-Type").unwrap(), CString::new("text/plain").unwrap())
            .add(CString::new("Content-Transfer-Encoding").unwrap(), CString::new("7bit").unwrap());

        let line = cheader.find("foo-bar");

        assert!(line.is_none());
    }

    #[test]
    pub fn find_with_default() {
        let cheader = CHeader::new()
            .add(CString::new("Content-Type").unwrap(), CString::new("text/plain").unwrap())
            .add(CString::new("Content-Transfer-Encoding").unwrap(), CString::new("7bit").unwrap());

        let value = cheader.get_value_or("foo-bar", "baaz");

        assert_eq!(value, "baaz");
    }

    #[test]
    pub fn remove_recipients() {
        let header = CHeader::new()
            .add(CString::new("To").unwrap(), CString::new("Alice <alice@example.com>, Bob <bob@other.org>").unwrap())
            .add(CString::new("Cc").unwrap(), CString::new("carol@example.com").unwrap());

        let modified_header = header.without_recipients(
            &vec![ Recipient::from("bob@other.org".to_string()) ]
        );
        let expected_header = CHeader::new()
            .add(CString::new("To").unwrap(), CString::new("Alice <alice@example.com>").unwrap())
            .add(CString::new("Cc").unwrap(), CString::new("carol@example.com").unwrap());
        assert_eq!(modified_header, expected_header);

        let modified_header = header.without_recipients(
            &vec![ Recipient::from("CaRol@exAmple.Com".to_string()) ]
        );
        let expected_header = CHeader::new()
            .add(CString::new("To").unwrap(), CString::new("Alice <alice@example.com>, Bob <bob@other.org>").unwrap());
        assert_eq!(modified_header, expected_header);
    }
}