pharos_plugins/proxy/
mod.rs

1//! Upstream DNS exchange. Independent cache shards hold idle connections only;
2//! no lock or manager task is held while a Query waits for its Response.
3//! DNS-over-TLS is deferred to M3.
4//! CoreDNS's `forward` performs the `prefer_udp` TC retry; Pharos does it here,
5//! so a `forward` caller must not retry the same truncated Response again.
6
7pub mod health;
8pub use health::{Health, HealthConfig};
9
10use pharos_proto::{Message, MessageError, NameError, Rcode, Type};
11use pharos_server::metrics::{ProxyConnectionCache, ProxyMetrics, ProxyRequestDuration, Registry};
12use std::{
13    collections::VecDeque,
14    fmt, io,
15    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
16    sync::{
17        Mutex, OnceLock,
18        atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering},
19    },
20    time::{Duration, Instant},
21};
22use tokio::{
23    io::{AsyncReadExt, AsyncWriteExt},
24    net::{TcpSocket, TcpStream, UdpSocket},
25    time::timeout,
26};
27
28const SHARDS: usize = 32;
29static NEXT_SHARD: AtomicUsize = AtomicUsize::new(0);
30thread_local! { static SHARD: usize = NEXT_SHARD.fetch_add(1, Ordering::Relaxed) % SHARDS; }
31const EXPIRE: Duration = Duration::from_secs(10);
32const READ_TIMEOUT: Duration = Duration::from_secs(2);
33
34#[cfg(target_os = "linux")]
35fn enable_recv_err(socket: &UdpSocket, address: SocketAddr) {
36    use nix::sys::socket::{
37        setsockopt,
38        sockopt::{Ipv4RecvErr, Ipv6RecvErr},
39    };
40    let result = if address.is_ipv4() {
41        setsockopt(socket, Ipv4RecvErr, &true)
42    } else {
43        setsockopt(socket, Ipv6RecvErr, &true)
44    };
45    if let Err(error) = result {
46        pharos_server::logging::log(
47            "WARN",
48            "proxy",
49            &format!("failed to enable UDP error queue: {error}"),
50        );
51    }
52}
53
54#[cfg(not(target_os = "linux"))]
55fn enable_recv_err(_: &UdpSocket, _: SocketAddr) {}
56
57async fn connect_tcp(address: SocketAddr, source: Option<IpAddr>) -> io::Result<TcpStream> {
58    let socket = if address.is_ipv4() {
59        TcpSocket::new_v4()?
60    } else {
61        TcpSocket::new_v6()?
62    };
63    if let Some(source) = source {
64        socket.bind(SocketAddr::new(source, 0))?;
65    }
66    socket.connect(address).await
67}
68
69fn is_overflow(error: &MessageError) -> bool {
70    matches!(
71        error,
72        MessageError::Truncated | MessageError::Name(NameError::Truncated)
73    )
74}
75
76#[derive(Clone, Copy, Debug, Eq, PartialEq)]
77pub enum Transport {
78    Udp,
79    Tcp,
80}
81
82#[derive(Clone, Copy, Debug, Default)]
83pub struct Options {
84    pub force_tcp: bool,
85    pub prefer_udp: bool,
86}
87
88#[derive(Debug)]
89pub enum Error {
90    InvalidRequest,
91    UnsupportedRequest,
92    NoHealthy,
93    NoForward,
94    Stopped,
95    Mismatch,
96    CachedClosed,
97    Io(io::Error),
98}
99impl fmt::Display for Error {
100    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
101        match self {
102            Self::InvalidRequest => f.write_str("invalid Query"),
103            Self::UnsupportedRequest => f.write_str("unsupported Query"),
104            Self::NoHealthy => f.write_str("no healthy Upstreams"),
105            Self::NoForward => f.write_str("no Upstream configured"),
106            Self::Stopped => f.write_str("proxy: transport stopped"),
107            Self::Mismatch => f.write_str("Upstream Response does not match Query"),
108            Self::CachedClosed => f.write_str("cached connection was closed by peer"),
109            Self::Io(error) => write!(f, "{error}"),
110        }
111    }
112}
113impl std::error::Error for Error {
114    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
115        match self {
116            Self::Io(error) => Some(error),
117            _ => None,
118        }
119    }
120}
121impl From<io::Error> for Error {
122    fn from(error: io::Error) -> Self {
123        Self::Io(error)
124    }
125}
126
127#[derive(Debug)]
128pub struct Exchange {
129    pub len: usize,
130    pub local_addr: SocketAddr,
131    pub transport: Transport,
132}
133
134struct Cached {
135    socket: Socket,
136    created: Instant,
137    used: Instant,
138}
139
140enum Socket {
141    Udp(UdpSocket),
142    Tcp(TcpStream),
143}
144
145impl Socket {
146    fn local_addr(&self) -> io::Result<SocketAddr> {
147        match self {
148            Self::Udp(s) => s.local_addr(),
149            Self::Tcp(s) => s.local_addr(),
150        }
151    }
152    async fn send(&mut self, wire: &[u8]) -> io::Result<()> {
153        match self {
154            Self::Udp(s) => {
155                s.send(wire).await?;
156                Ok(())
157            }
158            // RFC 1035 §4.2.2: TCP DNS messages have a two-octet length prefix.
159            Self::Tcp(s) => {
160                s.write_all(&(wire.len() as u16).to_be_bytes()).await?;
161                s.write_all(wire).await
162            }
163        }
164    }
165    async fn receive(&mut self, wire: &mut [u8]) -> io::Result<usize> {
166        match self {
167            Self::Udp(s) => s.recv(wire).await,
168            // RFC 1035 §4.2.2: TCP DNS messages have a two-octet length prefix.
169            Self::Tcp(s) => {
170                let len = usize::from(s.read_u16().await?);
171                if len > wire.len() {
172                    return Err(io::ErrorKind::InvalidData.into());
173                }
174                s.read_exact(&mut wire[..len]).await?;
175                Ok(len)
176            }
177        }
178    }
179}
180
181pub struct Proxy {
182    address: SocketAddr,
183    shards: [[Mutex<VecDeque<Cached>>; SHARDS]; 2],
184    expire: Duration,
185    max_age: Duration,
186    max_idle_conns: usize,
187    idle: [AtomicUsize; 2],
188    stopped: AtomicBool,
189    started: Instant,
190    next_sweep: AtomicU64,
191    read_timeout: Duration,
192    average_dial_nanos: AtomicI64,
193    source_address: Option<IpAddr>,
194    metrics: Option<Metrics>,
195    health: Health,
196}
197
198struct Metrics {
199    proxy: ProxyMetrics,
200    cache: [ProxyConnectionCache; 2],
201    durations: Box<[OnceLock<ProxyRequestDuration>]>,
202}
203
204impl Proxy {
205    #[must_use]
206    pub fn new(address: SocketAddr) -> Self {
207        Self {
208            address,
209            shards: std::array::from_fn(|_| std::array::from_fn(|_| Mutex::new(VecDeque::new()))),
210            expire: EXPIRE,
211            max_age: Duration::ZERO,
212            max_idle_conns: 0,
213            idle: [AtomicUsize::new(0), AtomicUsize::new(0)],
214            stopped: AtomicBool::new(false),
215            started: Instant::now(),
216            next_sweep: AtomicU64::new(EXPIRE.as_nanos() as u64),
217            read_timeout: READ_TIMEOUT,
218            average_dial_nanos: AtomicI64::new(15_000_000_000),
219            source_address: None,
220            metrics: None,
221            health: Health::new(address, 2),
222        }
223    }
224
225    /// Set the maximum idle time before an Upstream connection is discarded.
226    pub fn set_expire(&mut self, expire: Duration) {
227        self.expire = expire;
228        self.next_sweep
229            .store(self.sweep_deadline(), Ordering::Relaxed);
230    }
231
232    fn sweep_deadline(&self) -> u64 {
233        self.started
234            .elapsed()
235            .as_nanos()
236            .saturating_add(self.expire.as_nanos())
237            .min(u128::from(u64::MAX)) as u64
238    }
239
240    fn sweep(&self) {
241        let now = self.started.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64;
242        let next = self.next_sweep.load(Ordering::Relaxed);
243        if now < next
244            || self
245                .next_sweep
246                .compare_exchange(
247                    next,
248                    self.sweep_deadline(),
249                    Ordering::Relaxed,
250                    Ordering::Relaxed,
251                )
252                .is_err()
253        {
254            return;
255        }
256        // shortcut: without traffic expired sockets remain until stop; add a timer if idle cleanup is required.
257        for (protocol, shard) in self
258            .shards
259            .iter()
260            .enumerate()
261            .flat_map(|(p, s)| s.iter().map(move |s| (p, s)))
262        {
263            let mut queue = shard.lock().unwrap();
264            let before = queue.len();
265            queue.retain(|conn| {
266                conn.used.elapsed() <= self.expire
267                    && (self.max_age.is_zero() || conn.created.elapsed() <= self.max_age)
268            });
269            self.idle[protocol].fetch_sub(before - queue.len(), Ordering::Relaxed);
270        }
271    }
272
273    /// Stop probes and close all idle Upstream connections.
274    pub fn stop(&self) {
275        self.stopped.store(true, Ordering::Release);
276        self.health.shutdown();
277        for (protocol, shard) in self
278            .shards
279            .iter()
280            .enumerate()
281            .flat_map(|(p, s)| s.iter().map(move |s| (p, s)))
282        {
283            let mut queue = shard.lock().unwrap();
284            self.idle[protocol].fetch_sub(queue.len(), Ordering::Relaxed);
285            queue.clear();
286        }
287    }
288
289    /// Limit idle connections per Transport for this Upstream; zero means unlimited.
290    pub fn set_max_idle_conns(&mut self, max: usize) {
291        self.max_idle_conns = max;
292    }
293
294    /// Set maximum connection age from creation; zero means unlimited.
295    pub fn set_max_age(&mut self, max_age: Duration) {
296        self.max_age = max_age;
297    }
298
299    /// Set the Response read timeout (two seconds by default).
300    pub fn set_read_timeout(&mut self, read_timeout: Duration) {
301        self.read_timeout = read_timeout;
302    }
303
304    fn dial_timeout(&self) -> Duration {
305        let average = self.average_dial_nanos.load(Ordering::Relaxed);
306        Duration::from_nanos(if average < 1_000_000_000 {
307            1_000_000_000
308        } else if average < 15_000_000_000 {
309            (2 * average) as u64
310        } else {
311            30_000_000_000
312        })
313    }
314
315    fn update_dial_timeout(&self, duration: Duration) {
316        let average = self.average_dial_nanos.load(Ordering::Relaxed);
317        let observed = duration.as_nanos().min(i64::MAX as u128) as i64;
318        self.average_dial_nanos
319            .fetch_add((observed - average) / 4, Ordering::Relaxed);
320    }
321
322    async fn dial(&self, transport: Transport) -> io::Result<Socket> {
323        let start = Instant::now();
324        let result = timeout(self.dial_timeout(), async {
325            if transport == Transport::Tcp {
326                Ok(Socket::Tcp(
327                    connect_tcp(self.address, self.source_address).await?,
328                ))
329            } else {
330                let ip = if self.address.is_ipv4() {
331                    IpAddr::V4(Ipv4Addr::UNSPECIFIED)
332                } else {
333                    IpAddr::V6(Ipv6Addr::UNSPECIFIED)
334                };
335                let socket =
336                    UdpSocket::bind(SocketAddr::new(self.source_address.unwrap_or(ip), 0)).await?;
337                enable_recv_err(&socket, self.address);
338                socket.connect(self.address).await?;
339                Ok(Socket::Udp(socket))
340            }
341        })
342        .await;
343        self.update_dial_timeout(start.elapsed());
344        result.map_err(|_| io::ErrorKind::TimedOut)?
345    }
346
347    /// Bind new Upstream and health-check connections to this local IP.
348    pub fn set_source_address(&mut self, source: Option<IpAddr>) {
349        if source == self.source_address {
350            return;
351        }
352        for shard in self.shards.iter_mut().flatten() {
353            shard.get_mut().unwrap().clear();
354        }
355        for idle in &self.idle {
356            idle.store(0, Ordering::Relaxed);
357        }
358        self.source_address = source;
359        let mut config = self.health.config();
360        config.source_address = source;
361        self.health = Health::with_config(config);
362    }
363
364    /// Register the Upstream's metrics with the server's exporter.
365    pub fn set_metrics(&mut self, registry: &Registry, name: &str) {
366        let proxy = registry.proxy(name, &self.address.to_string());
367        let mut config = self.health.config();
368        config.metrics = Some(proxy.clone());
369        self.health = Health::with_config(config);
370        self.metrics = Some(Metrics {
371            cache: [proxy.connection_cache("udp"), proxy.connection_cache("tcp")],
372            durations: (0..4096).map(|_| OnceLock::new()).collect(),
373            proxy,
374        });
375    }
376
377    /// Configure health checks before starting exchanges.
378    pub fn set_health_config(&mut self, mut config: HealthConfig) {
379        assert_eq!(config.address, self.address);
380        if config.source_address.is_none() {
381            config.source_address = self.source_address;
382        }
383        if let Some(metrics) = &self.metrics {
384            config.metrics = Some(metrics.proxy.clone());
385        }
386        self.health = Health::with_config(config);
387    }
388
389    #[must_use]
390    pub fn down(&self) -> bool {
391        self.health.down()
392    }
393    pub async fn exchange(
394        &self,
395        query: &[u8],
396        response: &mut [u8],
397        transport: Transport,
398        options: Options,
399    ) -> Result<Exchange, Error> {
400        if self.stopped.load(Ordering::Acquire) {
401            return Err(Error::Stopped);
402        }
403        self.sweep();
404        let selected = if options.force_tcp {
405            Transport::Tcp
406        } else if options.prefer_udp {
407            Transport::Udp
408        } else {
409            transport
410        };
411        let result = self.exchange_once(query, response, selected).await;
412        // RFC 2181 §9: a client can retry a truncated response over TCP.
413        let result = match result {
414            Ok(_result)
415                if selected == Transport::Udp
416                    && options.prefer_udp
417                    && !options.force_tcp
418                    && response[2] & 0x02 != 0 =>
419            {
420                self.exchange_once(query, response, Transport::Tcp).await
421            }
422            other => other,
423        };
424        if matches!(result, Err(Error::Io(_))) {
425            self.health.start();
426        }
427        result
428    }
429
430    async fn exchange_once(
431        &self,
432        query: &[u8],
433        response: &mut [u8],
434        transport: Transport,
435    ) -> Result<Exchange, Error> {
436        let start = Instant::now();
437        let request = Message::parse(query).map_err(|_| Error::InvalidRequest)?;
438        if query.len() > response.len() || query.len() > u16::MAX as usize {
439            return Err(Error::InvalidRequest);
440        }
441        let protocol = usize::from(transport == Transport::Tcp);
442        let shard = SHARD.with(|index| *index);
443        let cached = {
444            let mut queue = self.shards[protocol][shard].lock().unwrap();
445            loop {
446                match queue.pop_front() {
447                    Some(cached) => {
448                        self.idle[protocol].fetch_sub(1, Ordering::Relaxed);
449                        if cached.used.elapsed() <= self.expire
450                            && (self.max_age.is_zero() || cached.created.elapsed() <= self.max_age)
451                        {
452                            break Some(cached);
453                        }
454                    }
455                    None => break None,
456                }
457            }
458        };
459        let reused = cached.is_some();
460        if let Some(metrics) = &self.metrics {
461            if reused {
462                metrics.cache[protocol].hit();
463            } else {
464                metrics.cache[protocol].miss();
465            }
466        }
467        // RFC 7766 §6.2.1: DNS clients should reuse persistent TCP connections.
468        let (mut socket, created) = if let Some(cached) = cached {
469            (cached.socket, cached.created)
470        } else {
471            (self.dial(transport).await?, Instant::now())
472        };
473        let mut id = [0; 2];
474        // RFC 8945 §5.5: forward messages with unshared TSIG keys unchanged.
475        if request
476            .additionals()
477            .any(|record| record.rtype == Type::TSIG)
478        {
479            id.copy_from_slice(&query[..2]);
480        } else {
481            // RFC 5452 §9.2: outgoing queries use unpredictable IDs.
482            getrandom::getrandom(&mut id)
483                .map_err(|_| Error::Io(io::Error::other("random ID generation failed")))?;
484        }
485        response[..query.len()].copy_from_slice(query);
486        response[..2].copy_from_slice(&id);
487        let local_addr = socket.local_addr()?;
488        let sent = timeout(READ_TIMEOUT, socket.send(&response[..query.len()]))
489            .await
490            .map_err(|_| Error::Io(io::ErrorKind::TimedOut.into()))?;
491        if let Err(err) = sent {
492            return Err(if reused && err.kind() == io::ErrorKind::UnexpectedEof {
493                Error::CachedClosed
494            } else {
495                Error::Io(err)
496            });
497        }
498        let received = timeout(self.read_timeout, async {
499            loop {
500                let len = socket.receive(response).await?;
501                if len < 12 && transport == Transport::Tcp {
502                    return Err(io::ErrorKind::InvalidData.into());
503                }
504                // RFC 5452 §9.1: match response IDs to outstanding queries.
505                if len >= 12 && response[..2] == id {
506                    return Ok::<_, io::Error>(len);
507                }
508            }
509        })
510        .await
511        .map_err(|_| Error::Io(io::ErrorKind::TimedOut.into()))?;
512        let mut length = received.map_err(|err| {
513            if reused && err.kind() == io::ErrorKind::UnexpectedEof {
514                Error::CachedClosed
515            } else {
516                Error::Io(err)
517            }
518        })?;
519        if let Err(error) = Message::parse(&response[..length]) {
520            if transport != Transport::Udp || !is_overflow(&error) {
521                return Err(Error::Io(io::ErrorKind::InvalidData.into()));
522            }
523            let counts: [u8; 6] = response[6..12].try_into().unwrap();
524            response[6..12].fill(0);
525            let question_matches = Message::parse(&response[..length]).is_ok_and(|answer| {
526                answer.qr()
527                    && request.question_count() == answer.question_count()
528                    && request.questions().zip(answer.questions()).all(|(a, b)| {
529                        a.name == b.name && a.qtype == b.qtype && a.qclass == b.qclass
530                    })
531            });
532            response[6..12].copy_from_slice(&counts);
533            if !question_matches {
534                return Err(Error::Io(io::ErrorKind::InvalidData.into()));
535            }
536            length = 12
537                + request
538                    .questions()
539                    .map(|q| q.name.wire_len() + 4)
540                    .sum::<usize>();
541            let flags = [response[2], response[3]];
542            response[..length].copy_from_slice(&query[..length]);
543            // RFC 5625 §4.4: a proxy that truncates a response must set TC.
544            response[2] = flags[0] | 0x02;
545            response[3] = flags[1];
546            response[6..12].fill(0);
547        }
548        let answer = Message::parse(&response[..length])
549            .map_err(|_| Error::Io(io::ErrorKind::InvalidData.into()))?;
550        if let Some(metrics) = &self.metrics {
551            let code = answer.rcode().0;
552            // shortcut: first observation of an RCODE allocates; pre-register codes if cold traffic matters.
553            metrics.durations[usize::from(code)]
554                .get_or_init(|| metrics.proxy.request_duration(&Rcode(code).to_string()))
555                .observe(start.elapsed().as_secs_f64());
556        }
557        // RFC 5452 §9.1: response names, classes, and types must match the query.
558        let same_questions = request.question_count() == answer.question_count()
559            && request
560                .questions()
561                .zip(answer.questions())
562                .all(|(a, b)| a.name == b.name && a.qtype == b.qtype && a.qclass == b.qclass);
563        if !answer.qr() || !same_questions {
564            return Err(Error::Mismatch);
565        }
566        // RFC 1035 §4.1.1: the response ID matches the original query ID.
567        response[..2].copy_from_slice(&query[..2]);
568        {
569            let mut queue = self.shards[protocol][shard].lock().unwrap();
570            if !self.stopped.load(Ordering::Acquire)
571                && self.idle[protocol]
572                    .try_update(Ordering::Relaxed, Ordering::Relaxed, |count| {
573                        (self.max_idle_conns == 0 || count < self.max_idle_conns).then(|| count + 1)
574                    })
575                    .is_ok()
576            {
577                queue.push_back(Cached {
578                    socket,
579                    created,
580                    used: Instant::now(),
581                });
582            }
583        }
584        Ok(Exchange {
585            len: length,
586            local_addr,
587            transport,
588        })
589    }
590}
591
592impl Drop for Proxy {
593    fn drop(&mut self) {
594        self.stop();
595    }
596}
597
598#[cfg(test)]
599mod tests {
600    use super::*;
601
602    #[cfg(target_os = "linux")]
603    #[tokio::test]
604    async fn udp_sockets_enable_error_queues_for_both_families() {
605        use nix::sys::socket::{
606            getsockopt,
607            sockopt::{Ipv4RecvErr, Ipv6RecvErr},
608        };
609        let v4 = UdpSocket::bind("127.0.0.1:0").await.unwrap();
610        enable_recv_err(&v4, v4.local_addr().unwrap());
611        assert!(getsockopt(&v4, Ipv4RecvErr).unwrap());
612        let v6 = UdpSocket::bind("[::1]:0").await.unwrap();
613        enable_recv_err(&v6, v6.local_addr().unwrap());
614        assert!(getsockopt(&v6, Ipv6RecvErr).unwrap());
615    }
616
617    #[test]
618    fn adaptive_dial_timeout_uses_weighted_average_and_bounds() {
619        let proxy = Proxy::new("127.0.0.1:53".parse().unwrap());
620        assert_eq!(proxy.dial_timeout(), Duration::from_secs(30));
621        proxy.update_dial_timeout(Duration::from_secs(3));
622        assert_eq!(proxy.dial_timeout(), Duration::from_secs(24));
623        for _ in 0..40 {
624            proxy.update_dial_timeout(Duration::ZERO);
625        }
626        assert_eq!(proxy.dial_timeout(), Duration::from_secs(1));
627        proxy.update_dial_timeout(Duration::from_secs(120));
628        assert_eq!(proxy.dial_timeout(), Duration::from_secs(30));
629    }
630
631    #[test]
632    fn only_truncated_message_and_name_errors_are_overflows() {
633        for (error, expected) in [
634            (MessageError::Truncated, true),
635            (MessageError::Name(NameError::Truncated), true),
636            (MessageError::Name(NameError::InvalidLabelType), false),
637            (MessageError::Name(NameError::PointerNotBackwards), false),
638            (MessageError::InvalidOpt, false),
639        ] {
640            assert_eq!(is_overflow(&error), expected, "{error:?}");
641        }
642    }
643}
pharos_plugins/proxy/mod.rs

pharos_plugins/proxy/
mod.rs
