hickory_resolver/name_server/
name_server_state.rs

1// Copyright 2015-2019 Benjamin Fry <benjaminfry@me.com>
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// https://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// https://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use std::cmp::Ordering;
9use std::sync::atomic::{self, AtomicU8};
10use std::sync::Arc;
11use std::time::Instant;
12
13use futures_util::lock::Mutex;
14use proto::op::Edns;
15
16pub(crate) struct NameServerState {
17    conn_state: AtomicU8,
18    remote_edns: Mutex<Arc<Option<Edns>>>,
19}
20
21/// State of a connection with a remote NameServer.
22#[derive(Debug, Eq, PartialEq, Copy, Clone)]
23#[repr(u8)]
24enum NameServerStateInner {
25    /// For some reason the connection failed. For UDP this would generally be a timeout
26    ///  for TCP this could be either Connection could never be established, or it
27    ///  failed at some point after. The Failed state should *not* be entered due to an
28    ///  error contained in a Message received from the server. In All cases to reestablish
29    ///  a new connection will need to be created.
30    Failed = 0,
31    /// Initial state, if Edns is not none, then Edns will be requested
32    Init = 1,
33    /// There has been successful communication with the remote.
34    ///  if no Edns is associated, then the remote does not support Edns
35    Established = 2,
36}
37
38impl From<NameServerStateInner> for u8 {
39    /// used for ordering purposes. The highest priority is placed on open connections
40    fn from(val: NameServerStateInner) -> Self {
41        val as Self
42    }
43}
44
45impl From<u8> for NameServerStateInner {
46    fn from(val: u8) -> Self {
47        match val {
48            2 => Self::Established,
49            1 => Self::Init,
50            _ => Self::Failed,
51        }
52    }
53}
54
55impl NameServerState {
56    fn store(&self, conn_state: NameServerStateInner) {
57        self.conn_state
58            .store(conn_state.into(), atomic::Ordering::Release);
59    }
60
61    fn load(&self) -> NameServerStateInner {
62        NameServerStateInner::from(self.conn_state.load(atomic::Ordering::Acquire))
63    }
64
65    /// Set at the new Init state
66    ///
67    /// If send_dns is some, this will be sent on the first request when it is established
68    pub(crate) fn init(_send_edns: Option<Edns>) -> Self {
69        // TODO: need to track send_edns
70        Self {
71            conn_state: AtomicU8::new(NameServerStateInner::Init.into()),
72            remote_edns: Mutex::new(Arc::new(None)),
73        }
74    }
75
76    /// Set at the new Init state
77    ///
78    /// If send_dns is some, this will be sent on the first request when it is established
79    pub(crate) fn reinit(&self, _send_edns: Option<Edns>) {
80        // eventually do this
81        // self.send_edns.lock() = send_edns;
82
83        self.store(NameServerStateInner::Init);
84    }
85
86    /// Transition to the Established state
87    ///
88    /// If remote_edns is Some, then it will be used to effect things like buffer sizes based on
89    ///   the remote's support.
90    pub(crate) fn establish(&self, remote_edns: Option<Edns>) {
91        if remote_edns.is_some() {
92            // best effort locking, we'll assume a different user of this connection is storing the same thing...
93            if let Some(mut current_edns) = self.remote_edns.try_lock() {
94                *current_edns = Arc::new(remote_edns)
95            }
96        }
97
98        self.store(NameServerStateInner::Established);
99    }
100
101    /// transition to the Failed state
102    ///
103    /// when is the time of the failure
104    ///
105    /// * when - deprecated
106    pub(crate) fn fail(&self, _when: /* FIXME: remove in 0.20 */ Instant) {
107        self.store(NameServerStateInner::Failed);
108    }
109
110    /// True if this is in the Failed state
111    pub(crate) fn is_failed(&self) -> bool {
112        NameServerStateInner::Failed == self.load()
113    }
114}
115
116impl Ord for NameServerStateInner {
117    fn cmp(&self, other: &Self) -> Ordering {
118        let (self_num, other_num) = (u8::from(*self), u8::from(*other));
119        self_num.cmp(&other_num)
120    }
121}
122
123impl PartialOrd for NameServerStateInner {
124    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
125        Some(self.cmp(other))
126    }
127}
128
129impl Ord for NameServerState {
130    fn cmp(&self, other: &Self) -> Ordering {
131        let other = other.load();
132        self.load().cmp(&other)
133    }
134}
135
136impl PartialOrd for NameServerState {
137    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
138        Some(self.cmp(other))
139    }
140}
141
142impl PartialEq for NameServerState {
143    fn eq(&self, other: &Self) -> bool {
144        self.load() == other.load()
145    }
146}
147
148impl Eq for NameServerState {}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use crate::name_server::NameServerState;
154
155    #[test]
156    fn test_state_cmp() {
157        let init = NameServerState::init(None);
158
159        let established = NameServerState::init(None);
160        established.establish(None);
161
162        let failed = NameServerState::init(None);
163        failed.fail(Instant::now());
164
165        assert_eq!(init.cmp(&init), Ordering::Equal);
166        assert_eq!(init.cmp(&established), Ordering::Less);
167        assert_eq!(init.cmp(&failed), Ordering::Greater);
168        assert_eq!(established.cmp(&established), Ordering::Equal);
169        assert_eq!(established.cmp(&failed), Ordering::Greater);
170        assert_eq!(failed.cmp(&failed), Ordering::Equal);
171    }
172}
hickory_resolver/name_server/name_server_state.rs

hickory_resolver/name_server/
name_server_state.rs
