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

use std::net::IpAddr;
use std::sync::Arc;
use std::collections::HashMap;

use futures::{BoxFuture, Future, failed};
use futures::stream::{BoxStream, Stream};

use stream_once::StreamOnce;
use {Name, Address, Error, Resolver, MemResolver, parse_name};


/// A builder/fluent interface to create a `Router`
pub struct RouterBuilder(Inner);

/// A helper which allows to resolve different names by different means
///
/// Create it using RouterBuilder
///
/// Note: router resolves any name only once. I.e. if name matches a suffix
/// it doesn't get resolved using a fallback resolver even if original resolver
/// returned error.
#[derive(Clone)]
pub struct Router(Arc<Inner>);


struct Inner {
    names: MemResolver,
    suffixes: HashMap<String, Box<Resolver>>,
    fallback: Option<Box<Resolver>>,
}

impl RouterBuilder {
    /// Start creating DNS router object
    pub fn new() -> RouterBuilder {
        RouterBuilder(Inner {
            names: MemResolver::new(),
            suffixes: HashMap::new(),
            fallback: None,
        })
    }
    /// Add a name that is resolved to a single IP (just like in `MemResolver`)
    pub fn add_ip(&mut self, name: Name, ip: IpAddr) -> &mut Self {
        self.0.names.add_host(name, ip);
        self
    }
    /// Add a suffix which will be resolved using specified resolver object
    ///
    /// This is useful for example to resolve all `*.consul` addresses against
    /// consul.
    ///
    /// Suffixes are always matched after dot in the name. Suffixes passed to
    /// to this function must *not* contain initial dot.
    ///
    /// If overlapping suffixes are specified, longest matching suffix wins.
    pub fn add_suffix<S, R>(&mut self, suffix: S, resolver: R) -> &mut Self
        where S: Into<String>, R: Resolver + 'static,
    {
        self.0.suffixes.insert(suffix.into(), Box::new(resolver));
        self
    }
    /// Add a default resolver
    ///
    /// Default resolver works as fallback when neither specific names nor
    /// suffixes matches.
    ///
    /// Note: when suffix matches but returns non existent domain name
    /// the default resolver is *not* called.
    pub fn add_default<R>(&mut self, resolver: R) -> &mut Self
        where R: Resolver + 'static,
    {
        self.0.fallback = Some(Box::new(resolver));
        self
    }
    /// Build real router object which can be used to resolve names
    pub fn into_resolver(self) -> Router {
        Router(Arc::new(self.0))
    }
}

impl Resolver for Router {
    fn resolve(&self, name: Name) -> BoxFuture<Address, Error> {
        if let Some((host, _)) = parse_name(name) {
            if self.0.names.contains_name(host) {
                return self.0.names.resolve(name);
            }
            if let Some(resolver) = self.0.suffixes.get(host) {
                return resolver.resolve(name);
            } else {
                for (idx, _) in host.match_indices('.') {
                    let suffix = &host[idx+1..];
                    if let Some(resolver) = self.0.suffixes.get(suffix) {
                        return resolver.resolve(name);
                    }
                }
                if let Some(ref resolver) = self.0.fallback {
                    return resolver.resolve(name);
                }
            }
        }
        failed(Error::NameNotFound).boxed()
    }
    fn subscribe(&self, name: Name) -> BoxStream<Address, Error> {
        if let Some((host, _)) = parse_name(name) {
            if self.0.names.contains_name(host) {
                return self.0.names.subscribe(name);
            }
            if let Some(resolver) = self.0.suffixes.get(host) {
                return resolver.subscribe(name);
            } else {
                for (idx, _) in host.match_indices('.') {
                    let suffix = &host[idx+1..];
                    if let Some(resolver) = self.0.suffixes.get(suffix) {
                        return resolver.subscribe(name);
                    }
                }
                if let Some(ref resolver) = self.0.fallback {
                    return resolver.subscribe(name);
                }
            }
        }
        StreamOnce::new(failed(Error::NameNotFound)).boxed()
    }
}