spftrace 0.3.0

// spftrace – utility for tracing SPF queries
// Copyright © 2022–2023 David Bürgin <dbuergin@gluet.ch>
//
// This program is free software: you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the Free Software
// Foundation, either version 3 of the License, or (at your option) any later
// version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
// details.
//
// You should have received a copy of the GNU General Public License along with
// this program. If not, see <https://www.gnu.org/licenses/>.

use crate::{config::Config, run};
use async_trait::async_trait;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use viaspf::{
    lookup::{Lookup, LookupError, LookupResult, Name},
    Sender,
};

#[derive(Default)]
struct MockLookup {
    lookup_a: Option<Box<dyn Fn(&Name) -> LookupResult<Vec<Ipv4Addr>> + Send + Sync + 'static>>,
    lookup_aaaa: Option<Box<dyn Fn(&Name) -> LookupResult<Vec<Ipv6Addr>> + Send + Sync + 'static>>,
    lookup_mx: Option<Box<dyn Fn(&Name) -> LookupResult<Vec<Name>> + Send + Sync + 'static>>,
    lookup_txt: Option<Box<dyn Fn(&Name) -> LookupResult<Vec<String>> + Send + Sync + 'static>>,
    lookup_ptr: Option<Box<dyn Fn(IpAddr) -> LookupResult<Vec<Name>> + Send + Sync + 'static>>,
}

impl MockLookup {
    fn new() -> Self {
        Default::default()
    }

    fn with_lookup_a(
        mut self,
        value: impl Fn(&Name) -> LookupResult<Vec<Ipv4Addr>> + Send + Sync + 'static,
    ) -> Self {
        self.lookup_a = Some(Box::new(value));
        self
    }

    fn with_lookup_aaaa(
        mut self,
        value: impl Fn(&Name) -> LookupResult<Vec<Ipv6Addr>> + Send + Sync + 'static,
    ) -> Self {
        self.lookup_aaaa = Some(Box::new(value));
        self
    }

    fn with_lookup_mx(
        mut self,
        value: impl Fn(&Name) -> LookupResult<Vec<Name>> + Send + Sync + 'static,
    ) -> Self {
        self.lookup_mx = Some(Box::new(value));
        self
    }

    fn with_lookup_txt(
        mut self,
        value: impl Fn(&Name) -> LookupResult<Vec<String>> + Send + Sync + 'static,
    ) -> Self {
        self.lookup_txt = Some(Box::new(value));
        self
    }

    fn with_lookup_ptr(
        mut self,
        value: impl Fn(IpAddr) -> LookupResult<Vec<Name>> + Send + Sync + 'static,
    ) -> Self {
        self.lookup_ptr = Some(Box::new(value));
        self
    }
}

#[async_trait]
impl Lookup for MockLookup {
    async fn lookup_a(&self, name: &Name) -> LookupResult<Vec<Ipv4Addr>> {
        self.lookup_a
            .as_ref()
            .map_or(Err(LookupError::NoRecords), |f| f(name))
    }

    async fn lookup_aaaa(&self, name: &Name) -> LookupResult<Vec<Ipv6Addr>> {
        self.lookup_aaaa
            .as_ref()
            .map_or(Err(LookupError::NoRecords), |f| f(name))
    }

    async fn lookup_mx(&self, name: &Name) -> LookupResult<Vec<Name>> {
        self.lookup_mx
            .as_ref()
            .map_or(Err(LookupError::NoRecords), |f| f(name))
    }

    async fn lookup_txt(&self, name: &Name) -> LookupResult<Vec<String>> {
        self.lookup_txt
            .as_ref()
            .map_or(Err(LookupError::NoRecords), |f| f(name))
    }

    async fn lookup_ptr(&self, ip: IpAddr) -> LookupResult<Vec<Name>> {
        self.lookup_ptr
            .as_ref()
            .map_or(Err(LookupError::NoRecords), |f| f(ip))
    }
}

#[tokio::test]
async fn readme_example() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.com." => Ok(vec!["v=spf1 mx include:spf.example.com ~all".into()]),
            "spf.example.com." => Ok(vec![
                "v=spf1 ip4:207.46.4.128/25 ip4:65.55.174.0/24 \
                ip6:2c0f:fb50:4000::/36 ip6:2001:4860:4000::/36 ~all".into(),
            ]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_mx(|name| match name.as_str() {
            "example.com." => Ok(vec![
                Name::new("mx1.example.com").unwrap(),
                Name::new("mx2.example.com").unwrap(),
            ]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_a(|name| match name.as_str() {
            "mx1.example.com." => Ok(vec![
                Ipv4Addr::new(216, 58, 192, 0),
                Ipv4Addr::new(65, 55, 52, 224),
                Ipv4Addr::new(207, 46, 116, 128),
            ]),
            "mx2.example.com." => Ok(vec![Ipv4Addr::new(65, 55, 238, 129)]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.com".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn exists_macros() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec![
                "v=spf1 this=works#?! exists:%{ir}.%{l1r+-}._spf.%{d} -all".into(),
            ]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_a(|name| match name.as_str() {
            "4.3.2.1.me._spf.example.org." => Ok(vec![Ipv4Addr::new(1, 2, 3, 4)]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me+lists@example.org".parse().unwrap();
    let ip = Ipv4Addr::new(1, 2, 3, 4).into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn invalid_redirect_target() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 redirect=%{i1}%{l}".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "-@example.org".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn multiple_modifiers() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 redirect=one.two redirect=three.four".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn multiple_spf_records() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 a -all".into(), "v=spf1 mx -all".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn neutral_default_result() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 a".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn recursive_spf_record() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 include:example.com".into()]),
            "example.com." => Ok(vec!["v=spf1 redirect=example.org".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn invalid_spf_record_syntax() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 include:example.com".into()]),
            "example.com." => Ok(vec!["v=spf1 redirect somewhere".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn target_name_not_found() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 include:example.com".into()]),
            "example.com." => Ok(vec!["v=spf1 include:i.dont.exist redirect=me.neither".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn ptr_mechanism() {
    let client_ip = Ipv6Addr::new(0x2001, 0x4, 0xca0, 0x103, 0, 0, 0, 0x9);

    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 ptr:mydomain.org -all".into()]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_ptr(move |ip| match ip {
            IpAddr::V6(addr) if addr == client_ip => {
                Ok(vec![
                    Name::new("someone.ch").unwrap(),
                    Name::new("mail.mydomain.org").unwrap(),
                    Name::new("other.mydomain.org").unwrap(),
                    Name::new("moreover.com").unwrap(),
                    Name::new("mailerr.mydomain.org").unwrap(),
                    Name::new("ptr6.someone.com").unwrap(),
                    Name::new("ptr7.someone.com").unwrap(),
                    Name::new("ptr8.someone.com").unwrap(),
                    Name::new("ptr9.someone.com").unwrap(),
                    Name::new("ptr10.someone.com").unwrap(),
                    Name::new("ptr11.someone.com").unwrap(),
                ])
            }
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_aaaa(move |name| match name.as_str() {
            "someone.ch." => Ok(vec![[0x2001, 0x4, 0xde, 0xabc, 0, 0, 0, 0].into()]),
            "mail.mydomain.org." => Ok(vec![client_ip]),
            "other.mydomain.org." => Ok(vec![
                [0x2001, 0x5, 0x6, 0xabc, 0, 0, 0, 2].into(),
                client_ip,
            ]),
            "mailerr.mydomain.org." => Err(LookupError::Dns(Some("DNS lookup failed".into()))),
            "ptr6.someone.com." => Ok(vec![[0x2001, 0x6, 0, 0, 0, 0, 0, 0].into()]),
            "ptr7.someone.com." => Ok(vec![[0x2001, 0x7, 0, 0, 0, 0, 0, 0].into()]),
            "ptr8.someone.com." => Ok(vec![[0x2001, 0x8, 0, 0, 0, 0, 0, 0].into()]),
            "ptr9.someone.com." => Ok(vec![[0x2001, 0x9, 0, 0, 0, 0, 0, 0].into()]),
            "ptr10.someone.com." => Ok(vec![[0x2001, 0xa, 0, 0, 0, 0, 0, 0].into()]),
            "ptr11.someone.com." => Ok(vec![[0x2001, 0xb, 0, 0, 0, 0, 0, 0].into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();

    run_test(config, sender, client_ip.into(), lookup).await;
}

#[tokio::test]
async fn ptr_lookup_error() {
    let client_ip = Ipv6Addr::new(0x2001, 0x4, 0xca0, 0x103, 0, 0, 0, 0x9);

    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 ptr -all".into()]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_ptr(move |ip| match ip {
            IpAddr::V6(addr) if addr == client_ip => Err(LookupError::Timeout),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();

    run_test(config, sender, client_ip.into(), lookup).await;
}

#[tokio::test]
async fn dual_cidr_length() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 +mx/24//64 -all".into()]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_mx(|name| match name.as_str() {
            "example.org." => Ok(vec![Name::new("mail.example.org").unwrap()]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_aaaa(|name| match name.as_str() {
            "mail.example.org." => Ok(vec![Ipv6Addr::new(1, 2, 3, 4, 0, 0, 0, 0)]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8).into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn exp_modifier() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 redirect=example.com".into()]),
            "example.com." => Ok(vec!["v=spf1 exp=exp._spf.%{d} -all +a:unused.com".into()]),
            "exp._spf.example.com." => Ok(vec!["%{r} says: %{l}: persona non grata.".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Config {
        hostname: Some("localhost".into()),
        ..Default::default()
    };
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn invalid_exp_modifier_syntax() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.com." => Ok(vec!["v=spf1 exp=exp._spf.example.com -all".into()]),
            "exp._spf.example.com." => Ok(vec!["%%%%%%%%%%%%%%%%%%%%%%%%%".into()]),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.com".parse().unwrap();
    let ip = Ipv4Addr::UNSPECIFIED.into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn p_macro_in_domain_spec() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 a:%{p}._validated.example.org".into()]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_a(|name| match name.as_str() {
            "mail.example.org."
            | "mail.other.org."
            | "mail.example.org._validated.example.org." => Ok(vec![Ipv4Addr::new(2, 3, 4, 5)]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_ptr(|ip| match ip {
            IpAddr::V4(addr) if addr.octets() == [2, 3, 4, 5] => {
                Ok(vec![
                    Name::new("mail.other.org").unwrap(),
                    Name::new("mail.example.org").unwrap(),
                ])
            }
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::new(2, 3, 4, 5).into();

    run_test(config, sender, ip, lookup).await;
}

#[tokio::test]
async fn dns_error() {
    let lookup = MockLookup::new()
        .with_lookup_txt(|name| match name.as_str() {
            "example.org." => Ok(vec!["v=spf1 mx -all".into()]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_mx(|name| match name.as_str() {
            "example.org." => Ok(vec![
                Name::new("mx1.example.org").unwrap(),
                Name::new("mx2.example.org").unwrap(),
            ]),
            _ => Err(LookupError::NoRecords),
        })
        .with_lookup_a(|name| match name.as_str() {
            "mx1.example.org." => Ok(vec![Ipv4Addr::new(1, 2, 3, 4)]),
            "mx2.example.org." => Err(LookupError::Dns(Some("what happened?".into()))),
            _ => Err(LookupError::NoRecords),
        });

    let config = Default::default();
    let sender = "me@example.org".parse().unwrap();
    let ip = Ipv4Addr::new(2, 3, 4, 5).into();

    run_test(config, sender, ip, lookup).await;
}

async fn run_test(config: Config, sender: Sender, ip: IpAddr, lookup: MockLookup) {
    let lookup = Box::new(lookup);
    let result = run::run_trace(config, sender, ip, Some(lookup)).await;
    assert_eq!(result, Ok(()));
}

#[tokio::test]
#[ignore = "depends on live DNS records"]
async fn live() {
    let config = Config {
        debug: true,
        time: true,
        ..Default::default()
    };
    let sender = "me@gluet.ch".parse().unwrap();
    let ip = Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8).into();

    let result = run::run_trace(config, sender, ip, None).await;

    assert_eq!(result, Ok(()));
}