diff --git a/Cargo.lock b/Cargo.lock index 6c757e08..44284a72 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -133,6 +133,7 @@ dependencies = [ "mailparse", "memchr", "nom", + "once_cell", "quick-error 2.0.1", "regex", "rsa", diff --git a/Cargo.toml b/Cargo.toml index 83d913a7..4ee1bab4 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/bytes.rs b/src/bytes.rs index b00a65a5..4c0fc4f1 100644 --- a/src/bytes.rs +++ b/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 { - let mut result = source.to_vec(); +pub(crate) fn replace_within_vec(result: &mut Vec, 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 { + 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() + ); } } diff --git a/src/canonicalization.rs b/src/canonicalization.rs index 18fe3643..3f7ce3fc 100644 --- a/src/canonicalization.rs +++ b/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 { +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 { - 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 { + 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 = 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 { - let mut out = Vec::new(); +pub(crate) fn canonicalize_header_simple(key: &str, value: &[u8], out: &mut Vec) { 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 { +pub(crate) fn canonicalize_header_relaxed(key: &str, value: &[u8], out: &mut Vec) { 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 { 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 { mod tests { use super::*; + fn header_relaxed(key: &str, value: &[u8]) -> Vec { + 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 { + 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 { + 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" + ); } } diff --git a/src/hash.rs b/src/hash.rs index 2ad4dce0..17348b08 100644 --- a/src/hash.rs +++ b/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 { + self.hasher.finalize_bytes() + } +} + +pub(crate) enum HashImpl { + Sha1(Sha1), + Sha256(Sha256), + #[cfg(test)] + Copy(Vec), +} + +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 { + 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>(data: T) -> Vec { - use sha1::{Digest, Sha1}; - - let mut hasher = Sha1::new(); - hasher.update(data); - hasher.finalize().to_vec() -} - -fn hash_sha256>(data: T) -> Vec { - 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 = 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 { - 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::() - .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 = 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, 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() ); }