spftrace 0.3.0

Utility for tracing SPF queries
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

// 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 async_trait::async_trait;
use std::{
    error::Error,
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    time::Duration,
};
use tokio::time::{self, Instant};
use trust_dns_resolver::{
    config::{ResolverConfig, ResolverOpts},
    system_conf, TokioAsyncResolver,
};
use viaspf::lookup::{Lookup, LookupError, LookupResult, Name};

/// Single-use, timeout Trust-DNS resolver.
///
/// We have disabled viaspf’s `tokio-timeout` feature and implemented timeout
/// per DNS query here, in order to only generate complete, not truncated
/// traces.
pub struct Resolver {
    pub inner: TokioAsyncResolver,
    pub deadline: Instant,
}

impl Resolver {
    pub fn new(timeout: Duration) -> Self {
        let (config, opts) = Default::default();

        Self::new_internal(timeout, config, opts)
    }

    pub fn new_with_system_config(timeout: Duration) -> Result<Self, Box<dyn Error>> {
        let (config, opts) = system_conf::read_system_conf()?;

        Ok(Self::new_internal(timeout, config, opts))
    }

    fn new_internal(timeout: Duration, config: ResolverConfig, mut opts: ResolverOpts) -> Self {
        opts.timeout = timeout;

        let inner = TokioAsyncResolver::tokio(config, opts);
        let deadline = Instant::now() + timeout;

        Self { inner, deadline }
    }
}

#[async_trait]
impl Lookup for Resolver {
    async fn lookup_a(&self, name: &Name) -> LookupResult<Vec<Ipv4Addr>> {
        time::timeout_at(self.deadline, self.inner.lookup_a(name))
            .await
            .map_err(|_| LookupError::Timeout)?
    }

    async fn lookup_aaaa(&self, name: &Name) -> LookupResult<Vec<Ipv6Addr>> {
        time::timeout_at(self.deadline, self.inner.lookup_aaaa(name))
            .await
            .map_err(|_| LookupError::Timeout)?
    }

    async fn lookup_mx(&self, name: &Name) -> LookupResult<Vec<Name>> {
        time::timeout_at(self.deadline, self.inner.lookup_mx(name))
            .await
            .map_err(|_| LookupError::Timeout)?
    }

    async fn lookup_txt(&self, name: &Name) -> LookupResult<Vec<String>> {
        time::timeout_at(self.deadline, self.inner.lookup_txt(name))
            .await
            .map_err(|_| LookupError::Timeout)?
    }

    async fn lookup_ptr(&self, ip: IpAddr) -> LookupResult<Vec<Name>> {
        time::timeout_at(self.deadline, self.inner.lookup_ptr(ip))
            .await
            .map_err(|_| LookupError::Timeout)?
    }
}