mtop_client/
discovery.rs

1use crate::core::MtopError;
2use crate::dns::{DnsClient, Message, MessageId, Name, RecordClass, RecordData, RecordType};
3use rustls_pki_types::ServerName;
4use std::cmp::Ordering;
5use std::collections::HashSet;
6use std::fmt;
7use std::net::{IpAddr, SocketAddr};
8use std::path::PathBuf;
9
10const DNS_A_PREFIX: &str = "dns+";
11const DNS_SRV_PREFIX: &str = "dnssrv+";
12const UNIX_SOCKET_PREFIX: &str = "/";
13
14/// Unique ID and address for a server in a Memcached cluster for indexing responses
15/// or errors and establishing connections.
16#[derive(Debug, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)]
17pub enum ServerID {
18    Name(String),
19    Socket(SocketAddr),
20    Path(PathBuf),
21}
22
23impl ServerID {
24    fn from_host_port<S>(host: S, port: u16) -> Self
25    where
26        S: AsRef<str>,
27    {
28        let host = host.as_ref();
29        if let Ok(ip) = host.parse::<IpAddr>() {
30            Self::Socket(SocketAddr::new(ip, port))
31        } else {
32            Self::Name(format!("{}:{}", host, port))
33        }
34    }
35}
36
37impl From<SocketAddr> for ServerID {
38    fn from(value: SocketAddr) -> Self {
39        Self::Socket(value)
40    }
41}
42
43impl From<(&str, u16)> for ServerID {
44    fn from(value: (&str, u16)) -> Self {
45        Self::from_host_port(value.0, value.1)
46    }
47}
48
49impl From<(String, u16)> for ServerID {
50    fn from(value: (String, u16)) -> Self {
51        Self::from_host_port(value.0, value.1)
52    }
53}
54
55impl From<PathBuf> for ServerID {
56    fn from(value: PathBuf) -> Self {
57        Self::Path(value)
58    }
59}
60
61impl fmt::Display for ServerID {
62    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
63        match self {
64            ServerID::Name(n) => n.fmt(f),
65            ServerID::Socket(s) => s.fmt(f),
66            ServerID::Path(p) => fmt::Debug::fmt(p, f),
67        }
68    }
69}
70
71/// An individual server that is part of a Memcached cluster.
72#[derive(Debug, Clone, Eq, PartialEq, Hash)]
73pub struct Server {
74    id: ServerID,
75    name: Option<ServerName<'static>>,
76}
77
78impl Server {
79    pub fn new(id: ServerID, name: ServerName<'static>) -> Self {
80        Self { id, name: Some(name) }
81    }
82
83    pub fn without_name(id: ServerID) -> Self {
84        Self { id, name: None }
85    }
86
87    pub fn id(&self) -> &ServerID {
88        &self.id
89    }
90
91    pub fn server_name(&self) -> &Option<ServerName<'static>> {
92        &self.name
93    }
94}
95
96impl PartialOrd for Server {
97    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
98        Some(self.cmp(other))
99    }
100}
101
102impl Ord for Server {
103    fn cmp(&self, other: &Self) -> Ordering {
104        self.id.cmp(&other.id)
105    }
106}
107
108impl fmt::Display for Server {
109    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
110        self.id.fmt(f)
111    }
112}
113
114/// Service discovery implementation for finding Memcached servers using DNS.
115///
116/// Different types of DNS records and different behaviors are used based on the
117/// presence of specific prefixes for hostnames. See `resolve_by_proto` for details.
118pub struct Discovery {
119    client: Box<dyn DnsClient + Send + Sync>,
120}
121
122impl Discovery {
123    pub fn new<C>(client: C) -> Self
124    where
125        C: DnsClient + Send + Sync + 'static,
126    {
127        Self {
128            client: Box::new(client),
129        }
130    }
131
132    /// Resolve a hostname to one or multiple Memcached servers based on DNS records
133    /// and/or the presence of certain prefixes on the hostnames.
134    ///
135    /// * `dns+` will resolve a hostname into multiple A and AAAA records and use the
136    ///   IP addresses from the records as Memcached servers.
137    /// * `dnssrv+` will resolve a hostname into multiple SRV records and use the
138    ///   unresolved targets from the SRV records as Memcached servers. Resolution of
139    ///   the targets to IP addresses will happen at connection time using the system
140    ///   resolver.
141    /// * `/` will resolve a hostname into a UNIX socket path and will use this path
142    ///   as a local Memcached server on a UNIX socket.
143    /// * No prefix with an IPv4 or IPv6 address will use the IP address as a Memcached
144    ///   server.
145    /// * No prefix with a non-IP address will use the host as a Memcached server.
146    ///   Resolution of the host to an IP address will happen at connection time using the
147    ///   system resolver.
148    pub async fn resolve_by_proto(&self, name: &str) -> Result<Vec<Server>, MtopError> {
149        if name.starts_with(DNS_A_PREFIX) {
150            Ok(self.resolve_a_aaaa(name.trim_start_matches(DNS_A_PREFIX)).await?)
151        } else if name.starts_with(DNS_SRV_PREFIX) {
152            Ok(self.resolve_srv(name.trim_start_matches(DNS_SRV_PREFIX)).await?)
153        } else if name.starts_with(UNIX_SOCKET_PREFIX) {
154            Ok(Self::resolve_unix_addr(name))
155        } else if let Ok(addr) = name.parse::<SocketAddr>() {
156            Ok(Self::resolve_socket_addr(name, addr)?)
157        } else {
158            Ok(Self::resolve_bare_host(name)?)
159        }
160    }
161
162    async fn resolve_srv(&self, name: &str) -> Result<Vec<Server>, MtopError> {
163        let (host, port) = Self::host_and_port(name)?;
164        let server_name = Self::server_name(host)?;
165        let name = host.parse()?;
166        let id = MessageId::random();
167
168        let res = self.client.resolve(id, name, RecordType::SRV, RecordClass::INET).await?;
169        Ok(Self::servers_from_answers(port, &server_name, &res))
170    }
171
172    async fn resolve_a_aaaa(&self, name: &str) -> Result<Vec<Server>, MtopError> {
173        let (host, port) = Self::host_and_port(name)?;
174        let server_name = Self::server_name(host)?;
175        let name: Name = host.parse()?;
176        let id = MessageId::random();
177
178        let res = self.client.resolve(id, name.clone(), RecordType::A, RecordClass::INET).await?;
179        let mut out = Self::servers_from_answers(port, &server_name, &res);
180
181        let res = self.client.resolve(id, name, RecordType::AAAA, RecordClass::INET).await?;
182        out.extend(Self::servers_from_answers(port, &server_name, &res));
183
184        Ok(out)
185    }
186
187    fn resolve_unix_addr(name: &str) -> Vec<Server> {
188        let path = PathBuf::from(name);
189        vec![Server::without_name(ServerID::from(path))]
190    }
191
192    fn resolve_socket_addr(name: &str, addr: SocketAddr) -> Result<Vec<Server>, MtopError> {
193        let (host, _port) = Self::host_and_port(name)?;
194        let server_name = Self::server_name(host)?;
195        Ok(vec![Server::new(ServerID::from(addr), server_name)])
196    }
197
198    fn resolve_bare_host(name: &str) -> Result<Vec<Server>, MtopError> {
199        let (host, port) = Self::host_and_port(name)?;
200        let server_name = Self::server_name(host)?;
201        Ok(vec![Server::new(ServerID::from((host, port)), server_name)])
202    }
203
204    fn servers_from_answers(port: u16, server_name: &ServerName<'static>, message: &Message) -> Vec<Server> {
205        let mut servers = HashSet::with_capacity(message.answers().len());
206
207        for answer in message.answers() {
208            let id = match answer.rdata() {
209                RecordData::A(data) => {
210                    let addr = SocketAddr::new(IpAddr::V4(data.addr()), port);
211                    ServerID::from(addr)
212                }
213                RecordData::AAAA(data) => {
214                    let addr = SocketAddr::new(IpAddr::V6(data.addr()), port);
215                    ServerID::from(addr)
216                }
217                RecordData::SRV(data) => {
218                    let target = data.target().to_string();
219
220                    ServerID::from((&target as &str, port))
221                }
222                _ => {
223                    tracing::warn!(message = "unexpected record data for answer", answer = ?answer);
224                    continue;
225                }
226            };
227
228            // Insert server into a HashSet to deduplicate them. We can potentially end up with
229            // duplicates when a SRV query returns multiple answers per hostname (such as when
230            // each host has more than a single port). Because we ignore the port number from the
231            // SRV answer we need to deduplicate here.
232            servers.insert(Server::new(id, server_name.to_owned()));
233        }
234
235        servers.into_iter().collect()
236    }
237
238    fn host_and_port(name: &str) -> Result<(&str, u16), MtopError> {
239        name.rsplit_once(':')
240            .ok_or_else(|| {
241                MtopError::configuration(format!(
242                    "invalid server name '{}', must be of the form 'host:port'",
243                    name
244                ))
245            })
246            // IPv6 addresses use brackets around them to disambiguate them from a port number.
247            // Since we're parsing the host and port, strip the brackets because they aren't
248            // needed or valid to include in a TLS ServerName.
249            .map(|(host, port)| (host.trim_start_matches('[').trim_end_matches(']'), port))
250            .and_then(|(host, port)| {
251                port.parse().map(|p| (host, p)).map_err(|e| {
252                    MtopError::configuration_cause(format!("unable to parse port number from '{}'", name), e)
253                })
254            })
255    }
256
257    fn server_name(host: &str) -> Result<ServerName<'static>, MtopError> {
258        ServerName::try_from(host)
259            .map(|s| s.to_owned())
260            .map_err(|e| MtopError::configuration_cause(format!("invalid server name '{}'", host), e))
261    }
262}
263
264impl fmt::Debug for Discovery {
265    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
266        f.debug_struct("Discovery").field("client", &"...").finish()
267    }
268}
269
270#[cfg(test)]
271mod test {
272    use super::{Discovery, ServerID};
273    use crate::core::MtopError;
274    use crate::dns::{
275        DnsClient, Flags, Message, MessageId, Name, Question, Record, RecordClass, RecordData, RecordDataA,
276        RecordDataAAAA, RecordDataSRV, RecordType,
277    };
278    use async_trait::async_trait;
279    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
280    use std::str::FromStr;
281    use tokio::sync::Mutex;
282
283    #[test]
284    fn test_server_id_from_ipv4_addr() {
285        let addr = SocketAddr::from((Ipv4Addr::LOCALHOST, 11211));
286        let id = ServerID::from(addr);
287        assert_eq!("127.0.0.1:11211", id.to_string());
288    }
289
290    #[test]
291    fn test_server_id_from_ipv6_addr() {
292        let addr = SocketAddr::from((Ipv6Addr::LOCALHOST, 11211));
293        let id = ServerID::from(addr);
294        assert_eq!("[::1]:11211", id.to_string());
295    }
296
297    #[test]
298    fn test_server_id_from_ipv4_pair() {
299        let pair = ("10.1.1.22", 11212);
300        let id = ServerID::from(pair);
301        assert_eq!("10.1.1.22:11212", id.to_string());
302    }
303
304    #[test]
305    fn test_server_id_from_ipv6_pair() {
306        let pair = ("::1", 11212);
307        let id = ServerID::from(pair);
308        assert_eq!("[::1]:11212", id.to_string());
309    }
310
311    #[test]
312    fn test_server_id_from_host_pair() {
313        let pair = ("cache.example.com", 11211);
314        let id = ServerID::from(pair);
315        assert_eq!("cache.example.com:11211", id.to_string());
316    }
317
318    struct MockDnsClient {
319        responses: Mutex<Vec<Message>>,
320    }
321
322    impl MockDnsClient {
323        fn new(responses: Vec<Message>) -> Self {
324            Self {
325                responses: Mutex::new(responses),
326            }
327        }
328    }
329
330    #[async_trait]
331    impl DnsClient for MockDnsClient {
332        async fn resolve(
333            &self,
334            _id: MessageId,
335            _name: Name,
336            _rtype: RecordType,
337            _rclass: RecordClass,
338        ) -> Result<Message, MtopError> {
339            let mut responses = self.responses.lock().await;
340            let res = responses.pop().unwrap();
341            Ok(res)
342        }
343    }
344
345    fn response_with_answers(rtype: RecordType, records: Vec<Record>) -> Message {
346        let flags = Flags::default().set_recursion_desired().set_recursion_available();
347        let mut message = Message::new(MessageId::random(), flags)
348            .add_question(Question::new(Name::from_str("example.com.").unwrap(), rtype));
349
350        for r in records {
351            message = message.add_answer(r);
352        }
353
354        message
355    }
356
357    #[tokio::test]
358    async fn test_dns_client_resolve_a_aaaa() {
359        let response_a = response_with_answers(
360            RecordType::A,
361            vec![Record::new(
362                Name::from_str("example.com.").unwrap(),
363                RecordType::A,
364                RecordClass::INET,
365                300,
366                RecordData::A(RecordDataA::new(Ipv4Addr::new(10, 1, 1, 1))),
367            )],
368        );
369
370        let response_aaaa = response_with_answers(
371            RecordType::AAAA,
372            vec![Record::new(
373                Name::from_str("example.com.").unwrap(),
374                RecordType::AAAA,
375                RecordClass::INET,
376                300,
377                RecordData::AAAA(RecordDataAAAA::new(Ipv6Addr::LOCALHOST)),
378            )],
379        );
380
381        let client = MockDnsClient::new(vec![response_a, response_aaaa]);
382        let discovery = Discovery::new(client);
383        let servers = discovery.resolve_by_proto("dns+example.com:11211").await.unwrap();
384
385        let ids = servers.iter().map(|s| s.id().clone()).collect::<Vec<_>>();
386        let id_a = ServerID::from("10.1.1.1:11211".parse::<SocketAddr>().unwrap());
387        let id_aaaa = ServerID::from("[::1]:11211".parse::<SocketAddr>().unwrap());
388
389        assert!(ids.contains(&id_a), "expected {:?} to contain {:?}", ids, id_a);
390        assert!(ids.contains(&id_aaaa), "expected {:?} to contain {:?}", ids, id_aaaa);
391    }
392
393    #[tokio::test]
394    async fn test_dns_client_resolve_srv() {
395        let response = response_with_answers(
396            RecordType::SRV,
397            vec![
398                Record::new(
399                    Name::from_str("_cache.example.com.").unwrap(),
400                    RecordType::SRV,
401                    RecordClass::INET,
402                    300,
403                    RecordData::SRV(RecordDataSRV::new(
404                        100,
405                        10,
406                        11211,
407                        Name::from_str("cache01.example.com.").unwrap(),
408                    )),
409                ),
410                Record::new(
411                    Name::from_str("_cache.example.com.").unwrap(),
412                    RecordType::SRV,
413                    RecordClass::INET,
414                    300,
415                    RecordData::SRV(RecordDataSRV::new(
416                        100,
417                        10,
418                        11211,
419                        Name::from_str("cache02.example.com.").unwrap(),
420                    )),
421                ),
422            ],
423        );
424
425        let client = MockDnsClient::new(vec![response]);
426        let discovery = Discovery::new(client);
427        let servers = discovery.resolve_by_proto("dnssrv+_cache.example.com:11211").await.unwrap();
428
429        let ids = servers.iter().map(|s| s.id().clone()).collect::<Vec<_>>();
430        let id1 = ServerID::from(("cache01.example.com.", 11211));
431        let id2 = ServerID::from(("cache02.example.com.", 11211));
432
433        assert!(ids.contains(&id1), "expected {:?} to contain {:?}", ids, id1);
434        assert!(ids.contains(&id2), "expected {:?} to contain {:?}", ids, id2);
435    }
436
437    #[tokio::test]
438    async fn test_dns_client_resolve_srv_dupes() {
439        let response = response_with_answers(
440            RecordType::SRV,
441            vec![
442                Record::new(
443                    Name::from_str("_cache.example.com.").unwrap(),
444                    RecordType::SRV,
445                    RecordClass::INET,
446                    300,
447                    RecordData::SRV(RecordDataSRV::new(
448                        100,
449                        10,
450                        11211,
451                        Name::from_str("cache01.example.com.").unwrap(),
452                    )),
453                ),
454                Record::new(
455                    Name::from_str("_cache.example.com.").unwrap(),
456                    RecordType::SRV,
457                    RecordClass::INET,
458                    300,
459                    RecordData::SRV(RecordDataSRV::new(
460                        100,
461                        10,
462                        9105,
463                        Name::from_str("cache01.example.com.").unwrap(),
464                    )),
465                ),
466            ],
467        );
468
469        let client = MockDnsClient::new(vec![response]);
470        let discovery = Discovery::new(client);
471        let servers = discovery.resolve_by_proto("dnssrv+_cache.example.com:11211").await.unwrap();
472
473        let ids = servers.iter().map(|s| s.id().clone()).collect::<Vec<_>>();
474        let id = ServerID::from(("cache01.example.com.", 11211));
475
476        assert_eq!(ids, vec![id]);
477    }
478
479    #[tokio::test]
480    async fn test_dns_client_resolve_socket_addr() {
481        let name = "127.0.0.2:11211";
482        let sock: SocketAddr = "127.0.0.2:11211".parse().unwrap();
483
484        let client = MockDnsClient::new(vec![]);
485        let discovery = Discovery::new(client);
486        let servers = discovery.resolve_by_proto(name).await.unwrap();
487
488        let ids = servers.iter().map(|s| s.id().clone()).collect::<Vec<_>>();
489        let id = ServerID::from(sock);
490
491        assert!(ids.contains(&id), "expected {:?} to contain {:?}", ids, id);
492    }
493
494    #[tokio::test]
495    async fn test_dns_client_resolve_bare_host() {
496        let name = "localhost:11211";
497
498        let client = MockDnsClient::new(vec![]);
499        let discovery = Discovery::new(client);
500        let servers = discovery.resolve_by_proto(name).await.unwrap();
501
502        let ids = servers.iter().map(|s| s.id().clone()).collect::<Vec<_>>();
503        let id = ServerID::from(("localhost", 11211));
504
505        assert!(ids.contains(&id), "expected {:?} to contain {:?}", ids, id);
506    }
507}
mtop_client/discovery.rs

mtop_client/
discovery.rs
