Skip to main content

mini_mail_auth/common/
headers.rs

1use std::{
2    iter::{Enumerate, Peekable},
3    slice::Iter,
4};
5
6pub trait HeaderStream<'x> {
7    fn next_header(&mut self) -> Option<(&'x [u8], &'x [u8])>;
8    fn body(&mut self) -> &'x [u8];
9}
10
11pub(crate) struct HeaderIterator<'x> {
12    message: &'x [u8],
13    iter: Peekable<Enumerate<Iter<'x, u8>>>,
14    start_pos: usize,
15}
16
17impl<'x> HeaderIterator<'x> {
18    pub fn new(message: &'x [u8]) -> Self {
19        HeaderIterator {
20            message,
21            iter: message.iter().enumerate().peekable(),
22            start_pos: 0,
23        }
24    }
25
26    pub fn body_offset(&mut self) -> Option<usize> {
27        self.iter.peek().map(|(pos, _)| *pos)
28    }
29}
30
31impl<'x> HeaderStream<'x> for HeaderIterator<'x> {
32    fn next_header(&mut self) -> Option<(&'x [u8], &'x [u8])> {
33        self.next()
34    }
35
36    fn body(&mut self) -> &'x [u8] {
37        self.body_offset()
38            .and_then(|offset| self.message.get(offset..))
39            .unwrap_or_default()
40    }
41}
42
43impl<'x> Iterator for HeaderIterator<'x> {
44    type Item = (&'x [u8], &'x [u8]);
45
46    fn next(&mut self) -> Option<Self::Item> {
47        let mut colon_pos = usize::MAX;
48        let mut last_ch = 0;
49
50        while let Some((pos, &ch)) = self.iter.next() {
51            if colon_pos == usize::MAX {
52                match ch {
53                    b':' => {
54                        colon_pos = pos;
55                    }
56                    b'\n' => {
57                        if last_ch == b'\r' || self.start_pos == pos {
58                            return None; // End of headers
59                        } else if self
60                            .iter
61                            .peek()
62                            .is_none_or(|(_, next_byte)| ![b' ', b'\t'].contains(next_byte))
63                        {
64                            let header_name = self
65                                .message
66                                .get(self.start_pos..pos + 1)
67                                .unwrap_or_default();
68                            self.start_pos = pos + 1;
69                            return Some((header_name, b""));
70                        }
71                    }
72                    _ => (),
73                }
74            } else if ch == b'\n'
75                && self
76                    .iter
77                    .peek()
78                    .is_none_or(|(_, next_byte)| ![b' ', b'\t'].contains(next_byte))
79            {
80                let header_name = self
81                    .message
82                    .get(self.start_pos..colon_pos)
83                    .unwrap_or_default();
84                let header_value = self.message.get(colon_pos + 1..pos + 1).unwrap_or_default();
85                self.start_pos = pos + 1;
86                return Some((header_name, header_value));
87            }
88            last_ch = ch;
89        }
90        None
91    }
92}
93
94// --- Writer Traits ---
95
96pub trait HeaderWriter: Sized {
97    fn write_header(&self, writer: &mut impl Writer);
98    fn to_header(&self) -> String {
99        let mut buf = Vec::new();
100        self.write_header(&mut buf);
101        String::from_utf8(buf).unwrap()
102    }
103}
104
105pub trait Writable {
106    fn write(self, writer: &mut impl Writer);
107}
108
109impl Writable for &[u8] {
110    fn write(self, writer: &mut impl Writer) {
111        writer.write(self);
112    }
113}
114
115pub trait Writer {
116    fn write(&mut self, buf: &[u8]);
117    fn write_len(&mut self, buf: &[u8], len: &mut usize) {
118        self.write(buf);
119        *len += buf.len();
120    }
121}
122
123impl Writer for Vec<u8> {
124    fn write(&mut self, buf: &[u8]) {
125        self.extend(buf);
126    }
127}