From 2f63c2440960cf17c995cb33cb035231feff0dbe Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Tue, 28 Feb 2023 08:18:46 -0700 Subject: [PATCH] refactor: move MX resolution to its own module --- Cargo.lock | 2 + crates/kumod/Cargo.toml | 2 + crates/kumod/src/dest_site.rs | 165 +------------ crates/kumod/src/http_server/inject_v1.rs | 2 +- crates/kumod/src/logging.rs | 2 +- crates/kumod/src/main.rs | 1 + crates/kumod/src/mx.rs | 282 ++++++++++++++++++++++ crates/kumod/src/smtp_server.rs | 2 +- crates/lruttl/src/lib.rs | 15 ++ 9 files changed, 311 insertions(+), 162 deletions(-) create mode 100644 crates/kumod/src/mx.rs diff --git a/Cargo.lock b/Cargo.lock index 3170d0cc..9dccfc15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1240,6 +1240,7 @@ dependencies = [ "gethostname", "k9", "lazy_static", + "lruttl", "mail-auth", "mail-builder", "mail-parser", @@ -1269,6 +1270,7 @@ dependencies = [ "tokio-rustls", "tracing", "tracing-subscriber", + "trust-dns-resolver", "webpki-roots", "zstd 0.12.3+zstd.1.5.2", ] diff --git a/crates/kumod/Cargo.toml b/crates/kumod/Cargo.toml index 49a37765..1f3d9db7 100644 --- a/crates/kumod/Cargo.toml +++ b/crates/kumod/Cargo.toml @@ -20,6 +20,7 @@ config = {path="../config"} gcd = "2.3" gethostname = "0.4" lazy_static = "1.4" +lruttl = {path="../lruttl"} mail-auth = "0.3" mail-builder = "0.2" message = {path="../message"} @@ -47,6 +48,7 @@ tokio = {version="1.25", features=["full"]} tokio-rustls = "0.23" tracing = "0.1" tracing-subscriber = {version="0.3", features=["env-filter"]} +trust-dns-resolver = "0.22" webpki-roots = "0.22" zstd = "0.12" diff --git a/crates/kumod/src/dest_site.rs b/crates/kumod/src/dest_site.rs index 39078898..53762ae8 100644 --- a/crates/kumod/src/dest_site.rs +++ b/crates/kumod/src/dest_site.rs @@ -1,11 +1,11 @@ use crate::egress_source::EgressSource; use crate::lifecycle::{Activity, ShutdownSubcription}; use crate::logging::{log_disposition, RecordType}; +use crate::mx::{MailExchanger, ResolvedAddress}; use crate::queue::{Queue, QueueManager}; use crate::spool::SpoolManager; use anyhow::Context; use config::load_config; -use mail_auth::{IpLookupStrategy, Resolver}; use message::message::QueueNameComponents; use message::Message; use mlua::prelude::*; @@ -13,7 +13,6 @@ use prometheus::{IntCounter, IntGauge}; use rfc5321::{ClientError, EnhancedStatusCode, ForwardPath, Response, ReversePath, SmtpClient}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, VecDeque}; -use std::net::IpAddr; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex as StdMutex}; use std::time::{Duration, Instant}; @@ -23,7 +22,6 @@ use tokio::task::JoinHandle; lazy_static::lazy_static! { static ref MANAGER: Mutex = Mutex::new(SiteManager::new()); - static ref RESOLVER: Mutex = Mutex::new(Resolver::new_system_conf().unwrap()); } #[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq, Copy)] @@ -124,30 +122,6 @@ pub struct SiteManager { sites: HashMap, } -async fn resolve_mx(domain_name: &str) -> anyhow::Result> { - let resolver = RESOLVER.lock().await; - match resolver.mx_lookup(domain_name).await { - Ok(mxs) if mxs.is_empty() => Ok(vec![domain_name.to_string()]), - Ok(mxs) => { - let mut hosts = vec![]; - for mx in mxs.iter() { - let mut hosts_this_pref: Vec = - mx.exchanges.iter().map(|s| s.to_string()).collect(); - hosts_this_pref.sort(); - hosts.append(&mut hosts_this_pref); - } - Ok(hosts) - } - err @ Err(mail_auth::Error::DnsRecordNotFound(_)) => { - match resolver.exists(domain_name).await { - Ok(true) => Ok(vec![domain_name.to_string()]), - _ => anyhow::bail!("{:#}", err.unwrap_err()), - } - } - Err(err) => anyhow::bail!("MX lookup for {domain_name} failed: {err:#}"), - } -} - impl SiteManager { pub fn new() -> Self { Self { @@ -168,8 +142,8 @@ impl SiteManager { egress_source: &str, ) -> anyhow::Result { let components = QueueNameComponents::parse(queue_name); - let mx = Arc::new(resolve_mx(components.domain).await?.into_boxed_slice()); - let name = factor_names(&mx); + let mx = MailExchanger::resolve(components.domain).await?; + let name = &mx.site_name; let name = format!("{egress_source}->{name}"); let egress_source = EgressSource::resolve(egress_source)?; @@ -289,7 +263,7 @@ struct DeliveryMetrics { pub struct DestinationSite { name: String, - mx: Arc>, + mx: Arc, ready: Arc>>, notify: Arc, connections: Vec>, @@ -395,39 +369,6 @@ impl DestinationSite { } } -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ResolvedAddress { - pub name: String, - pub addr: IpAddr, -} -async fn resolve_addresses(mx: &Arc>) -> Vec { - let mut result = vec![]; - - for mx_host in mx.iter() { - match RESOLVER - .lock() - .await - .ip_lookup(mx_host, IpLookupStrategy::default(), 32) - .await - { - Err(err) => { - tracing::error!("failed to resolve {mx_host}: {err:#}"); - continue; - } - Ok(addresses) => { - for addr in addresses { - result.push(ResolvedAddress { - name: mx_host.to_string(), - addr, - }); - } - } - } - } - result.reverse(); - result -} - struct Dispatcher { name: String, ready: Arc>>, @@ -447,7 +388,7 @@ struct Dispatcher { impl Dispatcher { async fn run( name: &str, - mx: Arc>, + mx: Arc, ready: Arc>>, notify: Arc, site_config: DestSiteConfig, @@ -462,7 +403,7 @@ impl Dispatcher { let activity = Activity::get()?; - let addresses = resolve_addresses(&mx).await; + let addresses = mx.resolve_addresses().await; let mut dispatcher = Self { name: name.to_string(), ready, @@ -824,104 +765,10 @@ fn ideal_connection_count(queue_size: usize, connection_limit: usize) -> usize { goal.ceil() as usize } -/// Given a list of host names, produce a pseudo-regex style alternation list -/// of the different elements of the hostnames. -/// The goal is to produce a more compact representation of the name list -/// with the common components factored out. -fn factor_names>(names: &[S]) -> String { - let mut max_element_count = 0; - - let mut elements: Vec> = vec![]; - - let mut split_names = vec![]; - for name in names { - let name = name.as_ref(); - let mut fields: Vec<_> = name.split('.').map(|s| s.to_lowercase()).collect(); - fields.reverse(); - max_element_count = max_element_count.max(fields.len()); - split_names.push(fields); - } - - fn add_element<'a>(elements: &mut Vec>, field: &'a str, i: usize) { - match elements.get_mut(i) { - Some(ele) => { - if !ele.contains(&field) { - ele.push(field); - } - } - None => { - elements.push(vec![field]); - } - } - } - - for fields in &split_names { - for (i, field) in fields.iter().enumerate() { - add_element(&mut elements, field, i); - } - for i in fields.len()..max_element_count { - add_element(&mut elements, "?", i); - } - } - - let mut result = vec![]; - for mut ele in elements { - let has_q = ele.contains(&"?"); - ele.retain(|&e| e != "?"); - let mut item_text = if ele.len() == 1 { - ele[0].to_string() - } else { - format!("({})", ele.join("|")) - }; - if has_q { - item_text.push('?'); - } - result.push(item_text); - } - result.reverse(); - - result.join(".") -} - #[cfg(test)] mod test { use super::*; - #[test] - fn name_factoring() { - assert_eq!( - factor_names(&[ - "mta5.am0.yahoodns.net", - "mta6.am0.yahoodns.net", - "mta7.am0.yahoodns.net" - ]), - "(mta5|mta6|mta7).am0.yahoodns.net".to_string() - ); - - // Verify that the case is normalized to lowercase - assert_eq!( - factor_names(&[ - "mta5.AM0.yahoodns.net", - "mta6.am0.yAHOodns.net", - "mta7.am0.yahoodns.net" - ]), - "(mta5|mta6|mta7).am0.yahoodns.net".to_string() - ); - - // When the names have mismatched lengths, do we produce - // something reasonable? - assert_eq!( - factor_names(&[ - "gmail-smtp-in.l.google.com", - "alt1.gmail-smtp-in.l.google.com", - "alt2.gmail-smtp-in.l.google.com", - "alt3.gmail-smtp-in.l.google.com", - "alt4.gmail-smtp-in.l.google.com", - ]), - "(alt1|alt2|alt3|alt4)?.gmail-smtp-in.l.google.com".to_string() - ); - } - #[test] fn connection_limit() { let sizes = [ diff --git a/crates/kumod/src/http_server/inject_v1.rs b/crates/kumod/src/http_server/inject_v1.rs index 42e30a83..dc7d904a 100644 --- a/crates/kumod/src/http_server/inject_v1.rs +++ b/crates/kumod/src/http_server/inject_v1.rs @@ -1,7 +1,7 @@ -use crate::dest_site::ResolvedAddress; use crate::http_server::auth::AuthKind; use crate::http_server::AppError; use crate::logging::{log_disposition, RecordType}; +use crate::mx::ResolvedAddress; use crate::queue::QueueManager; use crate::spool::{SpoolHandle, SpoolManager}; use anyhow::Context; diff --git a/crates/kumod/src/logging.rs b/crates/kumod/src/logging.rs index 8c5810f1..79d16e48 100644 --- a/crates/kumod/src/logging.rs +++ b/crates/kumod/src/logging.rs @@ -1,4 +1,4 @@ -use crate::dest_site::ResolvedAddress; +use crate::mx::ResolvedAddress; use anyhow::Context; use async_channel::{Receiver, Sender}; use chrono::{DateTime, Utc}; diff --git a/crates/kumod/src/main.rs b/crates/kumod/src/main.rs index 69c9af93..e9b26017 100644 --- a/crates/kumod/src/main.rs +++ b/crates/kumod/src/main.rs @@ -15,6 +15,7 @@ mod lifecycle; mod logging; mod metrics_helper; mod mod_kumo; +mod mx; mod queue; mod runtime; mod smtp_server; diff --git a/crates/kumod/src/mx.rs b/crates/kumod/src/mx.rs new file mode 100644 index 00000000..1661e8e8 --- /dev/null +++ b/crates/kumod/src/mx.rs @@ -0,0 +1,282 @@ +use lruttl::LruCacheWithTtl; +use serde::{Deserialize, Serialize}; +use std::net::IpAddr; +use std::sync::{Arc, Mutex as StdMutex}; +use std::time::Instant; +use trust_dns_resolver::error::{ResolveErrorKind, ResolveResult}; +use trust_dns_resolver::TokioAsyncResolver; + +lazy_static::lazy_static! { + static ref RESOLVER: TokioAsyncResolver = TokioAsyncResolver::tokio_from_system_conf().unwrap(); + static ref MX_CACHE: StdMutex>> = StdMutex::new(LruCacheWithTtl::new(64 * 1024)); + static ref IPV4_CACHE: StdMutex>>> = StdMutex::new(LruCacheWithTtl::new(1024)); + static ref IPV6_CACHE: StdMutex>>> = StdMutex::new(LruCacheWithTtl::new(1024)); + static ref IP_CACHE: StdMutex>>> = StdMutex::new(LruCacheWithTtl::new(1024)); +} + +#[derive(Clone)] +pub struct MailExchanger { + pub domain_name: String, + pub hosts: Vec, + pub site_name: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ResolvedAddress { + pub name: String, + pub addr: IpAddr, +} + +impl MailExchanger { + pub async fn resolve(domain_name: &str) -> anyhow::Result> { + if let Some(mx) = MX_CACHE.lock().unwrap().get(domain_name) { + return Ok(mx); + } + + let (hosts, expires) = match lookup_mx_record(domain_name).await { + Ok((hosts, expires)) if hosts.is_empty() => (vec![domain_name.to_string()], expires), + Ok((hosts, expires)) => (hosts, expires), + Err(err) if matches!(err.kind(), ResolveErrorKind::NoRecordsFound { .. }) => { + match ip_lookup(domain_name).await { + Ok((_addr, expires)) => (vec![domain_name.to_string()], expires), + Err(err) => anyhow::bail!("{err:#}"), + } + } + Err(err) => anyhow::bail!("MX lookup for {domain_name} failed: {err:#}"), + }; + + let site_name = factor_names(&hosts); + let mx = Self { + hosts, + domain_name: domain_name.to_string(), + site_name, + }; + + let mx = Arc::new(mx); + MX_CACHE + .lock() + .unwrap() + .insert(domain_name.to_string(), mx.clone(), expires); + Ok(mx) + } + + pub async fn resolve_addresses(&self) -> Vec { + let mut result = vec![]; + + for mx_host in &self.hosts { + match ip_lookup(mx_host).await { + Err(err) => { + tracing::error!("failed to resolve {mx_host}: {err:#}"); + continue; + } + Ok((addresses, _expires)) => { + for addr in addresses.iter() { + result.push(ResolvedAddress { + name: mx_host.to_string(), + addr: *addr, + }); + } + } + } + } + result.reverse(); + result + } +} + +async fn lookup_mx_record(domain_name: &str) -> ResolveResult<(Vec, Instant)> { + let mx_lookup = RESOLVER.mx_lookup(domain_name).await?; + let mx_records = mx_lookup.as_lookup().records(); + + struct ByPreference { + hosts: Vec, + pref: u16, + } + + let mut records: Vec = Vec::with_capacity(mx_records.len()); + for mx_record in mx_records { + if let Some(mx) = mx_record.data().and_then(|r| r.as_mx()) { + let pref = mx.preference(); + let host = mx.exchange().to_lowercase().to_string(); + + if let Some(record) = records.iter_mut().find(|r| r.pref == pref) { + record.hosts.push(host); + } else { + records.push(ByPreference { + hosts: vec![host], + pref, + }); + } + } + } + + // Sort by preference + records.sort_unstable_by(|a, b| a.pref.cmp(&b.pref)); + + // Sort the hosts at each preference level to produce the + // overall ordered list of hosts for this site + let mut hosts = vec![]; + for mut mx in records { + mx.hosts.sort(); + hosts.append(&mut mx.hosts); + } + + Ok((hosts, mx_lookup.valid_until())) +} + +pub async fn ip_lookup(key: &str) -> ResolveResult<(Arc>, Instant)> { + if let Some(value) = IP_CACHE.lock().unwrap().get_with_expiry(key) { + return Ok(value); + } + let (addr, exp) = match ipv4_lookup(key).await { + Ok((v4, exp)) => (v4, exp), + Err(_) => ipv6_lookup(key).await?, + }; + + IP_CACHE + .lock() + .unwrap() + .insert(key.to_string(), addr.clone(), exp); + Ok((addr, exp)) +} + +pub async fn ipv4_lookup(key: &str) -> ResolveResult<(Arc>, Instant)> { + if let Some(value) = IPV4_CACHE.lock().unwrap().get_with_expiry(key) { + return Ok(value); + } + + let ipv4_lookup = RESOLVER.ipv4_lookup(key).await?; + let ips = ipv4_lookup + .as_lookup() + .record_iter() + .filter_map(|r| (IpAddr::from(*r.data()?.as_a()?).into())) + .collect::>(); + + let ips = Arc::new(ips); + let expires = ipv4_lookup.valid_until(); + IPV4_CACHE + .lock() + .unwrap() + .insert(key.to_string(), ips.clone(), expires); + Ok((ips, expires)) +} + +pub async fn ipv6_lookup(key: &str) -> ResolveResult<(Arc>, Instant)> { + if let Some(value) = IPV6_CACHE.lock().unwrap().get_with_expiry(key) { + return Ok(value); + } + + let ipv6_lookup = RESOLVER.ipv4_lookup(key).await?; + let ips = ipv6_lookup + .as_lookup() + .record_iter() + .filter_map(|r| (IpAddr::from(*r.data()?.as_a()?)).into()) + .collect::>(); + + let ips = Arc::new(ips); + let expires = ipv6_lookup.valid_until(); + IPV6_CACHE + .lock() + .unwrap() + .insert(key.to_string(), ips.clone(), expires); + Ok((ips, expires)) +} + +/// Given a list of host names, produce a pseudo-regex style alternation list +/// of the different elements of the hostnames. +/// The goal is to produce a more compact representation of the name list +/// with the common components factored out. +fn factor_names>(names: &[S]) -> String { + let mut max_element_count = 0; + + let mut elements: Vec> = vec![]; + + let mut split_names = vec![]; + for name in names { + let name = name.as_ref(); + let mut fields: Vec<_> = name.split('.').map(|s| s.to_lowercase()).collect(); + fields.reverse(); + max_element_count = max_element_count.max(fields.len()); + split_names.push(fields); + } + + fn add_element<'a>(elements: &mut Vec>, field: &'a str, i: usize) { + match elements.get_mut(i) { + Some(ele) => { + if !ele.contains(&field) { + ele.push(field); + } + } + None => { + elements.push(vec![field]); + } + } + } + + for fields in &split_names { + for (i, field) in fields.iter().enumerate() { + add_element(&mut elements, field, i); + } + for i in fields.len()..max_element_count { + add_element(&mut elements, "?", i); + } + } + + let mut result = vec![]; + for mut ele in elements { + let has_q = ele.contains(&"?"); + ele.retain(|&e| e != "?"); + let mut item_text = if ele.len() == 1 { + ele[0].to_string() + } else { + format!("({})", ele.join("|")) + }; + if has_q { + item_text.push('?'); + } + result.push(item_text); + } + result.reverse(); + + result.join(".") +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn name_factoring() { + assert_eq!( + factor_names(&[ + "mta5.am0.yahoodns.net", + "mta6.am0.yahoodns.net", + "mta7.am0.yahoodns.net" + ]), + "(mta5|mta6|mta7).am0.yahoodns.net".to_string() + ); + + // Verify that the case is normalized to lowercase + assert_eq!( + factor_names(&[ + "mta5.AM0.yahoodns.net", + "mta6.am0.yAHOodns.net", + "mta7.am0.yahoodns.net" + ]), + "(mta5|mta6|mta7).am0.yahoodns.net".to_string() + ); + + // When the names have mismatched lengths, do we produce + // something reasonable? + assert_eq!( + factor_names(&[ + "gmail-smtp-in.l.google.com", + "alt1.gmail-smtp-in.l.google.com", + "alt2.gmail-smtp-in.l.google.com", + "alt3.gmail-smtp-in.l.google.com", + "alt4.gmail-smtp-in.l.google.com", + ]), + "(alt1|alt2|alt3|alt4)?.gmail-smtp-in.l.google.com".to_string() + ); + } +} diff --git a/crates/kumod/src/smtp_server.rs b/crates/kumod/src/smtp_server.rs index ab379f89..14d72cb8 100644 --- a/crates/kumod/src/smtp_server.rs +++ b/crates/kumod/src/smtp_server.rs @@ -1,6 +1,6 @@ -use crate::dest_site::ResolvedAddress; use crate::lifecycle::{Activity, ShutdownSubcription}; use crate::logging::{log_disposition, RecordType}; +use crate::mx::ResolvedAddress; use crate::queue::QueueManager; use crate::spool::{SpoolHandle, SpoolManager}; use anyhow::{anyhow, Context}; diff --git a/crates/lruttl/src/lib.rs b/crates/lruttl/src/lib.rs index 6560dd4d..c6d1023a 100644 --- a/crates/lruttl/src/lib.rs +++ b/crates/lruttl/src/lib.rs @@ -32,6 +32,21 @@ impl LruCacheWithTtl { } } + pub fn get_with_expiry(&self, name: &Q) -> Option<(V, Instant)> + where + K: Borrow, + Q: Hash + Eq, + { + let mut cache = self.cache.lock(); + let entry = cache.get_mut(name)?; + if Instant::now() < entry.expiration { + Some((entry.item.clone(), entry.expiration)) + } else { + cache.remove(name); + None + } + } + pub fn get(&self, name: &Q) -> Option where K: Borrow,