miltr_common/commands/
header.rs

1use std::borrow::Cow;
2
3use bytes::{BufMut, BytesMut};
4
5use crate::decoding::Parsable;
6use crate::encoding::Writable;
7use crate::InvalidData;
8use crate::ProtocolError;
9use miltr_utils::ByteParsing;
10
11/// An smtp header received
12#[derive(Clone, PartialEq, Debug, Default)]
13pub struct Header {
14    name: BytesMut,
15    value: BytesMut,
16}
17
18impl Header {
19    const CODE: u8 = b'L';
20
21    /// Create a Header from some bytes
22    #[must_use]
23    pub fn new(name: &[u8], value: &[u8]) -> Self {
24        Self {
25            name: BytesMut::from_iter(name),
26            value: BytesMut::from_iter(value),
27        }
28    }
29    /// The name of the received header
30    #[must_use]
31    pub fn name(&self) -> Cow<str> {
32        String::from_utf8_lossy(&self.name)
33    }
34
35    /// The value of the received header
36    #[must_use]
37    pub fn value(&self) -> Cow<str> {
38        String::from_utf8_lossy(&self.value)
39    }
40}
41
42impl Parsable for Header {
43    const CODE: u8 = Self::CODE;
44
45    fn parse(mut buffer: BytesMut) -> Result<Self, ProtocolError> {
46        let Some(name) = buffer.delimited(0) else {
47            return Err(InvalidData::new(
48                "Received header package without name terminated by null byte in it",
49                buffer,
50            )
51            .into());
52        };
53
54        let Some(value) = buffer.delimited(0) else {
55            return Err(InvalidData::new(
56                "Received header package without value terminated by null byte in it",
57                buffer,
58            )
59            .into());
60        };
61
62        Ok(Self { name, value })
63    }
64}
65
66impl Writable for Header {
67    fn write(&self, buffer: &mut BytesMut) {
68        buffer.extend_from_slice(&self.name);
69        buffer.put_u8(0);
70        buffer.extend_from_slice(&self.value);
71        buffer.put_u8(0);
72    }
73
74    fn len(&self) -> usize {
75        self.name.len() + 1 + self.value.len() + 1
76    }
77
78    fn code(&self) -> u8 {
79        Self::CODE
80    }
81
82    fn is_empty(&self) -> bool {
83        self.name.is_empty() && self.value.is_empty()
84    }
85}
86
87/// After all headers have been sent, end of header is sent
88#[derive(Clone, PartialEq, Debug, Default)]
89pub struct EndOfHeader;
90
91impl EndOfHeader {
92    const CODE: u8 = b'N';
93}
94
95impl Parsable for EndOfHeader {
96    const CODE: u8 = Self::CODE;
97
98    fn parse(_buffer: BytesMut) -> Result<Self, ProtocolError> {
99        Ok(Self)
100    }
101}
102
103impl Writable for EndOfHeader {
104    fn write(&self, _buffer: &mut BytesMut) {}
105
106    fn len(&self) -> usize {
107        0
108    }
109
110    fn code(&self) -> u8 {
111        Self::CODE
112    }
113
114    fn is_empty(&self) -> bool {
115        false
116    }
117}
118
119#[cfg(test)]
120mod test {
121    use super::*;
122    use crate::decoding::Parsable;
123    use pretty_assertions::assert_eq;
124    use rstest::rstest;
125
126    #[rstest]
127    #[case(BytesMut::from("name\0value\0"), Ok(Header {name: BytesMut::from("name"), value: BytesMut::from("value")} ))]
128    #[case(
129        BytesMut::from("name\0value"),
130        Err(InvalidData::new(
131            "Received header package without value terminated by null byte in it",
132            BytesMut::new()
133        ))
134    )]
135    #[case(
136        BytesMut::from("namevalue\0"),
137        Err(InvalidData::new(
138            "Received header package without value terminated by null byte in it",
139            BytesMut::new()
140        ))
141    )]
142    fn test_header(#[case] input: BytesMut, #[case] expected: Result<Header, InvalidData>) {
143        let parsed_header = Header::parse(input);
144
145        match (expected, parsed_header) {
146            (Err(expected), Err(ProtocolError::InvalidData(parsed))) => {
147                assert_eq!(expected.msg, parsed.msg);
148            }
149            (Ok(expected), Ok(parsed)) => assert_eq!(expected, parsed),
150            (expected, parsed) => panic!("Did not get expected:\n{expected:?}\n vs \n{parsed:?}"),
151        }
152    }
153    #[cfg(feature = "count-allocations")]
154    #[test]
155    fn test_parse_header() {
156        let buffer = BytesMut::from("name\0value\0");
157
158        let info = allocation_counter::measure(|| {
159            let res = Header::parse(buffer);
160
161            allocation_counter::opt_out(|| {
162                println!("{res:?}");
163                assert!(res.is_ok());
164            });
165        });
166
167        println!("{info:#?}");
168        assert_eq!(info.count_total, 1);
169    }
170}
miltr_common/commands/header.rs

miltr_common/commands/
header.rs
