mail_auth/dkim/
streaming.rs

1/*
2 * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
3 *
4 * SPDX-License-Identifier: Apache-2.0 OR MIT
5 */
6
7//! Streaming DKIM1 signing: the message is fed in chunks and only the header
8//! section is buffered, which keeps memory usage low for large messages.
9
10use super::{DkimSigner, Signature, canonicalize::BodyHasher, sign::SignableMessage};
11use crate::SystemTime;
12use crate::signer::Ready;
13use crate::{
14    Error,
15    crypto::{HashContext, HashImpl, SigningKey},
16    headers::HeaderIterator,
17};
18use memchr::memmem;
19
20/// A streaming DKIM1 signer, created by
21/// [`DkimSigner::sign_streaming`].
22///
23/// Avoids loading the entire message into memory before signing. Headers are
24/// buffered until the header/body boundary (`\r\n\r\n`) is found; body
25/// content is then canonicalized and hashed as it arrives. The resulting
26/// [`Signature`] is identical to the one [`DkimSigner::sign`] produces for
27/// the same message and signing time, except that streaming fails with
28/// [`Error::NoHeadersFound`] when none of the
29/// listed headers is present, where `sign` signs them as empty.
30///
31/// # Example
32///
33/// ```rust,no_run
34/// use mail_auth::{crypto::Ed25519Key, dkim::DkimSigner, headers::HeaderWriter};
35///
36/// # fn main() -> mail_auth::Result<()> {
37/// # let pkcs8_der: Vec<u8> = Vec::new();
38/// let signer = DkimSigner::from_key(Ed25519Key::from_pkcs8_der(&pkcs8_der)?)
39///     .domain("example.com")
40///     .selector("default")
41///     .headers(["From", "To", "Subject"]);
42///
43/// let mut stream = signer.sign_streaming();
44/// stream.write(b"From: sender@example.com\r\n");
45/// stream.write(b"To: recipient@example.com\r\n");
46/// stream.write(b"Subject: Test\r\n");
47/// stream.write(b"\r\n");
48/// stream.write(b"Body content here...");
49///
50/// let signature = stream.finish()?;
51/// let header = signature.to_header();
52/// # Ok(())
53/// # }
54/// ```
55pub struct DkimSigningStream<'a, T: SigningKey> {
56    template: Signature,
57    key: &'a T,
58    state: SigningState<<<T as SigningKey>::Hasher as HashImpl>::Context>,
59}
60
61enum SigningState<H> {
62    /// Accumulating headers until \r\n\r\n is found
63    ReadingHeaders { buffer: Vec<u8>, scanned: usize },
64    /// Header section buffered, now hashing body
65    HashingBody {
66        header_section: Vec<u8>,
67        body_hasher: BodyHasher<H>,
68    },
69    /// Finished or consumed
70    Done,
71}
72
73impl<T: SigningKey> DkimSigner<T, Ready> {
74    /// Creates a streaming signer that borrows this signer's key and copies
75    /// its settings.
76    ///
77    /// Feed raw message data with [`DkimSigningStream::write`], then call
78    /// [`DkimSigningStream::finish`] to obtain the signature.
79    ///
80    /// Headers are buffered until the header/body boundary (`\r\n\r\n`) is
81    /// detected. After that, body content is streamed through the hasher
82    /// without additional buffering.
83    ///
84    /// # Example
85    ///
86    /// ```rust,no_run
87    /// use mail_auth::{crypto::Ed25519Key, dkim::DkimSigner};
88    ///
89    /// # fn main() -> mail_auth::Result<()> {
90    /// # let pkcs8_der: Vec<u8> = Vec::new();
91    /// # let message_chunks: Vec<&[u8]> = Vec::new();
92    /// let signer = DkimSigner::from_key(Ed25519Key::from_pkcs8_der(&pkcs8_der)?)
93    ///     .domain("example.com")
94    ///     .selector("default")
95    ///     .headers(["From", "To", "Subject"]);
96    /// let mut stream = signer.sign_streaming();
97    /// for chunk in message_chunks {
98    ///     stream.write(chunk);
99    /// }
100    /// let signature = stream.finish()?;
101    /// # Ok(())
102    /// # }
103    /// ```
104    pub fn sign_streaming(&self) -> DkimSigningStream<'_, T> {
105        DkimSigningStream {
106            template: self.template.clone(),
107            key: &self.key,
108            state: SigningState::ReadingHeaders {
109                buffer: Vec::with_capacity(8192),
110                scanned: 0,
111            },
112        }
113    }
114}
115
116impl<T: SigningKey> DkimSigningStream<'_, T> {
117    /// Feeds the next chunk of raw message data to the signer.
118    ///
119    /// Data must be provided in order, starting with the headers. Chunks may
120    /// split lines, and the header/body boundary (`\r\n\r\n`), anywhere.
121    ///
122    /// While reading headers, all data is buffered. Once the header/body
123    /// boundary is detected, subsequent body data is streamed directly to the
124    /// hasher.
125    pub fn write(&mut self, chunk: &[u8]) {
126        match &mut self.state {
127            SigningState::ReadingHeaders { buffer, scanned } => {
128                buffer.extend_from_slice(chunk);
129
130                let Some(boundary_pos) =
131                    find_header_boundary(&buffer[*scanned..]).map(|pos| *scanned + pos)
132                else {
133                    *scanned = buffer.len().saturating_sub(3);
134                    return;
135                };
136
137                let mut header_section = std::mem::take(buffer);
138
139                let mut body_hasher = BodyHasher::new(
140                    <T::Hasher as HashImpl>::hasher(),
141                    self.template.cb,
142                    if self.template.l > 0 { u64::MAX } else { 0 },
143                );
144
145                body_hasher.write(&header_section[boundary_pos..]);
146                header_section.truncate(boundary_pos - 2);
147
148                self.state = SigningState::HashingBody {
149                    header_section,
150                    body_hasher,
151                };
152            }
153            SigningState::HashingBody { body_hasher, .. } => {
154                body_hasher.write(chunk);
155            }
156            SigningState::Done => {}
157        }
158    }
159
160    /// Consumes the stream and returns the DKIM signature.
161    ///
162    /// The current system time is used for `t=` (and as the base of `x=`
163    /// when an expiration is set). A stream that never saw the header/body
164    /// boundary is treated as a message with headers only and an empty body.
165    ///
166    /// # Errors
167    ///
168    /// Returns [`Error::NoHeadersFound`] when none of the signer's listed
169    /// headers was found, or [`Error::Crypto`] when the
170    /// signing operation fails.
171    pub fn finish(self) -> crate::Result<Signature>
172    where
173        <<T as SigningKey>::Hasher as HashImpl>::Context: HashContext,
174    {
175        let now = SystemTime::now()
176            .duration_since(SystemTime::UNIX_EPOCH)
177            .map(|d| d.as_secs())
178            .unwrap_or(0);
179        self.finish_at(now)
180    }
181
182    pub(crate) fn finish_at(mut self, now: u64) -> crate::Result<Signature>
183    where
184        <<T as SigningKey>::Hasher as HashImpl>::Context: HashContext,
185    {
186        match std::mem::replace(&mut self.state, SigningState::Done) {
187            SigningState::ReadingHeaders { mut buffer, .. } => {
188                let mut body_hasher = BodyHasher::new(
189                    <T::Hasher as HashImpl>::hasher(),
190                    self.template.cb,
191                    if self.template.l > 0 { u64::MAX } else { 0 },
192                );
193
194                let header_len = match find_header_boundary(&buffer) {
195                    Some(boundary_pos) => {
196                        body_hasher.write(&buffer[boundary_pos..]);
197                        boundary_pos - 2
198                    }
199                    None => {
200                        buffer.extend_from_slice(b"\r\n");
201                        buffer.len()
202                    }
203                };
204
205                let (hasher, body_len) = body_hasher.finish();
206                let body_hash = hasher.complete();
207
208                self.finish_with_parsed_data(&buffer[..header_len], body_hash, body_len, now)
209            }
210            SigningState::HashingBody {
211                header_section,
212                body_hasher,
213            } => {
214                let (hasher, body_len) = body_hasher.finish();
215                let body_hash = hasher.complete();
216                self.finish_with_parsed_data(&header_section, body_hash, body_len, now)
217            }
218            SigningState::Done => Err(Error::NoHeadersFound),
219        }
220    }
221
222    fn finish_with_parsed_data(
223        &self,
224        header_section: &[u8],
225        body_hash: crate::crypto::HashOutput,
226        body_len: u64,
227        now: u64,
228    ) -> crate::Result<Signature> {
229        if !self.template.a.is_dkim1()
230            || matches!(
231                self.template.atpsh,
232                Some(crate::crypto::HashAlgorithm::Sha512)
233            )
234        {
235            return Err(Error::Dkim(super::DkimError::UnsupportedAlgorithm));
236        }
237
238        let mut headers = Vec::with_capacity(self.template.h.len());
239        let mut found_headers = vec![false; self.template.h.len()];
240        let mut signed_headers = Vec::with_capacity(self.template.h.len());
241
242        for (name, value) in HeaderIterator::new(header_section) {
243            if let Some(pos) = self
244                .template
245                .h
246                .iter()
247                .position(|header| name.eq_ignore_ascii_case(header.as_bytes()))
248            {
249                headers.push((name, value));
250                found_headers[pos] = true;
251                signed_headers.push(std::str::from_utf8(name).unwrap_or_default().to_string());
252            }
253        }
254
255        if signed_headers.is_empty() {
256            return Err(Error::NoHeadersFound);
257        }
258
259        signed_headers.reverse();
260        for (header, found) in self.template.h.iter().zip(found_headers) {
261            if !found {
262                signed_headers.push(header.to_string());
263            }
264        }
265
266        let canonical_headers = self.template.ch.canonical_headers(headers);
267
268        let mut signature = self.template.clone();
269        signature.bh = body_hash.as_ref().to_vec();
270        signature.t = now;
271        signature.x = if signature.x > 0 {
272            now.saturating_add(signature.x)
273        } else {
274            0
275        };
276        signature.h = signed_headers;
277        if signature.l > 0 {
278            signature.l = body_len;
279        }
280
281        signature.b = self.key.sign(SignableMessage {
282            headers: canonical_headers,
283            signature: &signature,
284        })?;
285
286        Ok(signature)
287    }
288}
289
290/// Find the header/body boundary (\r\n\r\n) and return the position after it
291pub(crate) fn find_header_boundary(data: &[u8]) -> Option<usize> {
292    memmem::find(data, b"\r\n\r\n").map(|p| p + 4)
293}
294
295#[cfg(test)]
296#[allow(unused)]
297mod test {
298    use crate::{
299        crypto::{RsaKey, Sha256},
300        dkim::{Canonicalization, DkimSigner},
301        headers::HeaderIterator,
302    };
303
304    use rustls_pki_types::{PrivateKeyDer, PrivatePkcs1KeyDer, pem::PemObject};
305
306    const RSA_PRIVATE_KEY: &str = include_str!("../../resources/rsa-private.pem");
307    const TIMESTAMP: u64 = 1740000000;
308
309    #[test]
310    fn streaming_sign_matches_regular_sign() {
311        let message = concat!(
312            "From: bill@example.com\r\n",
313            "To: jdoe@example.com\r\n",
314            "Subject: TPS Report\r\n",
315            "\r\n",
316            "I'm going to need those TPS reports ASAP. ",
317            "So, if you could do that, that'd be great.\r\n"
318        );
319
320        let pk_rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
321            PrivatePkcs1KeyDer::from_pem_slice(RSA_PRIVATE_KEY.as_bytes()).unwrap(),
322        ))
323        .unwrap();
324
325        let signer = DkimSigner::from_key(pk_rsa)
326            .domain("example.com")
327            .selector("default")
328            .headers(["From", "To", "Subject"]);
329
330        let sig1 = signer
331            .sign_stream(HeaderIterator::new(message.as_bytes()), TIMESTAMP)
332            .unwrap();
333
334        let mut stream = signer.sign_streaming();
335        stream.write(message.as_bytes());
336        let sig2 = stream.finish_at(TIMESTAMP).unwrap();
337
338        assert_eq!(sig1.bh, sig2.bh, "Body hashes should match");
339        assert_eq!(sig1.h, sig2.h, "Signed headers should match");
340        assert_eq!(sig1.b, sig2.b, "Signatures should match");
341    }
342
343    #[test]
344    fn streaming_sign_multiple_chunks() {
345        let header = "From: bill@example.com\r\nTo: jdoe@example.com\r\nSubject: Test\r\n\r\n";
346        let body = "Hello World! This is the body.\r\n";
347
348        let pk_rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
349            PrivatePkcs1KeyDer::from_pem_slice(RSA_PRIVATE_KEY.as_bytes()).unwrap(),
350        ))
351        .unwrap();
352
353        let signer = DkimSigner::from_key(pk_rsa)
354            .domain("example.com")
355            .selector("default")
356            .headers(["From", "To", "Subject"]);
357
358        let full_message = format!("{}{}", header, body);
359        let reference_sig = signer.sign(full_message.as_bytes()).unwrap();
360
361        let mut stream = signer.sign_streaming();
362        stream.write(header.as_bytes());
363        stream.write(body.as_bytes());
364        let streamed_sig = stream.finish().unwrap();
365
366        assert_eq!(
367            reference_sig.bh, streamed_sig.bh,
368            "Body hashes should match"
369        );
370    }
371
372    #[test]
373    fn streaming_sign_chunked_body() {
374        let message = concat!(
375            "From: test@example.com\r\n",
376            "Subject: Chunked Test\r\n",
377            "\r\n",
378            "Line 1\r\n",
379            "Line 2\r\n",
380            "Line 3\r\n",
381        );
382
383        let pk_rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
384            PrivatePkcs1KeyDer::from_pem_slice(RSA_PRIVATE_KEY.as_bytes()).unwrap(),
385        ))
386        .unwrap();
387
388        let signer = DkimSigner::from_key(pk_rsa)
389            .domain("example.com")
390            .selector("default")
391            .headers(["From", "Subject"]);
392
393        let reference_sig = signer.sign(message.as_bytes()).unwrap();
394
395        for chunk_size in [1, 2, 5, 10, 20] {
396            let mut stream = signer.sign_streaming();
397            for chunk in message.as_bytes().chunks(chunk_size) {
398                stream.write(chunk);
399            }
400            let streamed_sig = stream.finish().unwrap();
401
402            assert_eq!(
403                reference_sig.bh, streamed_sig.bh,
404                "Body hash mismatch at chunk_size={}",
405                chunk_size
406            );
407        }
408    }
409
410    #[test]
411    fn streaming_sign_split_header_boundary() {
412        let pk_rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
413            PrivatePkcs1KeyDer::from_pem_slice(RSA_PRIVATE_KEY.as_bytes()).unwrap(),
414        ))
415        .unwrap();
416
417        let signer = DkimSigner::from_key(pk_rsa)
418            .domain("example.com")
419            .selector("default")
420            .headers(["From", "Subject"]);
421
422        let message = "From: test@example.com\r\nSubject: Test\r\n\r\nBody";
423        let reference_sig = signer.sign(message.as_bytes()).unwrap();
424
425        let mut stream = signer.sign_streaming();
426        stream.write(b"From: test@example.com\r\n");
427        stream.write(b"Subject: Test\r\n");
428        stream.write(b"\r\n");
429        stream.write(b"Body");
430        let streamed_sig = stream.finish().unwrap();
431
432        assert_eq!(reference_sig.bh, streamed_sig.bh);
433    }
434
435    #[test]
436    fn streaming_sign_empty_body() {
437        let message = "From: test@example.com\r\nSubject: Empty\r\n\r\n";
438        let pk_rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
439            PrivatePkcs1KeyDer::from_pem_slice(RSA_PRIVATE_KEY.as_bytes()).unwrap(),
440        ))
441        .unwrap();
442
443        let signer = DkimSigner::from_key(pk_rsa)
444            .domain("example.com")
445            .selector("default")
446            .headers(["From", "Subject"]);
447
448        let reference_sig = signer.sign(message.as_bytes()).unwrap();
449
450        let mut stream = signer.sign_streaming();
451        stream.write(message.as_bytes());
452        let streamed_sig = stream.finish().unwrap();
453
454        assert_eq!(reference_sig.bh, streamed_sig.bh);
455    }
456
457    #[test]
458    fn streaming_sign_simple_canonicalization() {
459        let message = concat!(
460            "From: test@example.com\r\n",
461            "Subject: Simple Canon Test\r\n",
462            "\r\n",
463            "Body with   spaces\r\n",
464        );
465
466        let pk_rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
467            PrivatePkcs1KeyDer::from_pem_slice(RSA_PRIVATE_KEY.as_bytes()).unwrap(),
468        ))
469        .unwrap();
470
471        let signer = DkimSigner::from_key(pk_rsa)
472            .domain("example.com")
473            .selector("default")
474            .headers(["From", "Subject"])
475            .header_canonicalization(Canonicalization::Simple)
476            .body_canonicalization(Canonicalization::Simple);
477
478        let reference_sig = signer
479            .sign_stream(HeaderIterator::new(message.as_bytes()), TIMESTAMP)
480            .unwrap();
481
482        let mut stream = signer.sign_streaming();
483        stream.write(message.as_bytes());
484        let streamed_sig = stream.finish_at(TIMESTAMP).unwrap();
485
486        assert_eq!(reference_sig.bh, streamed_sig.bh);
487        assert_eq!(reference_sig.b, streamed_sig.b);
488    }
489
490    #[test]
491    fn streaming_sign_folded_headers() {
492        let message = concat!(
493            "From: test@example.com\r\n",
494            "Subject: This is a very long subject line that\r\n",
495            " continues on the next line\r\n",
496            "\r\n",
497            "Body\r\n",
498        );
499        let pk_rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
500            PrivatePkcs1KeyDer::from_pem_slice(RSA_PRIVATE_KEY.as_bytes()).unwrap(),
501        ))
502        .unwrap();
503
504        let signer = DkimSigner::from_key(pk_rsa)
505            .domain("example.com")
506            .selector("default")
507            .headers(["From", "Subject"]);
508
509        let reference_sig = signer.sign(message.as_bytes()).unwrap();
510
511        let mut stream = signer.sign_streaming();
512        stream.write(message.as_bytes());
513        let streamed_sig = stream.finish().unwrap();
514
515        assert_eq!(reference_sig.bh, streamed_sig.bh);
516    }
517
518    #[test]
519    fn streaming_sign_no_matching_headers_error() {
520        let message = "X-Custom: value\r\n\r\nBody\r\n";
521
522        let pk_rsa = RsaKey::<Sha256>::from_key_der(PrivateKeyDer::Pkcs1(
523            PrivatePkcs1KeyDer::from_pem_slice(RSA_PRIVATE_KEY.as_bytes()).unwrap(),
524        ))
525        .unwrap();
526
527        let signer = DkimSigner::from_key(pk_rsa)
528            .domain("example.com")
529            .selector("default")
530            .headers(["From", "Subject"]);
531
532        let mut stream = signer.sign_streaming();
533        stream.write(message.as_bytes());
534        let result = stream.finish();
535
536        assert!(matches!(result, Err(crate::Error::NoHeadersFound)));
537    }
538}
mail_auth/dkim/streaming.rs

mail_auth/dkim/
streaming.rs
