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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

// 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 serde::Serialize;
use trust_dns_resolver::IntoName;
use trust_dns_resolver::Name;

use crate::resolver::{Error, ResolverResult};
use crate::RecordType;

/// UniQuery
///
/// Name's labels are all Rc, so clone is cheap
#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize)]
pub struct UniQuery {
    pub(crate) name: Name,
    pub(crate) record_type: RecordType,
}

impl UniQuery {
    pub fn new<N: IntoName>(name: N, record_type: RecordType) -> ResolverResult<UniQuery> {
        let name = name.into_name().map_err(Error::from)?;

        Ok(UniQuery { name, record_type })
    }

    pub fn name(&self) -> &Name {
        &self.name
    }

    pub fn record_type(&self) -> RecordType {
        self.record_type
    }
}

impl From<UniQuery> for MultiQuery {
    fn from(query: UniQuery) -> MultiQuery {
        MultiQuery {
            names: vec![query.name],
            record_types: vec![query.record_type],
        }
    }
}

/// MultiQuery allows to lookup multiple names for multiple record types
///
/// It can be easily constructed from a simple `UniQuery`
///
/// # Example
/// ```
/// # use mhost::resolver::{UniQuery, MultiQuery};
/// # use mhost::RecordType;
/// let query = UniQuery::new("www.example.com", RecordType::A).unwrap();
/// let multi_query: MultiQuery = query.into();
/// ```
#[derive(Debug, Clone)]
pub struct MultiQuery {
    pub(crate) names: Vec<Name>,
    pub(crate) record_types: Vec<RecordType>,
}

impl MultiQuery {
    pub fn new<N: IntoName, S: IntoIterator<Item = N>, T: IntoIterator<Item = RecordType>>(
        names: S,
        record_types: T,
    ) -> ResolverResult<MultiQuery> {
        let names: Vec<_> = names
            .into_iter()
            .map(|name| name.into_name().map_err(Error::from))
            .collect();
        let names: ResolverResult<Vec<_>> = names.into_iter().collect();
        let names = names?;
        let record_types = record_types.into_iter().collect();

        Ok(MultiQuery { names, record_types })
    }

    /// Lookup a single name for a single records type
    pub fn single<N: IntoName>(name: N, record_type: RecordType) -> ResolverResult<MultiQuery> {
        MultiQuery::new(vec![name], vec![record_type])
    }

    /// Lookup a multiple names for a single records type
    pub fn multi_name<N: IntoName, S: IntoIterator<Item = N>>(
        names: S,
        record_type: RecordType,
    ) -> ResolverResult<MultiQuery> {
        MultiQuery::new(names, vec![record_type])
    }

    /// Lookup a single name for a multiple records types
    pub fn multi_record<N: IntoName, T: IntoIterator<Item = RecordType>>(
        name: N,
        record_types: T,
    ) -> ResolverResult<MultiQuery> {
        MultiQuery::new(vec![name], record_types)
    }

    /// Converts this `MultiQuery` into individual `UniQuery`s
    pub fn into_uni_queries(self) -> Vec<UniQuery> {
        let mut queries = Vec::new();
        for name in self.names.iter() {
            for record_type in self.record_types.iter() {
                queries.push(UniQuery {
                    name: name.clone(),
                    record_type: *record_type,
                });
            }
        }

        queries
    }

    /// Returns number of names of this `MultiQuery`
    pub fn num_names(&self) -> usize {
        self.names.len()
    }

    /// Returns number of record types of this `MultiQuery`
    pub fn num_record_types(&self) -> usize {
        self.record_types.len()
    }
}