reduce heap allocations during hashing

This doesn't totally eliminate them, but for the body at least,
we can stream and hash the data and avoid making potentially
very many very large allocations in a busy signing application.

Cargo.lock

@@ -133,6 +133,7 @@ dependencies = [
  "mailparse",
  "memchr",
  "nom",
+ "once_cell",
  "quick-error 2.0.1",
  "regex",
  "rsa",

Cargo.toml

@@ -16,6 +16,7 @@ ed25519-dalek = "1.0.1"
 mailparse = "0.14"
 memchr = "2.5"
 quick-error = "2.0.1"
+once_cell = "1.17"
 nom = "7.1.0"
 chrono = { version = "0.4.19", default-features = false, features = ["clock", "std"] }
 trust-dns-resolver = "0.22"

src/bytes.rs

@@ -1,45 +1,35 @@
 ///! Various utility functions to operate on bytes
 pub(crate) use memchr::memmem::find;
 
-pub(crate) fn get_all_after<'a>(bytes: &'a [u8], end: &[u8]) -> &'a [u8] {
-    if let Some(mut end_index) = find(bytes, end) {
-        end_index += end.len();
-        &bytes[end_index..]
-    } else {
-        &[]
+pub(crate) fn replace(bytes: &mut [u8], from: u8, to: u8) {
+    let mut previous = 0;
+    while let Some(idx) = memchr::memchr(from, &bytes[previous..]) {
+        bytes[previous + idx] = to;
+        previous = idx + 1;
     }
 }
 
-pub(crate) fn replace(bytes: &mut [u8], from: char, to: char) {
-    for byte in bytes.iter_mut() {
-        if *byte == from as u8 {
-            *byte = to as u8;
-        }
-    }
-}
-
-pub(crate) fn replace_slice(source: &[u8], from: &[u8], to: &[u8]) -> Vec<u8> {
-    let mut result = source.to_vec();
+pub(crate) fn replace_within_vec(result: &mut Vec<u8>, from: &[u8], to: &[u8]) {
     let from_len = from.len();
     let to_len = to.len();
 
     let mut i = 0;
-    while i + from_len <= result.len() {
-        if result[i..].starts_with(from) {
-            result.splice(i..i + from_len, to.iter().cloned());
-            i += to_len;
-        } else {
-            i += 1;
-        }
+    while let Some(idx) = find(&result[i..], from) {
+        result.splice(idx + i..idx + i + from_len, to.iter().cloned());
+        i += idx + to_len;
     }
-
-    result
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    fn replace_slice(source: &[u8], from: &[u8], to: &[u8]) -> Vec<u8> {
+        let mut result = source.to_vec();
+        replace_within_vec(&mut result, from, to);
+        result
+    }
+
     #[test]
     fn it_find() {
         assert_eq!(find(&[97, 98, 99], &[1]), None);
@@ -47,10 +37,26 @@ mod tests {
         assert_eq!(find(&[97, 98, 99], &[97, 98]), Some(0));
     }
 
+    #[test]
+    fn test_replace() {
+        let mut data = b"abbcb".to_vec();
+        replace(&mut data, b'b', b'_');
+        assert_eq!(data, b"a__c_");
+    }
+
     #[test]
     fn it_replace_slice() {
         let source = "aba".as_bytes();
         assert_eq!(replace_slice(source, &[97], &[99]), "cbc".as_bytes());
         assert_eq!(replace_slice(source, &[97, 98], &[]), "a".as_bytes());
+
+        let source = "hello\r\nthere\r\n".as_bytes();
+        assert_eq!(replace_slice(source, b"\r\n", b""), "hellothere".as_bytes());
+
+        let source = "hello there".as_bytes();
+        assert_eq!(
+            replace_slice(source, b"\r\n", b""),
+            "hello there".as_bytes()
+        );
     }
 }

src/canonicalization.rs

@@ -1,5 +1,8 @@
-// Inspired from https://docs.rs/dkim/latest/src/dkim/canonicalization.rs.html
+//! Inspired from https://docs.rs/dkim/latest/src/dkim/canonicalization.rs.html
 use crate::bytes;
+use crate::hash::LimitHasher;
+use memchr::memmem::Finder;
+use once_cell::sync::Lazy;
 
 #[derive(PartialEq, Clone, Debug)]
 pub enum Type {
@@ -15,96 +18,140 @@ impl std::string::ToString for Type {
     }
 }
 
-/// Canonicalize body using the simple canonicalization algorithm.
-///
-/// The first argument **must** be the body of the mail.
-pub(crate) fn canonicalize_body_simple(mut body: &[u8]) -> Vec<u8> {
+fn do_body_simple<'a>(mut body: &'a [u8]) -> &'a [u8] {
     if body.is_empty() {
-        return b"\r\n".to_vec();
+        return b"\r\n";
     }
 
     while body.ends_with(b"\r\n\r\n") {
         body = &body[..body.len() - 2];
     }
 
-    body.to_vec()
-}
-
-/// https://datatracker.ietf.org/doc/html/rfc6376#section-3.4.3
-/// Canonicalize body using the relaxed canonicalization algorithm.  
-///
-/// The first argument **must** be the body of the mail.
-pub(crate) fn canonicalize_body_relaxed(body: &[u8]) -> Vec<u8> {
-    let mut body = body.to_vec();
-    // See https://tools.ietf.org/html/rfc6376#section-3.4.4 for implementation details
-
-    // Reduce all sequences of WSP within a line to a single SP character.
-    bytes::replace(&mut body, '\t', ' ');
-    let mut previous = false;
-    body.retain(|c| {
-        if *c == b' ' {
-            if previous {
-                false
-            } else {
-                previous = true;
-                true
-            }
-        } else {
-            previous = false;
-            true
-        }
-    });
-
-    // Ignore all whitespace at the end of lines. Implementations MUST NOT remove the CRLF at the end of the line.
-    while let Some(idx) = bytes::find(&body, b" \r\n") {
-        body.remove(idx);
-    }
-
-    // Ignore all empty lines at the end of the message body. "Empty line" is defined in Section 3.4.3.
-    while body.ends_with(b"\r\n\r\n") {
-        body.remove(body.len() - 1);
-        body.remove(body.len() - 1);
-    }
-
-    // If the body is non-empty but does not end with a CRLF, a CRLF is added. (For email, this is only possible when using extensions to SMTP or non-SMTP transport mechanisms.)
-    if !body.is_empty() && !body.ends_with(b"\r\n") {
-        body.push(b'\r');
-        body.push(b'\n');
-    }
-
     body
 }
 
+/// Canonicalize body using the simple canonicalization algorithm.
+pub(crate) fn canonicalize_body_simple(body: &[u8], hasher: &mut LimitHasher) {
+    let body = do_body_simple(body);
+    hasher.hash(body);
+}
+
+/// Helper for iterating lines using memmem
+struct IterLines<'haystack> {
+    haystack: &'haystack [u8],
+    inner: memchr::memmem::FindIter<'haystack, 'static>,
+    start: usize,
+    done: bool,
+}
+
+impl<'haystack> Iterator for IterLines<'haystack> {
+    type Item = &'haystack [u8];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.done {
+            return None;
+        }
+
+        match self.inner.next() {
+            Some(idx) => {
+                let line = &self.haystack[self.start..idx + 2];
+                self.start = idx + 2;
+                Some(line)
+            }
+            None => {
+                self.done = true;
+                let line = &self.haystack[self.start..];
+                if line.is_empty() {
+                    None
+                } else {
+                    Some(line)
+                }
+            }
+        }
+    }
+}
+
+fn iter_lines(haystack: &[u8]) -> IterLines {
+    static CRLF: Lazy<Finder> = Lazy::new(|| memchr::memmem::Finder::new("\r\n"));
+    IterLines {
+        haystack,
+        inner: CRLF.find_iter(haystack),
+        start: 0,
+        done: false,
+    }
+}
+
+/// https://datatracker.ietf.org/doc/html/rfc6376#section-3.4.3
+/// Canonicalize body using the relaxed canonicalization algorithm.
+pub(crate) fn apply_body_relaxed(mut body: &[u8], hasher: &mut LimitHasher) {
+    if body.is_empty() {
+        return;
+    }
+
+    // Ignore empty lines at the end of the message body
+    while body.ends_with(b"\r\n\r\n") {
+        body = &body[..body.len() - 2];
+    }
+
+    for mut line in iter_lines(body) {
+        // Ignore all whitespace at the end of the line
+        while let Some(c) = line.last() {
+            match c {
+                b' ' | b'\t' | b'\r' | b'\n' => {
+                    line = &line[0..line.len() - 1];
+                }
+                _ => break,
+            }
+        }
+
+        let mut prior = 0;
+        // Reduce all sequences of WSP within a line to a single SP character.
+        for idx in memchr::memchr2_iter(b' ', b'\t', line) {
+            if prior > 0 && idx == prior {
+                // Part of a run; ignore this one
+                prior = idx + 1;
+                continue;
+            }
+
+            // Found a new run of space(s).
+            // Emit the bytes ahead of this one
+            hasher.hash(&line[prior..idx]);
+            // and emit the canonical space
+            hasher.hash(b" ");
+
+            prior = idx + 1;
+        }
+        // and emit the remainder
+        hasher.hash(&line[prior..]);
+
+        // and canonical newline
+        hasher.hash(b"\r\n");
+    }
+}
+
 // https://datatracker.ietf.org/doc/html/rfc6376#section-3.4.1
-pub(crate) fn canonicalize_header_simple(key: &str, value: &[u8]) -> Vec<u8> {
-    let mut out = Vec::new();
+pub(crate) fn canonicalize_header_simple(key: &str, value: &[u8], out: &mut Vec<u8>) {
     out.extend_from_slice(key.as_bytes());
     out.extend_from_slice(b": ");
     out.extend_from_slice(value);
     out.extend_from_slice(b"\r\n");
-
-    out
 }
 
 // https://datatracker.ietf.org/doc/html/rfc6376#section-3.4.2
-pub(crate) fn canonicalize_header_relaxed(key: &str, value: &[u8]) -> Vec<u8> {
+pub(crate) fn canonicalize_header_relaxed(key: &str, value: &[u8], out: &mut Vec<u8>) {
     let key = key.to_lowercase();
     let key = key.trim_end();
-    let value = canonicalize_header_value_relaxed(value);
 
-    let mut out = Vec::new();
     out.extend_from_slice(key.as_bytes());
     out.extend_from_slice(b":");
-    out.extend_from_slice(&value);
+    out.extend_from_slice(&canonicalize_header_value_relaxed(value));
     out.extend_from_slice(b"\r\n");
-
-    out
 }
 
 fn canonicalize_header_value_relaxed(value: &[u8]) -> Vec<u8> {
     let mut value = value.to_vec();
-    bytes::replace(&mut value, '\t', ' ');
-    value = bytes::replace_slice(&value, b"\r\n", b"");
+    bytes::replace(&mut value, b'\t', b' ');
+    bytes::replace_within_vec(&mut value, b"\r\n", b"");
 
     while value.ends_with(b" ") {
         value.remove(value.len() - 1);
@@ -134,25 +181,59 @@ fn canonicalize_header_value_relaxed(value: &[u8]) -> Vec<u8> {
 mod tests {
     use super::*;
 
+    fn header_relaxed(key: &str, value: &[u8]) -> Vec<u8> {
+        let mut result = vec![];
+        canonicalize_header_relaxed(key, value, &mut result);
+        result
+    }
+
     #[test]
     fn test_canonicalize_header_relaxed() {
+        assert_eq!(header_relaxed("SUBJect", b" AbC\r\n"), b"subject:AbC\r\n");
         assert_eq!(
-            canonicalize_header_relaxed("SUBJect", b" AbC\r\n"),
-            b"subject:AbC\r\n"
-        );
-        assert_eq!(
-            canonicalize_header_relaxed("Subject \t", b"\t Your Name\t \r\n"),
+            header_relaxed("Subject \t", b"\t Your Name\t \r\n"),
             b"subject:Your Name\r\n"
         );
         assert_eq!(
-            canonicalize_header_relaxed("Subject \t", b"\t Kimi \t \r\n No \t\r\n Na Wa\r\n"),
+            header_relaxed("Subject \t", b"\t Kimi \t \r\n No \t\r\n Na Wa\r\n"),
             b"subject:Kimi No Na Wa\r\n"
         );
     }
 
+    fn body_relaxed(data: &[u8]) -> Vec<u8> {
+        let mut hasher = LimitHasher {
+            hasher: crate::hash::HashImpl::copy_data(),
+            limit: usize::MAX,
+            hashed: 0,
+        };
+        apply_body_relaxed(data, &mut hasher);
+        hasher.finalize_bytes()
+    }
+
+    fn body_simple(data: &[u8]) -> Vec<u8> {
+        let mut hasher = LimitHasher {
+            hasher: crate::hash::HashImpl::copy_data(),
+            limit: usize::MAX,
+            hashed: 0,
+        };
+        canonicalize_body_simple(data, &mut hasher);
+        hasher.finalize_bytes()
+    }
+
     #[test]
     fn test_canonicalize_body_relaxed() {
-        assert_eq!(canonicalize_body_relaxed(b"\r\n"), b"\r\n");
-        assert_eq!(canonicalize_body_relaxed(b"hey        \r\n"), b"hey\r\n");
+        assert_eq!(body_relaxed(b"\r\n"), b"\r\n");
+        assert_eq!(body_relaxed(b"hey        \r\n"), b"hey\r\n");
+        assert_eq!(body_relaxed(b" C \r\nD \t E\r\n\r\n\r\n"), b" C\r\nD E\r\n");
+    }
+
+    #[test]
+    fn test_canonicalize_body_simple() {
+        assert_eq!(body_simple(b"\r\n"), b"\r\n");
+        assert_eq!(body_simple(b"hey        \r\n"), b"hey        \r\n");
+        assert_eq!(
+            body_simple(b" C \r\nD \t E\r\n\r\n\r\n"),
+            b" C \r\nD \t E\r\n"
+        );
     }
 }

src/hash.rs

@@ -1,15 +1,20 @@
+use once_cell::sync::Lazy;
 use std::collections::HashMap;
 
 use base64::engine::general_purpose;
 use base64::Engine;
+use memchr::memmem::Finder;
+use sha1::Digest as _;
+use sha1::Sha1;
+use sha2::Sha256;
 use slog::debug;
 
 use crate::canonicalization::{
-    self, canonicalize_body_relaxed, canonicalize_body_simple, canonicalize_header_relaxed,
+    self, apply_body_relaxed, canonicalize_body_simple, canonicalize_header_relaxed,
     canonicalize_header_simple,
 };
 use crate::header::HEADER;
-use crate::{bytes, DKIMError, DKIMHeader};
+use crate::{DKIMError, DKIMHeader};
 
 #[derive(Debug, Clone)]
 pub enum HashAlgo {
@@ -18,25 +23,85 @@ pub enum HashAlgo {
     Ed25519Sha256,
 }
 
+pub(crate) struct LimitHasher {
+    pub limit: usize,
+    pub hashed: usize,
+    pub hasher: HashImpl,
+}
+
+impl LimitHasher {
+    pub fn hash(&mut self, bytes: &[u8]) {
+        let remain = self.limit - self.hashed;
+        let len = bytes.len().min(remain);
+        self.hasher.hash(&bytes[..len]);
+        self.hashed += len;
+    }
+
+    pub fn finalize(self) -> String {
+        self.hasher.finalize()
+    }
+
+    #[cfg(test)]
+    pub fn finalize_bytes(self) -> Vec<u8> {
+        self.hasher.finalize_bytes()
+    }
+}
+
+pub(crate) enum HashImpl {
+    Sha1(Sha1),
+    Sha256(Sha256),
+    #[cfg(test)]
+    Copy(Vec<u8>),
+}
+
+impl HashImpl {
+    pub fn from_algo(algo: HashAlgo) -> Self {
+        match algo {
+            HashAlgo::RsaSha1 => Self::Sha1(Sha1::new()),
+            HashAlgo::RsaSha256 | HashAlgo::Ed25519Sha256 => Self::Sha256(Sha256::new()),
+        }
+    }
+
+    #[cfg(test)]
+    pub fn copy_data() -> Self {
+        Self::Copy(vec![])
+    }
+
+    pub fn hash(&mut self, bytes: &[u8]) {
+        match self {
+            Self::Sha1(hasher) => hasher.update(bytes),
+            Self::Sha256(hasher) => hasher.update(bytes),
+            #[cfg(test)]
+            Self::Copy(data) => data.extend_from_slice(bytes),
+        }
+    }
+
+    pub fn finalize(self) -> String {
+        match self {
+            Self::Sha1(hasher) => general_purpose::STANDARD.encode(hasher.finalize()),
+            Self::Sha256(hasher) => general_purpose::STANDARD.encode(hasher.finalize()),
+            #[cfg(test)]
+            Self::Copy(data) => String::from_utf8_lossy(&data).into(),
+        }
+    }
+
+    pub fn finalize_bytes(self) -> Vec<u8> {
+        match self {
+            Self::Sha1(hasher) => hasher.finalize().to_vec(),
+            Self::Sha256(hasher) => hasher.finalize().to_vec(),
+            #[cfg(test)]
+            Self::Copy(data) => data,
+        }
+    }
+}
+
 /// Get the body part of an email
-fn get_body<'a>(email: &'a mailparse::ParsedMail<'a>) -> Result<&'a [u8], DKIMError> {
-    Ok(bytes::get_all_after(email.raw_bytes, b"\r\n\r\n"))
-}
-
-fn hash_sha1<T: AsRef<[u8]>>(data: T) -> Vec<u8> {
-    use sha1::{Digest, Sha1};
-
-    let mut hasher = Sha1::new();
-    hasher.update(data);
-    hasher.finalize().to_vec()
-}
-
-fn hash_sha256<T: AsRef<[u8]>>(data: T) -> Vec<u8> {
-    use sha2::{Digest, Sha256};
-
-    let mut hasher = Sha256::new();
-    hasher.update(data);
-    hasher.finalize().to_vec()
+fn get_body<'a>(email: &'a mailparse::ParsedMail<'a>) -> &'a [u8] {
+    static CRLFCRLF: Lazy<Finder> = Lazy::new(|| memchr::memmem::Finder::new("\r\n\r\n"));
+    CRLFCRLF
+        .find(email.raw_bytes)
+        .map(|idx| &email.raw_bytes[idx + 4..])
+        .unwrap_or(b"")
 }
 
 /// Returns the hash of message's body
@@ -47,26 +112,28 @@ pub(crate) fn compute_body_hash<'a>(
     hash_algo: HashAlgo,
     email: &'a mailparse::ParsedMail<'a>,
 ) -> Result<String, DKIMError> {
-    let body = get_body(email)?;
+    let body = get_body(email);
 
-    let mut canonicalized_body = if canonicalization_type == canonicalization::Type::Simple {
-        canonicalize_body_simple(&body)
-    } else {
-        canonicalize_body_relaxed(&body)
-    };
-    if let Some(length) = length {
-        let length = length
+    let limit = if let Some(length) = length {
+        length
             .parse::<usize>()
-            .map_err(|err| DKIMError::SignatureSyntaxError(format!("invalid length: {}", err)))?;
-        canonicalized_body.truncate(length);
+            .map_err(|err| DKIMError::SignatureSyntaxError(format!("invalid length: {}", err)))?
+    } else {
+        usize::MAX
     };
 
-    let hash = match hash_algo {
-        HashAlgo::RsaSha1 => hash_sha1(&canonicalized_body),
-        HashAlgo::RsaSha256 => hash_sha256(&canonicalized_body),
-        HashAlgo::Ed25519Sha256 => hash_sha256(&canonicalized_body),
+    let mut hasher = LimitHasher {
+        hasher: HashImpl::from_algo(hash_algo),
+        limit,
+        hashed: 0,
     };
-    Ok(general_purpose::STANDARD.encode(hash))
+
+    match canonicalization_type {
+        canonicalization::Type::Simple => canonicalize_body_simple(body, &mut hasher),
+        _ => apply_body_relaxed(body, &mut hasher),
+    };
+
+    Ok(hasher.finalize())
 }
 
 fn select_headers<'a>(
@@ -77,22 +144,19 @@ fn select_headers<'a>(
 
     let email_headers = &email.headers;
     let num_headers = email_headers.len();
-    let mut last_index: HashMap<String, usize> = HashMap::new();
+    let mut last_index: HashMap<&str, usize> = HashMap::new();
 
-    'outer: for name in dkim_header
-        .split(':')
-        .map(|h| h.trim().to_ascii_lowercase())
-    {
+    'outer: for name in dkim_header.split(':').map(|h| h.trim()) {
         let index = last_index.get(&name).unwrap_or(&num_headers);
-        for header in email_headers
+        for (header_index, header) in email_headers
             .iter()
             .enumerate()
             .rev()
             .skip(num_headers - index)
         {
-            if header.1.get_key_ref().eq_ignore_ascii_case(&name) {
-                signed_headers.push((header.1.get_key(), header.1.get_value_raw()));
-                last_index.insert(name, header.0);
+            if header.get_key_ref().eq_ignore_ascii_case(&name) {
+                signed_headers.push((header.get_key(), header.get_value_raw()));
+                last_index.insert(name, header_index);
                 continue 'outer;
             }
         }
@@ -112,15 +176,15 @@ pub(crate) fn compute_headers_hash<'a, 'b>(
     email: &'a mailparse::ParsedMail<'a>,
 ) -> Result<Vec<u8>, DKIMError> {
     let mut input = Vec::new();
+    let mut hasher = HashImpl::from_algo(hash_algo);
 
     // Add the headers defined in `h=` in the hash
     for (key, value) in select_headers(headers, email)? {
-        let canonicalized_value = if canonicalization_type == canonicalization::Type::Simple {
-            canonicalize_header_simple(&key, value)
+        if canonicalization_type == canonicalization::Type::Simple {
+            canonicalize_header_simple(&key, value, &mut input);
         } else {
-            canonicalize_header_relaxed(&key, value)
-        };
-        input.extend_from_slice(&canonicalized_value);
+            canonicalize_header_relaxed(&key, value, &mut input);
+        }
     }
 
     // Add the DKIM-Signature header in the hash. Remove the value of the
@@ -128,10 +192,11 @@ pub(crate) fn compute_headers_hash<'a, 'b>(
     {
         let sign = dkim_header.get_raw_tag("b").unwrap();
         let value = dkim_header.raw_bytes.replace(&sign, "");
-        let mut canonicalized_value = if canonicalization_type == canonicalization::Type::Simple {
-            canonicalize_header_simple(HEADER, value.as_bytes())
+        let mut canonicalized_value = vec![];
+        if canonicalization_type == canonicalization::Type::Simple {
+            canonicalize_header_simple(HEADER, value.as_bytes(), &mut canonicalized_value);
         } else {
-            canonicalize_header_relaxed(HEADER, value.as_bytes())
+            canonicalize_header_relaxed(HEADER, value.as_bytes(), &mut canonicalized_value);
         };
 
         // remove trailing "\r\n"
@@ -143,11 +208,8 @@ pub(crate) fn compute_headers_hash<'a, 'b>(
         debug!(logger, "headers to hash: {:?}", input);
     }
 
-    let hash = match hash_algo {
-        HashAlgo::RsaSha1 => hash_sha1(&input),
-        HashAlgo::RsaSha256 => hash_sha256(&input),
-        HashAlgo::Ed25519Sha256 => hash_sha256(&input),
-    };
+    hasher.hash(&input);
+    let hash = hasher.finalize_bytes();
     Ok(hash)
 }
 
@@ -411,7 +473,7 @@ Hello Alice
         let email =
             mailparse::parse_mail("Subject: A\r\n\r\nContent\n.hi\n.hello..".as_bytes()).unwrap();
         assert_eq!(
-            String::from_utf8_lossy(&get_body(&email).unwrap()),
+            String::from_utf8_lossy(get_body(&email)),
             "Content\n.hi\n.hello..".to_owned()
         );
     }