Commit
+122 -26 +/-3 browse
1 | diff --git a/src/common/headers.rs b/src/common/headers.rs |
2 | index 56b9ce0..4b24ea2 100644 |
3 | --- a/src/common/headers.rs |
4 | +++ b/src/common/headers.rs |
5 | @@ -23,6 +23,16 @@ impl<'x, T> Header<'x, T> { |
6 | } |
7 | } |
8 | |
9 | + pub trait HeaderStream<'x> { |
10 | + fn next_header(&mut self) -> Option<(&'x [u8], &'x [u8])>; |
11 | + fn body(&mut self) -> &'x [u8]; |
12 | + } |
13 | + |
14 | + pub(crate) struct ChainedHeaderIterator<'x, T: Iterator<Item = &'x [u8]>> { |
15 | + parts: T, |
16 | + iter: HeaderIterator<'x>, |
17 | + } |
18 | + |
19 | pub(crate) struct HeaderIterator<'x> { |
20 | message: &'x [u8], |
21 | iter: Peekable<Enumerate<Iter<'x, u8>>>, |
22 | @@ -96,6 +106,18 @@ impl<'x> HeaderIterator<'x> { |
23 | } |
24 | } |
25 | |
26 | + impl<'x> HeaderStream<'x> for HeaderIterator<'x> { |
27 | + fn next_header(&mut self) -> Option<(&'x [u8], &'x [u8])> { |
28 | + self.next() |
29 | + } |
30 | + |
31 | + fn body(&mut self) -> &'x [u8] { |
32 | + self.body_offset() |
33 | + .and_then(|offset| self.message.get(offset..)) |
34 | + .unwrap_or_default() |
35 | + } |
36 | + } |
37 | + |
38 | impl<'x> Iterator for HeaderIterator<'x> { |
39 | type Item = (&'x [u8], &'x [u8]); |
40 | |
41 | @@ -153,6 +175,30 @@ impl<'x> Iterator for HeaderIterator<'x> { |
42 | } |
43 | } |
44 | |
45 | + impl<'x, T: Iterator<Item = &'x [u8]>> ChainedHeaderIterator<'x, T> { |
46 | + pub fn new(mut parts: T) -> Self { |
47 | + ChainedHeaderIterator { |
48 | + iter: HeaderIterator::new(parts.next().unwrap()), |
49 | + parts, |
50 | + } |
51 | + } |
52 | + } |
53 | + |
54 | + impl<'x, T: Iterator<Item = &'x [u8]>> HeaderStream<'x> for ChainedHeaderIterator<'x, T> { |
55 | + fn next_header(&mut self) -> Option<(&'x [u8], &'x [u8])> { |
56 | + if let Some(header) = self.iter.next_header() { |
57 | + Some(header) |
58 | + } else { |
59 | + self.iter = HeaderIterator::new(self.parts.next()?); |
60 | + self.iter.next_header() |
61 | + } |
62 | + } |
63 | + |
64 | + fn body(&mut self) -> &'x [u8] { |
65 | + self.iter.body() |
66 | + } |
67 | + } |
68 | + |
69 | impl<'x> Iterator for HeaderParser<'x> { |
70 | type Item = (AuthenticatedHeader<'x>, &'x [u8]); |
71 | |
72 | @@ -367,7 +413,7 @@ const MSGID: u64 = (b'm' as u64) |
73 | mod test { |
74 | use crate::common::headers::{AuthenticatedHeader, HeaderParser}; |
75 | |
76 | - use super::HeaderIterator; |
77 | + use super::{ChainedHeaderIterator, HeaderIterator, HeaderStream}; |
78 | |
79 | #[test] |
80 | fn header_iterator() { |
81 | @@ -480,4 +526,33 @@ mod test { |
82 | assert!(parser.has_message_id); |
83 | assert_eq!(parser.num_received, 3); |
84 | } |
85 | + |
86 | + #[test] |
87 | + fn chained_header_iterator() { |
88 | + let parts = vec![ |
89 | + &b"From: a\nTo: b\nEmpty:\nMulti: 1\n 2\n"[..], |
90 | + &b"Subject: c\nReceived: d\n\nhey"[..], |
91 | + ]; |
92 | + let mut headers = vec![ |
93 | + ("From", " a\n"), |
94 | + ("To", " b\n"), |
95 | + ("Empty", "\n"), |
96 | + ("Multi", " 1\n 2\n"), |
97 | + ("Subject", " c\n"), |
98 | + ("Received", " d\n"), |
99 | + ] |
100 | + .into_iter(); |
101 | + let mut it = ChainedHeaderIterator::new(parts.iter().copied()); |
102 | + |
103 | + while let Some((k, v)) = it.next_header() { |
104 | + assert_eq!( |
105 | + ( |
106 | + std::str::from_utf8(k).unwrap(), |
107 | + std::str::from_utf8(v).unwrap() |
108 | + ), |
109 | + headers.next().unwrap() |
110 | + ); |
111 | + } |
112 | + assert_eq!(it.body(), b"hey"); |
113 | + } |
114 | } |
115 | diff --git a/src/dkim/canonicalize.rs b/src/dkim/canonicalize.rs |
116 | index 16b6b4f..a0a8bd1 100644 |
117 | --- a/src/dkim/canonicalize.rs |
118 | +++ b/src/dkim/canonicalize.rs |
119 | @@ -10,7 +10,7 @@ |
120 | |
121 | use crate::common::{ |
122 | crypto::{HashContext, HashImpl}, |
123 | - headers::{HeaderIterator, Writer}, |
124 | + headers::{HeaderStream, Writer}, |
125 | }; |
126 | |
127 | use super::{Canonicalization, Signature}; |
128 | @@ -147,18 +147,17 @@ impl Canonicalization { |
129 | |
130 | impl Signature { |
131 | #[allow(clippy::while_let_on_iterator)] |
132 | - pub(crate) fn canonicalize( |
133 | + pub(crate) fn canonicalize<'x>( |
134 | &self, |
135 | - message: &[u8], |
136 | + mut message: impl HeaderStream<'x>, |
137 | header_hasher: &mut impl Writer, |
138 | body_hasher: &mut impl Writer, |
139 | ) -> (usize, Vec<String>) { |
140 | - let mut headers_it = HeaderIterator::new(message); |
141 | let mut headers = Vec::with_capacity(self.h.len()); |
142 | let mut found_headers = vec![false; self.h.len()]; |
143 | let mut signed_headers = Vec::with_capacity(self.h.len()); |
144 | |
145 | - for (name, value) in &mut headers_it { |
146 | + while let Some((name, value)) = message.next_header() { |
147 | if let Some(pos) = self |
148 | .h |
149 | .iter() |
150 | @@ -170,10 +169,7 @@ impl Signature { |
151 | } |
152 | } |
153 | |
154 | - let body = headers_it |
155 | - .body_offset() |
156 | - .and_then(|pos| message.get(pos..)) |
157 | - .unwrap_or_default(); |
158 | + let body = message.body(); |
159 | let body_len = body.len(); |
160 | self.ch |
161 | .canonicalize_headers(&mut headers.into_iter().rev(), header_hasher); |
162 | diff --git a/src/dkim/sign.rs b/src/dkim/sign.rs |
163 | index ca2f0af..43c2e6e 100644 |
164 | --- a/src/dkim/sign.rs |
165 | +++ b/src/dkim/sign.rs |
166 | @@ -14,7 +14,10 @@ use mail_builder::encoders::base64::base64_encode; |
167 | |
168 | use super::{DkimSigner, Done, Signature}; |
169 | use crate::{ |
170 | - common::crypto::{HashContext, SigningKey}, |
171 | + common::{ |
172 | + crypto::{HashContext, SigningKey}, |
173 | + headers::{ChainedHeaderIterator, HeaderIterator, HeaderStream}, |
174 | + }, |
175 | Error, |
176 | }; |
177 | |
178 | @@ -22,8 +25,23 @@ impl<T: SigningKey> DkimSigner<T, Done> { |
179 | /// Signs a message. |
180 | #[inline(always)] |
181 | pub fn sign(&self, message: &[u8]) -> crate::Result<Signature> { |
182 | - self.sign_( |
183 | - message, |
184 | + self.sign_stream( |
185 | + HeaderIterator::new(message), |
186 | + SystemTime::now() |
187 | + .duration_since(SystemTime::UNIX_EPOCH) |
188 | + .map(|d| d.as_secs()) |
189 | + .unwrap_or(0), |
190 | + ) |
191 | + } |
192 | + |
193 | + #[inline(always)] |
194 | + /// Signs a chained message. |
195 | + pub fn sign_chained<'x>( |
196 | + &self, |
197 | + chunks: impl Iterator<Item = &'x [u8]>, |
198 | + ) -> crate::Result<Signature> { |
199 | + self.sign_stream( |
200 | + ChainedHeaderIterator::new(chunks), |
201 | SystemTime::now() |
202 | .duration_since(SystemTime::UNIX_EPOCH) |
203 | .map(|d| d.as_secs()) |
204 | @@ -31,7 +49,11 @@ impl<T: SigningKey> DkimSigner<T, Done> { |
205 | ) |
206 | } |
207 | |
208 | - fn sign_(&self, message: &[u8], now: u64) -> crate::Result<Signature> { |
209 | + fn sign_stream<'x>( |
210 | + &self, |
211 | + message: impl HeaderStream<'x>, |
212 | + now: u64, |
213 | + ) -> crate::Result<Signature> { |
214 | let mut body_hasher = self.key.hasher(); |
215 | let mut header_hasher = self.key.hasher(); |
216 | |
217 | @@ -81,6 +103,7 @@ mod test { |
218 | use crate::{ |
219 | common::{ |
220 | crypto::{Ed25519Key, RsaKey, Sha256}, |
221 | + headers::HeaderIterator, |
222 | parse::TxtRecordParser, |
223 | verify::DomainKey, |
224 | }, |
225 | @@ -123,14 +146,16 @@ GMot/L2x0IYyMLAz6oLWh2hm7zwtb0CgOrPo1ke44hFYnfc= |
226 | .domain("stalw.art") |
227 | .selector("default") |
228 | .headers(["From", "To", "Subject"]) |
229 | - .sign_( |
230 | - concat!( |
231 | - "From: hello@stalw.art\r\n", |
232 | - "To: dkim@stalw.art\r\n", |
233 | - "Subject: Testing DKIM!\r\n\r\n", |
234 | - "Here goes the test\r\n\r\n" |
235 | - ) |
236 | - .as_bytes(), |
237 | + .sign_stream( |
238 | + HeaderIterator::new( |
239 | + concat!( |
240 | + "From: hello@stalw.art\r\n", |
241 | + "To: dkim@stalw.art\r\n", |
242 | + "Subject: Testing DKIM!\r\n\r\n", |
243 | + "Here goes the test\r\n\r\n" |
244 | + ) |
245 | + .as_bytes(), |
246 | + ), |
247 | 311923920, |
248 | ) |
249 | .unwrap(); |
250 | @@ -289,7 +314,7 @@ GMot/L2x0IYyMLAz6oLWh2hm7zwtb0CgOrPo1ke44hFYnfc= |
251 | .headers(["From", "To", "Subject"]) |
252 | .expiration(12345) |
253 | .reporting(true) |
254 | - .sign_(message.as_bytes(), 12345) |
255 | + .sign_stream(HeaderIterator::new(message.as_bytes()), 12345) |
256 | .unwrap(), |
257 | message, |
258 | Err(super::Error::SignatureExpired), |
259 | @@ -309,7 +334,7 @@ GMot/L2x0IYyMLAz6oLWh2hm7zwtb0CgOrPo1ke44hFYnfc= |
260 | .headers(["From", "To", "Subject"]) |
261 | .atps("example.com") |
262 | .atpsh(HashAlgorithm::Sha256) |
263 | - .sign_(message.as_bytes(), 12345) |
264 | + .sign_stream(HeaderIterator::new(message.as_bytes()), 12345) |
265 | .unwrap(), |
266 | message, |
267 | Err(super::Error::DnsRecordNotFound(ResponseCode::NXDomain)), |
268 | @@ -330,7 +355,7 @@ GMot/L2x0IYyMLAz6oLWh2hm7zwtb0CgOrPo1ke44hFYnfc= |
269 | .headers(["From", "To", "Subject"]) |
270 | .atps("example.com") |
271 | .atpsh(HashAlgorithm::Sha256) |
272 | - .sign_(message.as_bytes(), 12345) |
273 | + .sign_stream(HeaderIterator::new(message.as_bytes()), 12345) |
274 | .unwrap(), |
275 | message, |
276 | Ok(()), |
277 | @@ -350,7 +375,7 @@ GMot/L2x0IYyMLAz6oLWh2hm7zwtb0CgOrPo1ke44hFYnfc= |
278 | .selector("default") |
279 | .headers(["From", "To", "Subject"]) |
280 | .atps("example.com") |
281 | - .sign_(message.as_bytes(), 12345) |
282 | + .sign_stream(HeaderIterator::new(message.as_bytes()), 12345) |
283 | .unwrap(), |
284 | message, |
285 | Ok(()), |