aws_smithy_runtime/client/dns/
caching.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use std::{
7    io::{Error as IoError, ErrorKind as IoErrorKind},
8    net::IpAddr,
9    time::Duration,
10};
11
12use aws_smithy_runtime_api::client::dns::{DnsFuture, ResolveDns, ResolveDnsError};
13use hickory_resolver::{
14    config::{NameServerConfigGroup, ResolverConfig},
15    name_server::TokioConnectionProvider,
16    Resolver,
17};
18
19/// DNS resolver that uses [hickory_resolver] and caches DNS entries in memory.
20///
21/// This resolver requires a [tokio] runtime to function and isn't available for WASM targets.
22#[non_exhaustive]
23#[derive(Debug)]
24pub struct CachingDnsResolver {
25    resolver: Resolver<TokioConnectionProvider>,
26}
27
28impl Default for CachingDnsResolver {
29    /// Constructs a new Tokio based [ResolveDns] with the system configuration.
30    /// This uses `/etc/resolv.conf` on Unix OSes and registry settings on Windows.
31    fn default() -> Self {
32        Self {
33            resolver: Resolver::builder_tokio().expect("In tokio runtime").build(),
34        }
35    }
36}
37
38impl CachingDnsResolver {
39    /// Creates a new DNS resolver that caches IP addresses in memory.
40    pub fn builder() -> CachingDnsResolverBuilder {
41        CachingDnsResolverBuilder {
42            nameservers: None,
43            timeout: None,
44            attempts: None,
45            cache_size: None,
46            num_concurrent_reqs: None,
47        }
48    }
49
50    /// Flush the cache
51    pub fn clear_cache(&self) {
52        self.resolver.clear_cache();
53    }
54}
55
56impl ResolveDns for CachingDnsResolver {
57    fn resolve_dns<'a>(&'a self, name: &'a str) -> DnsFuture<'a> {
58        DnsFuture::new(async move {
59            let result = self.resolver.lookup_ip(name).await;
60
61            match result {
62                Ok(ips) => Ok(ips.into_iter().collect()),
63                Err(failure) => Err(ResolveDnsError::new(IoError::new(
64                    IoErrorKind::Other,
65                    failure,
66                ))),
67            }
68        })
69    }
70}
71
72pub struct CachingDnsResolverBuilder {
73    nameservers: Option<Nameservers>,
74    timeout: Option<Duration>,
75    attempts: Option<usize>,
76    cache_size: Option<usize>,
77    num_concurrent_reqs: Option<usize>,
78}
79
80struct Nameservers {
81    ips: Vec<IpAddr>,
82    port: u16,
83}
84
85impl CachingDnsResolverBuilder {
86    /// Configure upstream nameservers to use for resolution. Defaults to the system
87    /// configuration.
88    pub fn nameservers(mut self, ips: &[IpAddr], port: u16) -> Self {
89        self.nameservers = Some(Nameservers {
90            ips: ips.to_vec(),
91            port,
92        });
93        self
94    }
95
96    /// Specify the timeout for a request. Defaults to 5 seconds.
97    pub fn timeout(mut self, timeout: Duration) -> Self {
98        self.timeout = Some(timeout);
99        self
100    }
101
102    /// Number of retries after lookup failure before giving up. Defaults to 2.
103    pub fn attempts(mut self, attempts: usize) -> Self {
104        self.attempts = Some(attempts);
105        self
106    }
107
108    /// Cache size is in number of records (some records can be large). Defaults to 32.
109    /// Note: cache items are not evicted by new items, they are only evicted when their
110    /// TTL expires.
111    pub fn cache_size(mut self, cache_size: usize) -> Self {
112        self.cache_size = Some(cache_size);
113        self
114    }
115
116    /// Number of concurrent requests per query.
117    ///
118    /// Where more than one nameserver is configured, this configures the resolver
119    /// to send queries to a number of servers in parallel. Defaults to 2. Setting
120    /// to 0 or 1 will execute requests serially.
121    pub fn num_concurrent_reqs(mut self, num_concurrent_reqs: usize) -> Self {
122        self.num_concurrent_reqs = Some(num_concurrent_reqs);
123        self
124    }
125
126    pub fn build(self) -> CachingDnsResolver {
127        let mut builder = if let Some(nameservers) = self.nameservers {
128            let nameserver_config =
129                NameServerConfigGroup::from_ips_clear(&nameservers.ips, nameservers.port, true);
130            let resolver_config = ResolverConfig::from_parts(None, vec![], nameserver_config);
131
132            Resolver::builder_with_config(resolver_config, TokioConnectionProvider::default())
133        } else {
134            Resolver::builder_tokio().expect("Successfully read system config")
135        };
136
137        let opts = builder.options_mut();
138
139        if let Some(timeout) = self.timeout {
140            opts.timeout = timeout;
141        }
142
143        if let Some(attempts) = self.attempts {
144            opts.attempts = attempts;
145        }
146
147        if let Some(cache_size) = self.cache_size {
148            opts.cache_size = cache_size;
149        }
150
151        if let Some(num_concurrent_reqs) = self.num_concurrent_reqs {
152            opts.num_concurrent_reqs = num_concurrent_reqs;
153        }
154
155        CachingDnsResolver {
156            resolver: builder.build(),
157        }
158    }
159}