mhost 0.11.3

// Copyright 2017-2021 Lukas Pustina <lukas@pustina.de>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use std::collections::{BTreeMap, HashSet};
use std::marker::PhantomData;

use crate::resolver::lookup::LookupResult;
use crate::resolver::{Error, Lookups};
use crate::RecordType;

use super::*;

#[derive(Debug)]
pub struct LookupsStats<'a> {
    pub responses: usize,
    pub nxdomains: usize,
    pub timeout_errors: usize,
    pub refuse_errors: usize,
    pub servfail_errors: usize,
    pub total_errors: usize,
    pub rr_type_counts: BTreeMap<RecordType, usize>,
    pub responding_servers: usize,
    pub response_time_summary: Summary<u128>,
    // This is used to please the borrow checker as we currently don't use a borrowed value with lifetime 'a
    phantom: PhantomData<&'a usize>,
}

#[derive(Debug)]
struct Counts {
    responses: usize,
    nxdomains: usize,
    timeout_errors: usize,
    refuse_errors: usize,
    servfail_errors: usize,
    total_errors: usize,
}

impl<'a> Statistics<'a> for Lookups {
    type StatsOut = LookupsStats<'a>;

    fn statistics(&'a self) -> Self::StatsOut {
        let counts = count_result_types(self);
        let rr_type_counts = count_rr_types(self);
        let responding_servers = count_responding_servers(self);
        let response_times: Vec<_> = self
            .iter()
            .filter_map(|x| x.result().response())
            .map(|x| x.response_time().as_millis())
            .collect();
        let response_time_summary = Summary::summary(response_times.as_slice());

        LookupsStats {
            responses: counts.responses,
            nxdomains: counts.nxdomains,
            timeout_errors: counts.timeout_errors,
            refuse_errors: counts.refuse_errors,
            servfail_errors: counts.servfail_errors,
            total_errors: counts.total_errors,
            rr_type_counts,
            responding_servers,
            response_time_summary,
            phantom: PhantomData,
        }
    }
}

fn count_rr_types(lookups: &Lookups) -> BTreeMap<RecordType, usize> {
    let mut type_counts = BTreeMap::new();

    for l in lookups.iter() {
        if let Some(response) = l.result().response() {
            for r in response.records() {
                let type_count = type_counts.entry(r.record_type()).or_insert(0);
                *type_count += 1;
            }
        }
    }

    type_counts
}

fn count_result_types(lookups: &Lookups) -> Counts {
    let mut responses: usize = 0;
    let mut nxdomains: usize = 0;
    let mut timeout_errors: usize = 0;
    let mut refuse_errors: usize = 0;
    let mut servfail_errors: usize = 0;
    let mut total_errors: usize = 0;

    for l in lookups.iter() {
        match l.result() {
            LookupResult::Response { .. } => responses += 1,
            LookupResult::NxDomain { .. } => nxdomains += 1,
            LookupResult::Error(Error::Timeout) => {
                timeout_errors += 1;
                total_errors += 1
            }
            LookupResult::Error(Error::QueryRefused) => {
                refuse_errors += 1;
                total_errors += 1
            }
            LookupResult::Error(Error::ServerFailure) => {
                servfail_errors += 1;
                total_errors += 1
            }
            LookupResult::Error { .. } => total_errors += 1,
        }
    }

    Counts {
        responses,
        nxdomains,
        timeout_errors,
        refuse_errors,
        servfail_errors,
        total_errors,
    }
}

fn count_responding_servers(lookups: &Lookups) -> usize {
    let server_set: HashSet<_> = lookups
        .iter()
        .filter(|x| x.result().is_response())
        .map(|x| x.name_server().to_string())
        .collect();

    server_set.len()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::nameserver::NameServerConfig;
    use crate::resolver::lookup::{Lookup, LookupResult, NxDomain, Response};
    use crate::resolver::{Error, Lookups, UniQuery};
    use crate::resources::rdata::{Name, MX};
    use crate::resources::{RData, Record};
    use std::net::Ipv4Addr;
    use std::sync::Arc;
    use std::time::Duration;

    fn make_ns(addr: &str) -> Arc<NameServerConfig> {
        Arc::new(NameServerConfig::from_str(addr).unwrap())
    }

    fn make_query() -> UniQuery {
        UniQuery::new("example.com.", RecordType::A).unwrap()
    }

    fn make_a_record() -> Record {
        Record::new_for_test(
            Name::from_utf8("example.com.").unwrap(),
            RecordType::A,
            300,
            RData::A(Ipv4Addr::new(1, 2, 3, 4)),
        )
    }

    fn make_mx_record() -> Record {
        Record::new_for_test(
            Name::from_utf8("example.com.").unwrap(),
            RecordType::MX,
            300,
            RData::MX(MX::new(10, Name::from_utf8("mail.example.com.").unwrap())),
        )
    }

    #[test]
    fn statistics_empty_lookups() {
        let lookups = Lookups::empty();
        let stats = lookups.statistics();

        assert_eq!(stats.responses, 0);
        assert_eq!(stats.nxdomains, 0);
        assert_eq!(stats.timeout_errors, 0);
        assert_eq!(stats.refuse_errors, 0);
        assert_eq!(stats.servfail_errors, 0);
        assert_eq!(stats.total_errors, 0);
        assert!(stats.rr_type_counts.is_empty());
        assert_eq!(stats.responding_servers, 0);
        assert!(stats.response_time_summary.min.is_none());
        assert!(stats.response_time_summary.max.is_none());
    }

    #[test]
    fn statistics_counts_responses_and_nxdomains() {
        let ns = make_ns("udp:8.8.8.8:53");
        let lookups = Lookups::new(vec![
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::Response(Response::new_for_test(vec![make_a_record()], Duration::from_millis(10))),
            ),
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::Response(Response::new_for_test(vec![make_a_record()], Duration::from_millis(20))),
            ),
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::NxDomain(NxDomain::new_for_test(Duration::from_millis(15))),
            ),
        ]);
        let stats = lookups.statistics();

        assert_eq!(stats.responses, 2);
        assert_eq!(stats.nxdomains, 1);
    }

    #[test]
    fn statistics_counts_error_types() {
        let ns = make_ns("udp:8.8.8.8:53");
        let lookups = Lookups::new(vec![
            Lookup::new_for_test(make_query(), ns.clone(), LookupResult::Error(Error::Timeout)),
            Lookup::new_for_test(make_query(), ns.clone(), LookupResult::Error(Error::QueryRefused)),
            Lookup::new_for_test(make_query(), ns.clone(), LookupResult::Error(Error::ServerFailure)),
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::Error(Error::ResolveError {
                    reason: "test".to_string(),
                }),
            ),
        ]);
        let stats = lookups.statistics();

        assert_eq!(stats.timeout_errors, 1);
        assert_eq!(stats.refuse_errors, 1);
        assert_eq!(stats.servfail_errors, 1);
        assert_eq!(stats.total_errors, 4);
    }

    #[test]
    fn statistics_counts_rr_types() {
        let ns = make_ns("udp:8.8.8.8:53");
        let lookups = Lookups::new(vec![
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::Response(Response::new_for_test(
                    vec![make_a_record(), make_a_record()],
                    Duration::from_millis(10),
                )),
            ),
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::Response(Response::new_for_test(
                    vec![make_mx_record()],
                    Duration::from_millis(20),
                )),
            ),
        ]);
        let stats = lookups.statistics();

        assert_eq!(stats.rr_type_counts[&RecordType::A], 2);
        assert_eq!(stats.rr_type_counts[&RecordType::MX], 1);
        assert_eq!(stats.rr_type_counts.len(), 2);
    }

    #[test]
    fn statistics_counts_responding_servers() {
        let ns1 = make_ns("udp:8.8.8.8:53");
        let ns2 = make_ns("udp:1.1.1.1:53");
        let lookups = Lookups::new(vec![
            Lookup::new_for_test(
                make_query(),
                ns1.clone(),
                LookupResult::Response(Response::new_for_test(vec![make_a_record()], Duration::from_millis(10))),
            ),
            Lookup::new_for_test(
                make_query(),
                ns2.clone(),
                LookupResult::Response(Response::new_for_test(vec![make_a_record()], Duration::from_millis(20))),
            ),
            // NxDomain should not count as a responding server
            Lookup::new_for_test(
                make_query(),
                make_ns("udp:9.9.9.9:53"),
                LookupResult::NxDomain(NxDomain::new_for_test(Duration::from_millis(15))),
            ),
        ]);
        let stats = lookups.statistics();

        assert_eq!(stats.responding_servers, 2);
    }

    #[test]
    fn statistics_response_time_summary() {
        let ns = make_ns("udp:8.8.8.8:53");
        let lookups = Lookups::new(vec![
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::Response(Response::new_for_test(vec![make_a_record()], Duration::from_millis(50))),
            ),
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::Response(Response::new_for_test(vec![make_a_record()], Duration::from_millis(10))),
            ),
            Lookup::new_for_test(
                make_query(),
                ns.clone(),
                LookupResult::Response(Response::new_for_test(
                    vec![make_a_record()],
                    Duration::from_millis(100),
                )),
            ),
        ]);
        let stats = lookups.statistics();

        assert_eq!(stats.response_time_summary.min, Some(10));
        assert_eq!(stats.response_time_summary.max, Some(100));
    }
}