refactor: move MX resolution to its own module

This commit is contained in:
Wez Furlong
2023-02-28 08:18:46 -07:00
parent 93ff9e1991
commit 2f63c24409
9 changed files with 311 additions and 162 deletions
Generated
+2
View File
@@ -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",
]
+2
View File
@@ -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"
+6 -159
View File
@@ -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<SiteManager> = Mutex::new(SiteManager::new());
static ref RESOLVER: Mutex<Resolver> = Mutex::new(Resolver::new_system_conf().unwrap());
}
#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq, Copy)]
@@ -124,30 +122,6 @@ pub struct SiteManager {
sites: HashMap<String, SiteHandle>,
}
async fn resolve_mx(domain_name: &str) -> anyhow::Result<Vec<String>> {
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<String> =
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<SiteHandle> {
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<Box<[String]>>,
mx: Arc<MailExchanger>,
ready: Arc<StdMutex<VecDeque<Message>>>,
notify: Arc<Notify>,
connections: Vec<JoinHandle<()>>,
@@ -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<Box<[String]>>) -> Vec<ResolvedAddress> {
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<StdMutex<VecDeque<Message>>>,
@@ -447,7 +388,7 @@ struct Dispatcher {
impl Dispatcher {
async fn run(
name: &str,
mx: Arc<Box<[String]>>,
mx: Arc<MailExchanger>,
ready: Arc<StdMutex<VecDeque<Message>>>,
notify: Arc<Notify>,
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<S: AsRef<str>>(names: &[S]) -> String {
let mut max_element_count = 0;
let mut elements: Vec<Vec<&str>> = 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<Vec<&'a str>>, 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 = [
+1 -1
View File
@@ -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;
+1 -1
View File
@@ -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};
+1
View File
@@ -15,6 +15,7 @@ mod lifecycle;
mod logging;
mod metrics_helper;
mod mod_kumo;
mod mx;
mod queue;
mod runtime;
mod smtp_server;
+282
View File
@@ -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<LruCacheWithTtl<String, Arc<MailExchanger>>> = StdMutex::new(LruCacheWithTtl::new(64 * 1024));
static ref IPV4_CACHE: StdMutex<LruCacheWithTtl<String, Arc<Vec<IpAddr>>>> = StdMutex::new(LruCacheWithTtl::new(1024));
static ref IPV6_CACHE: StdMutex<LruCacheWithTtl<String, Arc<Vec<IpAddr>>>> = StdMutex::new(LruCacheWithTtl::new(1024));
static ref IP_CACHE: StdMutex<LruCacheWithTtl<String, Arc<Vec<IpAddr>>>> = StdMutex::new(LruCacheWithTtl::new(1024));
}
#[derive(Clone)]
pub struct MailExchanger {
pub domain_name: String,
pub hosts: Vec<String>,
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<Arc<Self>> {
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<ResolvedAddress> {
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<String>, Instant)> {
let mx_lookup = RESOLVER.mx_lookup(domain_name).await?;
let mx_records = mx_lookup.as_lookup().records();
struct ByPreference {
hosts: Vec<String>,
pref: u16,
}
let mut records: Vec<ByPreference> = 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<Vec<IpAddr>>, 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<Vec<IpAddr>>, 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::<Vec<_>>();
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<Vec<IpAddr>>, 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::<Vec<_>>();
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<S: AsRef<str>>(names: &[S]) -> String {
let mut max_element_count = 0;
let mut elements: Vec<Vec<&str>> = 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<Vec<&'a str>>, 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()
);
}
}
+1 -1
View File
@@ -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};
+15
View File
@@ -32,6 +32,21 @@ impl<K: Hash + Eq, V: Clone> LruCacheWithTtl<K, V> {
}
}
pub fn get_with_expiry<Q: ?Sized>(&self, name: &Q) -> Option<(V, Instant)>
where
K: Borrow<Q>,
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<Q: ?Sized>(&self, name: &Q) -> Option<V>
where
K: Borrow<Q>,