pharos_plugins/proxy/
health.rs

1//! On-demand upstream health checks, independent of the forwarding transport.
2
3use pharos_proto::{Message, Name, Opcode};
4use pharos_server::metrics::ProxyMetrics;
5use std::{
6    io,
7    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
8    sync::{
9        Arc, Mutex,
10        atomic::{AtomicU32, Ordering},
11    },
12    time::Duration,
13};
14use tokio::{
15    io::{AsyncReadExt, AsyncWriteExt},
16    net::UdpSocket,
17    task::JoinHandle,
18    time::{sleep, timeout},
19};
20
21const DIAL_TIMEOUT: Duration = Duration::from_secs(2);
22const IO_TIMEOUT: Duration = Duration::from_secs(1);
23
24/// Settings for on-demand checks of one upstream.
25#[derive(Clone)]
26pub struct HealthConfig {
27    pub address: SocketAddr,
28    pub domain: Name,
29    pub recursive: bool,
30    pub force_tcp: bool,
31    pub source_address: Option<IpAddr>,
32    pub max_fails: u32,
33    pub interval: Duration,
34    pub metrics: Option<ProxyMetrics>,
35}
36
37impl HealthConfig {
38    /// Defaults to recursive `. IN NS` checks over UDP every 500ms.
39    #[must_use]
40    pub fn new(address: SocketAddr, max_fails: u32) -> Self {
41        Self {
42            address,
43            domain: ".".parse().expect("root is a valid DNS name"),
44            recursive: true,
45            force_tcp: false,
46            source_address: None,
47            max_fails,
48            interval: Duration::from_millis(500),
49            metrics: None,
50        }
51    }
52}
53
54#[derive(Default)]
55struct State {
56    running: bool,
57    shutdown: bool,
58    generation: u64,
59    task: Option<JoinHandle<()>>,
60}
61
62struct Shared {
63    config: HealthConfig,
64    state: Mutex<State>,
65    fails: AtomicU32,
66}
67
68/// A single upstream's on-demand probe and consecutive failed-check count.
69pub struct Health {
70    shared: Arc<Shared>,
71}
72
73impl Health {
74    pub(super) fn config(&self) -> HealthConfig {
75        self.shared.config.clone()
76    }
77    #[must_use]
78    pub fn new(address: SocketAddr, max_fails: u32) -> Self {
79        Self::with_config(HealthConfig::new(address, max_fails))
80    }
81
82    #[must_use]
83    pub fn with_config(config: HealthConfig) -> Self {
84        Self {
85            shared: Arc::new(Shared {
86                config,
87                state: Mutex::new(State::default()),
88                fails: AtomicU32::new(0),
89            }),
90        }
91    }
92
93    /// Starts one probe loop after an exchange error; concurrent calls are coalesced.
94    /// Must be called from a Tokio runtime. `max_fails == 0` disables probing.
95    pub fn start(&self) {
96        if self.shared.config.max_fails == 0 {
97            return;
98        }
99        let mut state = self.shared.state.lock().unwrap();
100        if state.shutdown {
101            return;
102        }
103        if state.running && state.task.as_ref().is_some_and(|task| !task.is_finished()) {
104            return;
105        }
106        state.running = true;
107        state.generation = state.generation.wrapping_add(1);
108        let generation = state.generation;
109        let shared = Arc::clone(&self.shared);
110        state.task = Some(tokio::spawn(async move {
111            loop {
112                let healthy = check(&shared.config).await.is_ok();
113                {
114                    let mut state = shared.state.lock().unwrap();
115                    if state.generation != generation {
116                        return;
117                    }
118                    if healthy {
119                        shared.fails.store(0, Ordering::Relaxed);
120                        state.running = false;
121                        return;
122                    }
123                    let _ =
124                        shared
125                            .fails
126                            .try_update(Ordering::Relaxed, Ordering::Relaxed, |fails| {
127                                Some(fails.saturating_add(1))
128                            });
129                    if let Some(metrics) = &shared.config.metrics {
130                        metrics.healthcheck_failure();
131                    }
132                }
133                sleep(shared.config.interval).await;
134            }
135        }));
136    }
137
138    /// Whether failed checks strictly exceed the configured threshold.
139    #[must_use]
140    pub fn down(&self) -> bool {
141        self.shared.config.max_fails != 0 && self.fails() > self.shared.config.max_fails
142    }
143
144    /// Consecutive failed checks, reset to zero by a successful check.
145    #[must_use]
146    pub fn fails(&self) -> u32 {
147        self.shared.fails.load(Ordering::Relaxed)
148    }
149
150    /// Cancels any current probe and closes its socket on cancellation.
151    pub fn stop(&self) {
152        let mut state = self.shared.state.lock().unwrap();
153        if let Some(task) = state.task.take() {
154            task.abort();
155        }
156        state.running = false;
157        state.generation = state.generation.wrapping_add(1);
158    }
159
160    pub(super) fn shutdown(&self) {
161        let mut state = self.shared.state.lock().unwrap();
162        state.shutdown = true;
163        if let Some(task) = state.task.take() {
164            task.abort();
165        }
166        state.running = false;
167        state.generation = state.generation.wrapping_add(1);
168    }
169}
170
171impl Drop for Health {
172    fn drop(&mut self) {
173        self.stop();
174    }
175}
176
177async fn check(config: &HealthConfig) -> io::Result<()> {
178    // RFC 5452 §9.2: outgoing queries use unpredictable IDs.
179    let mut id = [0; 2];
180    getrandom::getrandom(&mut id).map_err(|_| io::Error::other("random ID generation failed"))?;
181    let mut query = Vec::with_capacity(12 + config.domain.as_wire().len() + 4);
182    query.extend_from_slice(&id);
183    // RFC 1035 §4.1.1: RD requests recursive query processing.
184    query.extend_from_slice(&[u8::from(config.recursive), 0, 0, 1, 0, 0, 0, 0, 0, 0]);
185    query.extend_from_slice(config.domain.as_wire());
186    // RFC 1035 §3.2.2: NS has type 2.
187    // RFC 1035 §3.2.4: IN has class 1.
188    query.extend_from_slice(&[0, 2, 0, 1]);
189
190    let response = if config.force_tcp {
191        let mut stream = timeout(
192            DIAL_TIMEOUT,
193            super::connect_tcp(config.address, config.source_address),
194        )
195        .await??;
196        timeout(IO_TIMEOUT, async {
197            // RFC 1035 §4.2.2: TCP DNS messages have a two-octet length prefix.
198            stream
199                .write_all(&(query.len() as u16).to_be_bytes())
200                .await?;
201            stream.write_all(&query).await
202        })
203        .await??;
204        timeout(IO_TIMEOUT, async {
205            // RFC 1035 §4.2.2: TCP DNS messages have a two-octet length prefix.
206            let len = stream.read_u16().await?;
207            let mut response = vec![0; usize::from(len)];
208            stream.read_exact(&mut response).await?;
209            Ok::<_, io::Error>(response)
210        })
211        .await??
212    } else {
213        let ip = if config.address.is_ipv4() {
214            IpAddr::V4(Ipv4Addr::UNSPECIFIED)
215        } else {
216            IpAddr::V6(Ipv6Addr::UNSPECIFIED)
217        };
218        let socket = timeout(DIAL_TIMEOUT, async {
219            let socket =
220                UdpSocket::bind(SocketAddr::new(config.source_address.unwrap_or(ip), 0)).await?;
221            super::enable_recv_err(&socket, config.address);
222            socket.connect(config.address).await?;
223            Ok::<_, io::Error>(socket)
224        })
225        .await??;
226        timeout(IO_TIMEOUT, socket.send(&query)).await??;
227        let mut response = vec![0; usize::from(u16::MAX)];
228        let len = timeout(IO_TIMEOUT, socket.recv(&mut response)).await??;
229        response.truncate(len);
230        response
231    };
232
233    // CoreDNS treats even a partially decoded response as healthy if its header is sane.
234    if Message::parse(&response).is_ok()
235        || (response.len() >= 12
236            && (response[2] & 0x80 != 0 || (response[2] >> 3) & 15 == Opcode::QUERY.0))
237    {
238        Ok(())
239    } else {
240        Err(io::ErrorKind::InvalidData.into())
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use super::*;
247    use pharos_proto::{Class, Message, Type};
248    use tokio::{
249        net::{TcpListener, UdpSocket},
250        time::timeout,
251    };
252
253    #[tokio::test]
254    async fn default_probe_sends_recursive_root_ns_over_udp() {
255        let server = UdpSocket::bind("127.0.0.1:0").await.unwrap();
256        let health = Health::new(server.local_addr().unwrap(), 2);
257        health.start();
258        let mut wire = [0; 512];
259        let (len, peer) = timeout(Duration::from_secs(1), server.recv_from(&mut wire))
260            .await
261            .unwrap()
262            .unwrap();
263        let message = Message::parse(&wire[..len]).unwrap();
264        let question = message.questions().next().unwrap();
265        assert_eq!(question.name.to_string(), ".");
266        assert_eq!(question.qtype, Type::NS);
267        assert_eq!(question.qclass, Class::IN);
268        assert!(message.rd());
269        wire[2] |= 0x80;
270        server.send_to(&wire[..len], peer).await.unwrap();
271    }
272
273    #[tokio::test]
274    async fn failed_probe_updates_injected_metric_and_recovery_resets_fails() {
275        let server = UdpSocket::bind("127.0.0.1:0").await.unwrap();
276        let registry = pharos_server::metrics::Registry::new("test", "test", "test");
277        let mut config = HealthConfig::new(server.local_addr().unwrap(), 1);
278        config.metrics = Some(registry.proxy("forward", "test-upstream"));
279        config.interval = Duration::from_millis(25);
280        let health = Health::with_config(config);
281        health.start();
282        let mut wire = [0; 512];
283        let (_, peer) = timeout(Duration::from_secs(1), server.recv_from(&mut wire))
284            .await
285            .unwrap()
286            .unwrap();
287        server.send_to(&[0, 1], peer).await.unwrap();
288        timeout(Duration::from_secs(1), async {
289            while health.fails() < 1 {
290                tokio::task::yield_now().await;
291            }
292        })
293        .await
294        .unwrap();
295        assert!(!health.down());
296        let (len, peer) = timeout(Duration::from_secs(1), server.recv_from(&mut wire))
297            .await
298            .unwrap()
299            .unwrap();
300        server.send_to(&[0, 1], peer).await.unwrap();
301        timeout(Duration::from_secs(1), async {
302            while !health.down() {
303                tokio::task::yield_now().await;
304            }
305        })
306        .await
307        .unwrap();
308        let (len2, peer) = timeout(Duration::from_secs(1), server.recv_from(&mut wire))
309            .await
310            .unwrap()
311            .unwrap();
312        assert_eq!(len, len2);
313        wire[2] |= 0x80;
314        wire[3] = 2; // SERVFAIL is still a response from a live upstream.
315        server.send_to(&wire[..len2], peer).await.unwrap();
316        timeout(Duration::from_secs(1), async {
317            while health.fails() != 0 {
318                tokio::task::yield_now().await;
319            }
320        })
321        .await
322        .unwrap();
323        let mut text = String::new();
324        registry.encode(&mut text);
325        assert!(text.contains(
326            "pharos_proxy_healthcheck_failures_total{proxy_name=\"forward\",to=\"test-upstream\"} 2"
327        ));
328        assert!(!health.down());
329    }
330
331    #[tokio::test]
332    async fn configured_probe_sends_nonrecursive_name_over_tcp() {
333        let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
334        let mut config = HealthConfig::new(server.local_addr().unwrap(), 2);
335        config.domain = "example.org.".parse().unwrap();
336        config.recursive = false;
337        config.force_tcp = true;
338        let health = Health::with_config(config);
339        health.start();
340        let (mut stream, _) = timeout(Duration::from_secs(1), server.accept())
341            .await
342            .unwrap()
343            .unwrap();
344        let len = stream.read_u16().await.unwrap();
345        let mut wire = vec![0; usize::from(len)];
346        stream.read_exact(&mut wire).await.unwrap();
347        let message = Message::parse(&wire).unwrap();
348        let question = message.questions().next().unwrap();
349        assert_eq!(question.name.to_string(), "example.org.");
350        assert_eq!(question.qtype, Type::NS);
351        assert_eq!(question.qclass, Class::IN);
352        assert!(!message.rd());
353        wire[2] |= 0x80;
354        wire[3] = 5;
355        stream.write_u16(len).await.unwrap();
356        stream.write_all(&wire).await.unwrap();
357    }
358
359    #[tokio::test]
360    async fn repeated_starts_share_one_probe_and_stop_closes_it() {
361        let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
362        let mut config = HealthConfig::new(server.local_addr().unwrap(), 1);
363        config.force_tcp = true;
364        let health = Health::with_config(config);
365        health.start();
366        let (mut stream, _) = timeout(Duration::from_secs(1), server.accept())
367            .await
368            .unwrap()
369            .unwrap();
370        let len = stream.read_u16().await.unwrap();
371        let mut wire = vec![0; usize::from(len)];
372        stream.read_exact(&mut wire).await.unwrap();
373        for _ in 0..10 {
374            health.start();
375        }
376        assert!(
377            timeout(Duration::from_millis(50), server.accept())
378                .await
379                .is_err()
380        );
381        health.stop();
382        assert_eq!(
383            timeout(Duration::from_secs(1), stream.read(&mut wire))
384                .await
385                .unwrap()
386                .unwrap(),
387            0
388        );
389        health.start();
390        timeout(Duration::from_secs(1), server.accept())
391            .await
392            .unwrap()
393            .unwrap();
394    }
395
396    #[tokio::test]
397    async fn zero_max_fails_disables_checks() {
398        let server = UdpSocket::bind("127.0.0.1:0").await.unwrap();
399        let health = Health::new(server.local_addr().unwrap(), 0);
400        health.start();
401        let mut wire = [0; 512];
402        assert!(
403            timeout(Duration::from_millis(50), server.recv_from(&mut wire))
404                .await
405                .is_err()
406        );
407        assert_eq!(health.fails(), 0);
408        assert!(!health.down());
409    }
410
411    #[tokio::test]
412    async fn tcp_read_timeout_counts_as_a_failed_probe() {
413        let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
414        let mut config = HealthConfig::new(server.local_addr().unwrap(), 1);
415        config.force_tcp = true;
416        let health = Health::with_config(config);
417        health.start();
418        let (_stream, _) = timeout(Duration::from_secs(1), server.accept())
419            .await
420            .unwrap()
421            .unwrap();
422        timeout(Duration::from_secs(2), async {
423            while health.fails() == 0 {
424                tokio::task::yield_now().await;
425            }
426        })
427        .await
428        .unwrap();
429        assert_eq!(health.fails(), 1);
430        assert!(!health.down());
431    }
432}
pharos_plugins/proxy/health.rs

pharos_plugins/proxy/
health.rs
