miltr_common/modifications/
headers.rs

1//! Add, change or insert smtp headers
2
3use std::borrow::Cow;
4
5use bytes::{BufMut, BytesMut};
6
7use crate::commands::Header;
8use crate::decoding::Parsable;
9use crate::encoding::Writable;
10use crate::error::STAGE_DECODING;
11use crate::{NotEnoughData, ProtocolError};
12use miltr_utils::ByteParsing;
13
14/// Add a header
15#[derive(Debug, Clone)]
16pub struct AddHeader {
17    header: Header,
18}
19
20impl AddHeader {
21    const CODE: u8 = b'h';
22
23    /// Create a Header from some bytes
24    #[must_use]
25    pub fn new(name: &[u8], value: &[u8]) -> Self {
26        Self {
27            header: Header::new(name, value),
28        }
29    }
30
31    /// The name of the header
32    #[must_use]
33    pub fn name(&self) -> Cow<str> {
34        self.header.name()
35    }
36
37    /// The value of the header
38    #[must_use]
39    pub fn value(&self) -> Cow<str> {
40        self.header.value()
41    }
42}
43
44impl Parsable for AddHeader {
45    const CODE: u8 = Self::CODE;
46
47    fn parse(buffer: BytesMut) -> Result<Self, ProtocolError> {
48        let header = Header::parse(buffer)?;
49
50        Ok(Self { header })
51    }
52}
53
54impl Writable for AddHeader {
55    fn write(&self, buffer: &mut BytesMut) {
56        self.header.write(buffer);
57    }
58
59    fn len(&self) -> usize {
60        self.header.len()
61    }
62
63    fn code(&self) -> u8 {
64        Self::CODE
65    }
66    fn is_empty(&self) -> bool {
67        self.header.is_empty()
68    }
69}
70
71/// Change an existing header
72#[derive(Debug, Clone)]
73pub struct ChangeHeader {
74    /// The index in a list of headers sharing `name` which to change
75    ///
76    /// Headers can be set multiple times. This index is only valid in the
77    /// context of headers with the same name.
78    index: u32,
79
80    header: Header,
81}
82
83impl ChangeHeader {
84    const CODE: u8 = b'm';
85
86    /// Create a Header from some bytes
87    #[must_use]
88    pub fn new(index: u32, name: &[u8], value: &[u8]) -> Self {
89        Self {
90            index,
91            header: Header::new(name, value),
92        }
93    }
94
95    /// The name of the header
96    #[must_use]
97    pub fn name(&self) -> Cow<str> {
98        self.header.name()
99    }
100
101    /// The value of the header
102    #[must_use]
103    pub fn value(&self) -> Cow<str> {
104        self.header.value()
105    }
106
107    /// The index in a list of headers sharing `name` which to change
108    ///
109    /// Headers can be set multiple times. This index is only valid in the
110    /// context of headers with the same name.
111    #[must_use]
112    pub fn index(&self) -> u32 {
113        self.index
114    }
115}
116
117impl Parsable for ChangeHeader {
118    const CODE: u8 = Self::CODE;
119
120    fn parse(mut buffer: BytesMut) -> Result<Self, ProtocolError> {
121        let Some(index) = buffer.safe_get_u32() else {
122            return Err(NotEnoughData::new(
123                STAGE_DECODING,
124                "ChangeHeader",
125                "Index byte missing",
126                1,
127                0,
128                buffer,
129            )
130            .into());
131        };
132        let header = Header::parse(buffer)?;
133
134        Ok(Self { index, header })
135    }
136}
137
138impl Writable for ChangeHeader {
139    ///index : uint32 is Index of the occurrence of this header.
140    /// index has to > 0.
141    ///(Note that the "index" above is per-name--i.e. a 3 in this field
142    ///indicates that the modification is to be applied to the third such
143    ///header matching the supplied "name" field.  A zero length string for
144    ///"value", leaving only a single NUL byte, indicates that the header
145    ///should be deleted entirely.)
146    fn write(&self, buffer: &mut BytesMut) {
147        let index = u32::to_be_bytes(self.index);
148        buffer.put_slice(&index);
149        self.header.write(buffer);
150    }
151
152    fn len(&self) -> usize {
153        4 + self.header.len()
154    }
155
156    fn code(&self) -> u8 {
157        Self::CODE
158    }
159    fn is_empty(&self) -> bool {
160        self.header.is_empty()
161    }
162}
163
164/// Insert header at a specified position (modification action)
165#[derive(Debug, Clone)]
166pub struct InsertHeader {
167    index: u32,
168    header: Header,
169}
170
171impl InsertHeader {
172    const CODE: u8 = b'i';
173
174    /// Create a Header from some bytes
175    #[must_use]
176    pub fn new(index: u32, name: &[u8], value: &[u8]) -> Self {
177        Self {
178            index,
179            header: Header::new(name, value),
180        }
181    }
182
183    /// The name of the header
184    #[must_use]
185    pub fn name(&self) -> Cow<str> {
186        self.header.name()
187    }
188
189    /// The value of the header
190    #[must_use]
191    pub fn value(&self) -> Cow<str> {
192        self.header.value()
193    }
194
195    /// The list index at which to insert this header
196    #[must_use]
197    pub fn index(&self) -> u32 {
198        self.index
199    }
200}
201
202impl Parsable for InsertHeader {
203    const CODE: u8 = Self::CODE;
204
205    fn parse(mut buffer: BytesMut) -> Result<Self, ProtocolError> {
206        let Some(index) = buffer.safe_get_u32() else {
207            return Err(NotEnoughData::new(
208                STAGE_DECODING,
209                "InsertHeader",
210                "Index byte missing",
211                1,
212                0,
213                buffer,
214            )
215            .into());
216        };
217        let header = Header::parse(buffer)?;
218
219        Ok(Self { index, header })
220    }
221}
222
223impl Writable for InsertHeader {
224    ///index is Index into header list where insertion should occur
225    fn write(&self, buffer: &mut BytesMut) {
226        let index = u32::to_be_bytes(self.index);
227        buffer.put_slice(&index);
228        self.header.write(buffer);
229    }
230
231    fn len(&self) -> usize {
232        4 + self.header.len()
233    }
234
235    fn code(&self) -> u8 {
236        Self::CODE
237    }
238    fn is_empty(&self) -> bool {
239        self.header.is_empty()
240    }
241}
242
243#[cfg(test)]
244mod test {
245
246    use super::*;
247    use pretty_assertions::assert_eq;
248    use rstest::rstest;
249
250    #[test]
251    fn test_add_header() {
252        let mut buffer = BytesMut::from("h");
253        let add_header = AddHeader {
254            header: Header::new(b"name", b"value"),
255        };
256
257        add_header.write(&mut buffer);
258        assert_eq!(buffer, BytesMut::from("hname\0value\0"));
259    }
260
261    #[rstest]
262    #[case((1, String::from("name"), String::from("value")), BytesMut::from("m\0\0\0\x01name\0value\0"))]
263    #[case((0, String::from("name"), String::from("value")), BytesMut::from("m\0\0\0\0name\0value\0"))]
264    #[case((2, String::from("name"), String::from("\0")), BytesMut::from("m\0\0\0\x02name\0\0\0"))]
265    fn test_change_header(#[case] input: (u32, String, String), #[case] expected: BytesMut) {
266        let mut buffer = BytesMut::from("m");
267
268        let change_header = ChangeHeader {
269            index: input.0,
270            header: Header::new(input.1.as_bytes(), input.2.as_bytes()),
271        };
272
273        change_header.write(&mut buffer);
274
275        assert_eq!(buffer, expected);
276    }
277    #[rstest]
278    #[case((1, String::from("name"), String::from("value")), BytesMut::from("i\0\0\0\x01name\0value\0"))]
279    #[case((0, String::from("name"), String::from("value")), BytesMut::from("i\0\0\0\0name\0value\0"))]
280    #[case((2, String::from("name"), String::from("\0")), BytesMut::from("i\0\0\0\x02name\0\0\0"))]
281    fn test_insert_header(#[case] input: (u32, String, String), #[case] expected: BytesMut) {
282        let mut buffer = BytesMut::from("i");
283
284        let change_header = ChangeHeader {
285            index: input.0,
286            header: Header::new(input.1.as_bytes(), input.2.as_bytes()),
287        };
288
289        change_header.write(&mut buffer);
290
291        assert_eq!(buffer, expected);
292    }
293}