1 // Copyright 2015-2019 Benjamin Fry <benjaminfry@me.com> 2 // 3 // Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or 4 // http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or 5 // http://opensource.org/licenses/MIT>, at your option. This file may not be 6 // copied, modified, or distributed except according to those terms. 7 8 use std::cmp::Ordering; 9 use std::sync::atomic::{self, AtomicU8}; 10 use std::sync::Arc; 11 use std::time::Instant; 12 13 use futures::lock::Mutex; 14 use proto::op::Edns; 15 16 pub 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)] 24 enum 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 38 impl From<NameServerStateInner> for u8 { 39 /// used for ordering purposes. The highest priority is placed on open connections from(val: NameServerStateInner) -> u840 fn from(val: NameServerStateInner) -> u8 { 41 val as u8 42 } 43 } 44 45 impl From<u8> for NameServerStateInner { from(val: u8) -> Self46 fn from(val: u8) -> Self { 47 match val { 48 2 => NameServerStateInner::Established, 49 1 => NameServerStateInner::Init, 50 _ => NameServerStateInner::Failed, 51 } 52 } 53 } 54 55 impl NameServerState { store(&self, conn_state: NameServerStateInner)56 fn store(&self, conn_state: NameServerStateInner) { 57 self.conn_state 58 .store(conn_state.into(), atomic::Ordering::Release); 59 } 60 load(&self) -> NameServerStateInner61 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 init(_send_edns: Option<Edns>) -> Self68 pub fn init(_send_edns: Option<Edns>) -> Self { 69 // TODO: need to track send_edns 70 NameServerState { 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 reinit(&self, _send_edns: Option<Edns>)79 pub 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. establish(&self, remote_edns: Option<Edns>)90 pub 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 fail(&self, _when: Instant)106 pub 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 is_failed(&self) -> bool111 pub(crate) fn is_failed(&self) -> bool { 112 NameServerStateInner::Failed == self.load() 113 } 114 } 115 116 impl Ord for NameServerStateInner { cmp(&self, other: &Self) -> Ordering117 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 123 impl PartialOrd for NameServerStateInner { partial_cmp(&self, other: &Self) -> Option<Ordering>124 fn partial_cmp(&self, other: &Self) -> Option<Ordering> { 125 Some(self.cmp(other)) 126 } 127 } 128 129 impl Ord for NameServerState { cmp(&self, other: &Self) -> Ordering130 fn cmp(&self, other: &Self) -> Ordering { 131 let other = other.load(); 132 self.load().cmp(&other) 133 } 134 } 135 136 impl PartialOrd for NameServerState { partial_cmp(&self, other: &Self) -> Option<Ordering>137 fn partial_cmp(&self, other: &Self) -> Option<Ordering> { 138 Some(self.cmp(other)) 139 } 140 } 141 142 impl PartialEq for NameServerState { eq(&self, other: &Self) -> bool143 fn eq(&self, other: &Self) -> bool { 144 self.load() == other.load() 145 } 146 } 147 148 impl Eq for NameServerState {} 149 150 #[cfg(test)] 151 mod tests { 152 use super::*; 153 use crate::name_server::NameServerState; 154 155 #[test] test_state_cmp()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 } 173