miltr_common/commands/
mail.rs

1use std::borrow::Cow;
2
3use bytes::{BufMut, BytesMut};
4
5use crate::decoding::Parsable;
6use crate::encoding::Writable;
7use crate::{InvalidData, ProtocolError};
8use miltr_utils::ByteParsing;
9
10/// Information about a mail to be processed
11#[derive(Clone, PartialEq, Debug, Default)]
12pub struct Mail {
13    sender: BytesMut,
14    esmtp_args: Option<BytesMut>,
15}
16
17impl From<&[u8]> for Mail {
18    fn from(value: &[u8]) -> Self {
19        Self {
20            sender: BytesMut::from_iter(value),
21            esmtp_args: None,
22        }
23    }
24}
25
26impl Mail {
27    const CODE: u8 = b'M';
28    /// The sender of this email
29    #[must_use]
30    pub fn sender(&self) -> Cow<str> {
31        String::from_utf8_lossy(&self.sender)
32    }
33
34    /// Optionally set additional esmtp args.
35    ///
36    /// If those are empty, an empty vector is returned.
37    #[must_use]
38    pub fn esmtp_args(&self) -> Vec<Cow<str>> {
39        let Some(args) = &self.esmtp_args else {
40            return Vec::new();
41        };
42
43        args[..]
44            .split(|&b| b == 0)
45            .map(String::from_utf8_lossy)
46            .collect()
47    }
48}
49
50impl Parsable for Mail {
51    const CODE: u8 = Self::CODE;
52
53    fn parse(mut buffer: BytesMut) -> Result<Self, ProtocolError> {
54        let Some(sender) = buffer.delimited(0) else {
55            return Err(InvalidData::new(
56                "Null-byte missing in mail package to sender hostname",
57                buffer,
58            )
59            .into());
60        };
61
62        let esmtp_args = {
63            if buffer.is_empty() {
64                None
65            } else {
66                Some(buffer)
67            }
68        };
69
70        Ok(Self { sender, esmtp_args })
71    }
72}
73
74impl Writable for Mail {
75    fn write(&self, buffer: &mut BytesMut) {
76        buffer.extend_from_slice(&self.sender);
77        buffer.put_u8(0);
78        if let Some(b) = &self.esmtp_args {
79            buffer.extend_from_slice(b);
80        }
81    }
82
83    fn len(&self) -> usize {
84        self.sender.len()
85            + 1
86            + self
87                .esmtp_args
88                .as_ref()
89                .map(BytesMut::len)
90                .unwrap_or_default()
91    }
92
93    fn code(&self) -> u8 {
94        Self::CODE
95    }
96
97    fn is_empty(&self) -> bool {
98        self.sender.is_empty() && self.esmtp_args.is_some()
99    }
100}
101
102/// SMTP Data command has been sent
103#[derive(Clone, PartialEq, Debug, Default)]
104pub struct Data;
105
106impl Data {
107    const CODE: u8 = b'T';
108}
109
110impl Parsable for Data {
111    const CODE: u8 = Self::CODE;
112
113    fn parse(_buffer: BytesMut) -> Result<Self, ProtocolError> {
114        Ok(Self)
115    }
116}
117
118impl Writable for Data {
119    fn write(&self, _buffer: &mut BytesMut) {}
120
121    fn len(&self) -> usize {
122        0
123    }
124
125    fn code(&self) -> u8 {
126        Self::CODE
127    }
128
129    fn is_empty(&self) -> bool {
130        false
131    }
132}
133
134#[cfg(test)]
135mod test {
136    use super::*;
137    use crate::decoding::Parsable;
138    use pretty_assertions::assert_eq;
139    use rstest::rstest;
140
141    #[rstest]
142    #[case(BytesMut::from("sender\0arg1\0arg2"), Ok( Mail {sender: BytesMut::from("sender"), esmtp_args: Some(BytesMut::from("arg1\0arg2"))}))]
143    #[case(
144        BytesMut::from("senderarg1arg2"),
145        Err(InvalidData::new(
146            "Null-byte missing in mail package to sender hostname",
147            BytesMut::new(),
148        ))
149    )]
150    fn test_mail(#[case] input: BytesMut, #[case] expected: Result<Mail, InvalidData>) {
151        let parsed_mail = Mail::parse(input);
152
153        match parsed_mail {
154            Ok(mail) => {
155                let expected_mail = expected.unwrap();
156                assert_eq!(mail.sender, expected_mail.sender);
157                assert_eq!(mail.esmtp_args, expected_mail.esmtp_args);
158
159                //test function mail.esmtp_args()
160                let vec: Vec<Cow<str>> = vec![
161                    String::from_utf8_lossy(b"arg1"),
162                    String::from_utf8_lossy(b"arg2"),
163                ];
164                assert_eq!(mail.esmtp_args(), vec);
165            }
166
167            Err(ProtocolError::InvalidData(e)) => {
168                assert_eq!(e.msg, expected.unwrap_err().msg);
169            }
170            _ => panic!("Wrong error received"),
171        }
172    }
173
174    #[cfg(feature = "count-allocations")]
175    #[test]
176    fn test_parse_mail() {
177        use super::Mail;
178
179        let buffer = BytesMut::from("sender\0arg1\0arg2");
180        let info = allocation_counter::measure(|| {
181            let res = Mail::parse(buffer);
182            allocation_counter::opt_out(|| {
183                println!("{res:?}");
184                assert!(res.is_ok());
185            });
186        });
187
188        println!("{}", &info.count_total);
189        assert_eq!(info.count_total, 1);
190    }
191}