dkim: remove another extraneous allocation, tidy up canon interface

crates/dkim/benches/sign.rs

@@ -1,5 +1,5 @@
 use cfdkim::canonicalization::Type;
-use cfdkim::{DkimPrivateKey, SignerBuilder};
+use cfdkim::{DkimPrivateKey, ParsedEmail, SignerBuilder};
 use chrono::TimeZone;
 use criterion::{black_box, criterion_group, criterion_main, Criterion, SamplingMode, Throughput};
 use rsa::pkcs1::DecodeRsaPrivateKey;
@@ -48,10 +48,11 @@ suscipit mattis risus non placerat. Sed imperdiet, nisi et laoreet imperdiet, ur
 ante, non ultricies enim lorem cursus ante. Mauris euismod turpis eu tristique lobortis. Mauris
 aliquet eu tortor nec hendrerit. Aliquam in arcu ac erat venenatis pretium at sit amet magna. Lorem
 ipsum dolor sit a.
-        "#;
+        "#
+    .replace("\n", "\r\n");
 
     for canon in [Type::Simple, Type::Relaxed] {
-        let email = mailparse::parse_mail(email_text.as_bytes()).unwrap();
+        let email = ParsedEmail::parse_bytes(email_text.as_bytes()).unwrap();
 
         let private_key =
             rsa::RsaPrivateKey::read_pkcs1_pem_file(Path::new("./test/keys/2022.private")).unwrap();

crates/dkim/src/bytes.rs

@@ -1,62 +0,0 @@
-///! Various utility functions to operate on bytes
-pub(crate) use memchr::memmem::find;
-
-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_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 let Some(idx) = find(&result[i..], from) {
-        result.splice(idx + i..idx + i + from_len, to.iter().cloned());
-        i += idx + to_len;
-    }
-}
-
-#[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);
-        assert_eq!(find(&[97, 98, 99], &[97]), Some(0));
-        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()
-        );
-    }
-}

crates/dkim/src/canonicalization.rs

@@ -1,5 +1,4 @@
 //! 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;
@@ -17,6 +16,20 @@ impl Type {
             Self::Relaxed => "relaxed",
         }
     }
+
+    pub(crate) fn canon_body(&self, body: &[u8], hasher: &mut LimitHasher) {
+        match self {
+            Self::Simple => body_simple(body, hasher),
+            Self::Relaxed => body_relaxed(body, hasher),
+        }
+    }
+
+    pub(crate) fn canon_header_into(&self, key: &str, value: &[u8], out: &mut Vec<u8>) {
+        match self {
+            Self::Simple => canonicalize_header_simple(key, value, out),
+            Self::Relaxed => canonicalize_header_relaxed(key, value, out),
+        }
+    }
 }
 
 fn do_body_simple<'a>(mut body: &'a [u8]) -> &'a [u8] {
@@ -32,7 +45,7 @@ fn do_body_simple<'a>(mut body: &'a [u8]) -> &'a [u8] {
 }
 
 /// Canonicalize body using the simple canonicalization algorithm.
-pub(crate) fn canonicalize_body_simple(body: &[u8], hasher: &mut LimitHasher) {
+fn body_simple(body: &[u8], hasher: &mut LimitHasher) {
     let body = do_body_simple(body);
     hasher.hash(body);
 }
@@ -84,7 +97,7 @@ fn iter_lines(haystack: &[u8]) -> IterLines {
 
 /// 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) {
+fn body_relaxed(mut body: &[u8], hasher: &mut LimitHasher) {
     if body.is_empty() {
         return;
     }
@@ -96,14 +109,7 @@ pub(crate) fn apply_body_relaxed(mut body: &[u8], hasher: &mut LimitHasher) {
 
     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,
-            }
-        }
+        line = trim_ws_end(line);
 
         let mut prior = 0;
         // Reduce all sequences of WSP within a line to a single SP character.
@@ -131,7 +137,7 @@ pub(crate) fn apply_body_relaxed(mut body: &[u8], hasher: &mut LimitHasher) {
 }
 
 // https://datatracker.ietf.org/doc/html/rfc6376#section-3.4.1
-pub(crate) fn canonicalize_header_simple(key: &str, value: &[u8], out: &mut Vec<u8>) {
+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);
@@ -139,43 +145,55 @@ pub(crate) fn canonicalize_header_simple(key: &str, value: &[u8], out: &mut Vec<
 }
 
 // https://datatracker.ietf.org/doc/html/rfc6376#section-3.4.2
-pub(crate) fn canonicalize_header_relaxed(key: &str, value: &[u8], out: &mut Vec<u8>) {
+fn canonicalize_header_relaxed(key: &str, value: &[u8], out: &mut Vec<u8>) {
     let key = key.to_lowercase();
     let key = key.trim_end();
 
     out.extend_from_slice(key.as_bytes());
     out.extend_from_slice(b":");
-    out.extend_from_slice(&canonicalize_header_value_relaxed(value));
+
+    let value = trim_ws_start(trim_ws_end(value));
+    let mut space_run = false;
+    for &c in value {
+        match c {
+            b'\r' | b'\n' => {}
+            b' ' | b'\t' => {
+                if space_run {
+                    continue;
+                }
+                space_run = true;
+                out.push(b' ');
+            }
+            _ => {
+                space_run = false;
+                out.push(c);
+            }
+        }
+    }
+
     out.extend_from_slice(b"\r\n");
 }
 
-fn canonicalize_header_value_relaxed(value: &[u8]) -> Vec<u8> {
-    let mut value = value.to_vec();
-    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);
-    }
-    while value.starts_with(b" ") {
-        value.remove(0);
-    }
-    let mut previous = false;
-    value.retain(|c| {
-        if *c == b' ' {
-            if previous {
-                false
-            } else {
-                previous = true;
-                true
-            }
-        } else {
-            previous = false;
-            true
+fn trim_ws_start(mut line: &[u8]) -> &[u8] {
+    while let Some(c) = line.first() {
+        match c {
+            b' ' | b'\t' | b'\r' | b'\n' => line = &line[1..],
+            _ => break,
         }
-    });
+    }
+    line
+}
 
-    value
+fn trim_ws_end(mut line: &[u8]) -> &[u8] {
+    while let Some(c) = line.last() {
+        match c {
+            b' ' | b'\t' | b'\r' | b'\n' => {
+                line = &line[0..line.len() - 1];
+            }
+            _ => break,
+        }
+    }
+    line
 }
 
 #[cfg(test)]
@@ -207,7 +225,7 @@ mod tests {
             limit: usize::MAX,
             hashed: 0,
         };
-        apply_body_relaxed(data, &mut hasher);
+        super::body_relaxed(data, &mut hasher);
         hasher.finalize_bytes()
     }
 
@@ -217,7 +235,7 @@ mod tests {
             limit: usize::MAX,
             hashed: 0,
         };
-        canonicalize_body_simple(data, &mut hasher);
+        super::body_simple(data, &mut hasher);
         hasher.finalize_bytes()
     }
 

crates/dkim/src/hash.rs

@@ -1,9 +1,5 @@
-use crate::canonicalization::{
-    self, apply_body_relaxed, canonicalize_body_simple, canonicalize_header_relaxed,
-    canonicalize_header_simple,
-};
 use crate::header::HEADER;
-use crate::{DKIMError, DKIMHeader, ParsedEmail};
+use crate::{canonicalization, DKIMError, DKIMHeader, ParsedEmail};
 use base64::engine::general_purpose;
 use base64::Engine;
 use sha1::{Digest as _, Sha1};
@@ -123,10 +119,7 @@ pub(crate) fn compute_body_hash<'a>(
         hashed: 0,
     };
 
-    match canonicalization_type {
-        canonicalization::Type::Simple => canonicalize_body_simple(body, &mut hasher),
-        _ => apply_body_relaxed(body, &mut hasher),
-    };
+    canonicalization_type.canon_body(body, &mut hasher);
 
     Ok(hasher.finalize())
 }
@@ -174,11 +167,7 @@ pub(crate) fn compute_headers_hash<'a, 'b>(
 
     // Add the headers defined in `h=` in the hash
     for (key, value) in select_headers(headers, email)? {
-        if canonicalization_type == canonicalization::Type::Simple {
-            canonicalize_header_simple(&key, value, &mut input);
-        } else {
-            canonicalize_header_relaxed(&key, value, &mut input);
-        }
+        canonicalization_type.canon_header_into(&key, value, &mut input);
     }
 
     // Add the DKIM-Signature header in the hash. Remove the value of the
@@ -187,11 +176,7 @@ pub(crate) fn compute_headers_hash<'a, 'b>(
         let sign = dkim_header.get_required_raw_tag("b");
         let value = dkim_header.raw_bytes.replace(&sign, "");
         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(), &mut canonicalized_value);
-        };
+        canonicalization_type.canon_header_into(HEADER, value.as_bytes(), &mut canonicalized_value);
 
         // remove trailing "\r\n"
         canonicalized_value.truncate(canonicalized_value.len() - 2);

crates/dkim/src/lib.rs

@@ -13,7 +13,6 @@ use mailparse::MailHeaderMap;
 #[macro_use]
 extern crate quick_error;
 
-mod bytes;
 pub mod canonicalization;
 pub mod dns;
 mod errors;