pharos_server/
lib.rs

1#![forbid(unsafe_code)]
2//! Plugin Chain contract. A Plugin returns an Outcome rather than writing to a
3//! Listener; earlier Plugins may inspect and modify the caller-owned Response
4//! buffer after awaiting Next. The Listener handles unanswered codes and errors.
5
6use std::{
7    any::Any,
8    collections::{HashMap, HashSet},
9    fmt,
10    future::Future,
11    io,
12    net::{Ipv4Addr, SocketAddr},
13    pin::Pin,
14    sync::Arc,
15    time::Duration,
16};
17
18use arc_swap::ArcSwap;
19use pharos_proto::{BuildError, Message, Name, Rcode, Type};
20use tokio::{
21    net::{TcpListener, UdpSocket},
22    sync::watch,
23    task::JoinHandle,
24};
25
26mod admission;
27pub mod http;
28mod lifecycle;
29pub mod logging;
30pub mod metrics;
31mod outcome;
32mod panic;
33mod process_metrics;
34mod quota;
35pub use admission::{Admission, admit};
36pub use lifecycle::{Lifecycle, Phase};
37pub use outcome::{finish_outcome, response_limit};
38pub use panic::{contain_panic, panics_total};
39mod router;
40mod server;
41mod startup;
42pub use router::{Router, ServerBlock, ViewFilter};
43pub use server::{Chain, Server};
44pub use startup::startup_domains;
45mod tcp;
46mod udp;
47#[cfg(target_os = "linux")]
48mod udp_linux;
49
50/// A statically dispatched continuation of the Plugin Chain.
51pub trait Next {
52    fn run(self, cx: &mut Context<'_>) -> impl Future<Output = Outcome> + Send;
53}
54
55/// A Plugin may await Next and inspect or modify its outcome and Response.
56pub trait Plugin {
57    fn serve<N: Next + Send>(
58        &self,
59        cx: &mut Context<'_>,
60        next: N,
61    ) -> impl Future<Output = Outcome> + Send;
62
63    /// Return a readiness check when this Plugin participates in readiness.
64    fn readiness(&self) -> Option<&dyn Readiness> {
65        None
66    }
67}
68
69/// A Plugin's readiness, evaluated outside the Query hot path.
70pub trait Readiness: Send + Sync {
71    fn ready(&self) -> bool;
72}
73
74/// Transport named by a Server Block Key.
75#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
76pub enum Transport {
77    Dns,
78    Tls,
79    Quic,
80    Grpc,
81    Https,
82    Https3,
83}
84
85/// Protocol used by a Query.
86#[derive(Clone, Copy, Debug, Eq, PartialEq)]
87pub enum Protocol {
88    Udp,
89    Tcp,
90}
91
92/// A placeholder for a future Zone Transfer stream.
93#[derive(Clone, Copy, Debug, Eq, PartialEq)]
94pub struct Stream;
95
96/// Explicit result of running a Plugin Chain.
97#[derive(Clone, Copy, Debug, Eq, PartialEq)]
98pub enum Outcome {
99    Response,
100    Stream(Stream),
101    Drop,
102    Unanswered(Rcode),
103    Error(Error),
104}
105
106/// Errors for the Listener to log and answer with SERVFAIL.
107#[derive(Clone, Copy, Debug, Eq, PartialEq)]
108pub enum Error {
109    Build(BuildError),
110    InvalidQuery,
111    ResponseBufferTooSmall,
112    NotImplemented,
113    NoNextPlugin,
114    PluginPanic,
115    SelfLookupDepth,
116    SelfLookupUnavailable,
117    SelfLookupDropped,
118}
119
120impl fmt::Display for Error {
121    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
122        match self {
123            Self::Build(error) => write!(f, "{error}"),
124            Self::InvalidQuery => f.write_str("invalid Query"),
125            Self::ResponseBufferTooSmall => f.write_str("Response buffer too small"),
126            Self::NotImplemented => f.write_str("plugin is not implemented"),
127            Self::NoNextPlugin => f.write_str("no next plugin found"),
128            Self::PluginPanic => f.write_str("plugin panicked"),
129            Self::SelfLookupDepth => f.write_str("Self Lookup depth exceeded"),
130            Self::SelfLookupUnavailable => f.write_str("no Server available for Self Lookup"),
131            Self::SelfLookupDropped => f.write_str("Self Lookup produced no Response"),
132        }
133    }
134}
135
136impl std::error::Error for Error {}
137
138/// Caller-provided, fixed-capacity Metadata slots.
139pub struct Metadata<'a> {
140    slots: &'a mut [Option<(&'a str, &'a str)>],
141    inline_name: Option<&'static str>,
142    inline_value: [u8; 64],
143    inline_len: usize,
144    // shortcut: 12 copied labels and 1 KiB suffice for Kubernetes; expand for more providers.
145    copied: [Option<(&'static str, usize, usize)>; 12],
146    copied_data: [u8; 1024],
147    copied_len: usize,
148    source: Option<Arc<dyn MetadataSource>>,
149    enabled: bool,
150    self_lookup: Option<SelfLookupState<'a>>,
151}
152
153impl<'a> Metadata<'a> {
154    pub fn new(slots: &'a mut [Option<(&'a str, &'a str)>]) -> Self {
155        Self {
156            slots,
157            inline_name: None,
158            inline_value: [0; 64],
159            inline_len: 0,
160            copied: [None; 12],
161            copied_data: [0; 1024],
162            copied_len: 0,
163            source: None,
164            enabled: true,
165            self_lookup: None,
166        }
167    }
168
169    /// Retain an existing shared source of Metadata without allocating per Query.
170    pub fn set_source(&mut self, source: Arc<dyn MetadataSource>) {
171        self.source = Some(source);
172    }
173
174    /// Enable public Plugin Metadata for the current Server Block.
175    pub fn set_enabled(&mut self, enabled: bool) {
176        self.enabled = enabled;
177    }
178
179    #[must_use]
180    pub fn enabled(&self) -> bool {
181        self.enabled
182    }
183
184    /// Construct fresh Metadata while retaining detached Self Lookup routing.
185    pub fn with_self_lookup(
186        slots: &'a mut [Option<(&'a str, &'a str)>],
187        handle: Option<&'a SelfLookupHandle>,
188    ) -> Self {
189        let mut metadata = Self::new(slots);
190        metadata.self_lookup = handle.map(|handle| SelfLookupState {
191            server: handle,
192            depth: handle.depth,
193        });
194        metadata
195    }
196
197    /// Insert or replace a named value; returns false when no slot is free.
198    pub fn insert(&mut self, name: &'a str, value: &'a str) -> bool {
199        let slot = self
200            .slots
201            .iter()
202            .position(|slot| slot.is_some_and(|(existing, _)| existing == name))
203            .or_else(|| self.slots.iter().position(Option::is_none));
204        let Some(slot) = slot else { return false };
205        self.remove_copied(name);
206        if self.inline_name == Some(name) {
207            self.inline_name = None;
208        }
209        self.slots[slot] = Some((name, value));
210        true
211    }
212
213    /// Copy a short value whose lifetime is shorter than the Query's.
214    pub fn insert_inline(&mut self, name: &'static str, value: &str) -> bool {
215        if value.len() > self.inline_value.len()
216            || self.inline_name.is_some_and(|existing| existing != name)
217        {
218            return false;
219        }
220        self.remove_copied(name);
221        for slot in self.slots.iter_mut() {
222            if slot.is_some_and(|(key, _)| key == name) {
223                *slot = None;
224            }
225        }
226        // shortcut: one copied value suffices for forward; add slots for other producers.
227        self.inline_name = Some(name);
228        self.inline_len = value.len();
229        self.inline_value[..value.len()].copy_from_slice(value.as_bytes());
230        true
231    }
232
233    /// Copy a value from a temporary Plugin buffer without allocating.
234    pub fn insert_copied(&mut self, name: &'static str, value: &str) -> bool {
235        let previous = self
236            .copied
237            .iter()
238            .flatten()
239            .find(|(key, _, _)| *key == name)
240            .map_or(0, |(_, start, end)| end - start);
241        if value.len() > self.copied_data.len() - (self.copied_len - previous) {
242            return false;
243        }
244        let Some(index) = self
245            .copied
246            .iter()
247            .position(|slot| slot.is_none_or(|(key, _, _)| key == name))
248        else {
249            return false;
250        };
251        self.remove_copied(name);
252        if self.inline_name == Some(name) {
253            self.inline_name = None;
254        }
255        for slot in self.slots.iter_mut() {
256            if slot.is_some_and(|(key, _)| key == name) {
257                *slot = None;
258            }
259        }
260        let end = self.copied_len + value.len();
261        self.copied_data[self.copied_len..end].copy_from_slice(value.as_bytes());
262        self.copied[index] = Some((name, self.copied_len, end));
263        self.copied_len = end;
264        true
265    }
266
267    fn remove_copied(&mut self, name: &str) {
268        let Some(index) = self
269            .copied
270            .iter()
271            .position(|slot| slot.is_some_and(|(key, _, _)| key == name))
272        else {
273            return;
274        };
275        let (_, start, end) = self.copied[index].take().unwrap();
276        self.copied_data.copy_within(end..self.copied_len, start);
277        self.copied_len -= end - start;
278        for (_, other_start, other_end) in self.copied.iter_mut().flatten() {
279            if *other_start >= end {
280                *other_start -= end - start;
281                *other_end -= end - start;
282            }
283        }
284    }
285
286    #[must_use]
287    pub fn get(&self, name: &str) -> Option<&str> {
288        if let Some((_, start, end)) = self
289            .copied
290            .iter()
291            .flatten()
292            .find(|(key, _, _)| *key == name)
293        {
294            return std::str::from_utf8(&self.copied_data[*start..*end]).ok();
295        }
296        if self.inline_name == Some(name) {
297            return std::str::from_utf8(&self.inline_value[..self.inline_len]).ok();
298        }
299        self.slots
300            .iter()
301            .flatten()
302            .find_map(|(key, value)| (*key == name).then_some(*value))
303            .or_else(|| self.source.as_ref().and_then(|source| source.get(name)))
304    }
305}
306
307/// A shared source whose values remain valid for a Query.
308pub trait MetadataSource: Send + Sync {
309    fn get(&self, label: &str) -> Option<&str>;
310}
311
312/// Borrowed Query state and reusable, preallocated Response buffer.
313pub struct Context<'a> {
314    pub query: Message<'a>,
315    pub client: SocketAddr,
316    /// Listener address; wildcard addresses are not resolved per Query.
317    pub local: SocketAddr,
318    pub transport: Transport,
319    pub protocol: Protocol,
320    pub domain: &'a Name,
321    pub metadata: Metadata<'a>,
322    pub response: &'a mut Vec<u8>,
323    pub max_response_size: usize,
324}
325
326/// The Listener state and depth associated with a Query.
327struct SelfLookupState<'a> {
328    server: &'a dyn SelfLookupServer,
329    depth: u8,
330}
331
332// Type erasure keeps Context nongeneric without allocating on ordinary Queries.
333trait SelfLookupServer: Sync {
334    fn lookup<'a>(
335        &'a self,
336        wire: &'a [u8],
337        peer: SocketAddr,
338        protocol: Protocol,
339        depth: u8,
340        response: &'a mut Vec<u8>,
341    ) -> Pin<Box<dyn Future<Output = Action> + Send + 'a>>;
342
343    fn owned(&self) -> Option<Arc<dyn OwnedSelfLookupServer>> {
344        None
345    }
346}
347
348trait OwnedSelfLookupServer: Send + Sync {
349    fn lookup<'a>(
350        self: Arc<Self>,
351        wire: &'a [u8],
352        peer: SocketAddr,
353        protocol: Protocol,
354        depth: u8,
355        response: &'a mut Vec<u8>,
356    ) -> Pin<Box<dyn Future<Output = Action> + Send + 'a>>;
357}
358
359/// Owned Listener state for a detached Query, including its Self Lookup depth.
360pub struct SelfLookupHandle {
361    server: Arc<dyn OwnedSelfLookupServer>,
362    depth: u8,
363}
364
365impl SelfLookupServer for SelfLookupHandle {
366    fn lookup<'a>(
367        &'a self,
368        wire: &'a [u8],
369        peer: SocketAddr,
370        protocol: Protocol,
371        depth: u8,
372        response: &'a mut Vec<u8>,
373    ) -> Pin<Box<dyn Future<Output = Action> + Send + 'a>> {
374        Arc::clone(&self.server).lookup(wire, peer, protocol, depth, response)
375    }
376
377    fn owned(&self) -> Option<Arc<dyn OwnedSelfLookupServer>> {
378        Some(Arc::clone(&self.server))
379    }
380}
381
382impl Context<'_> {
383    /// Retain this Query's Listener state for a detached task, if available.
384    pub fn self_lookup_handle(&self) -> Option<SelfLookupHandle> {
385        let state = self.metadata.self_lookup.as_ref()?;
386        Some(SelfLookupHandle {
387            server: state.server.owned()?,
388            depth: state.depth,
389        })
390    }
391
392    /// Route a new Query through this Query's Listener and capture its Response.
393    pub async fn self_lookup(&self, name: &Name, qtype: Type) -> Result<Vec<u8>, Error> {
394        let state = self
395            .metadata
396            .self_lookup
397            .as_ref()
398            .ok_or(Error::SelfLookupUnavailable)?;
399        if state.depth >= 10 {
400            return Err(Error::SelfLookupDepth);
401        }
402        let query = self
403            .query
404            .with_question(name, qtype)
405            .map_err(Error::Build)?;
406        let mut response = Vec::with_capacity(usize::from(u16::MAX));
407        match state
408            .server
409            .lookup(
410                &query,
411                self.client,
412                self.protocol,
413                state.depth + 1,
414                &mut response,
415            )
416            .await
417        {
418            Action::Respond => Ok(response),
419            Action::Drop => Err(Error::SelfLookupDropped),
420        }
421    }
422}
423
424#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
425pub struct ListenerKey {
426    pub transport: Transport,
427    pub address: SocketAddr,
428}
429
430fn canonical_peer(mut peer: SocketAddr) -> SocketAddr {
431    let ip = peer.ip().to_canonical();
432    if ip != peer.ip() {
433        peer.set_ip(ip);
434    }
435    peer
436}
437
438#[derive(Clone, Copy, Debug, Eq, PartialEq)]
439pub enum Action {
440    Respond,
441    Drop,
442}
443
444/// Async, statically dispatched Query handling. Clear and reuse the supplied
445/// Response buffer; return `Drop` to suppress a Response.
446/// Handlers must not block the runtime; TCP Queries on a connection may run
447/// concurrently.
448pub trait Handler: Send + Sync + 'static {
449    /// Number of UDP worker sockets for this Listener.
450    fn udp_workers(&self) -> usize {
451        quota::worker_count()
452    }
453    /// Whether this HINFO Query must run concurrently with UDP reception.
454    fn parallel_query(&self, _query: &[u8], _peer: SocketAddr) -> Option<watch::Receiver<bool>> {
455        None
456    }
457    fn handle(
458        &self,
459        query: &[u8],
460        peer: SocketAddr,
461        protocol: Protocol,
462        response: &mut Vec<u8>,
463    ) -> impl Future<Output = Action> + Send;
464
465    /// Handle a Query while retaining access to the Listener's owned state.
466    fn handle_shared(
467        self: &Arc<Self>,
468        query: &[u8],
469        peer: SocketAddr,
470        protocol: Protocol,
471        response: &mut Vec<u8>,
472    ) -> impl Future<Output = Action> + Send {
473        self.handle(query, peer, protocol, response)
474    }
475}
476
477/// A DNS Listener serving UDP and TCP on the same address.
478pub struct Listener {
479    address: SocketAddr,
480    stop: watch::Sender<bool>,
481    tasks: Vec<JoinHandle<io::Result<()>>>,
482    udp_stop: watch::Sender<bool>,
483    udp_tasks: Vec<JoinHandle<io::Result<()>>>,
484}
485
486impl Listener {
487    pub async fn bind<H: Handler>(address: SocketAddr, handler: H) -> io::Result<Self> {
488        let workers = handler.udp_workers();
489        Self::bind_with_ipv6(address, handler, workers, bind_ipv6_dual, udp::bind).await
490    }
491
492    async fn bind_with_ipv6<H: Handler>(
493        address: SocketAddr,
494        handler: H,
495        workers: usize,
496        bind_ipv6: fn(SocketAddr) -> io::Result<TcpListener>,
497        bind_udp: fn(SocketAddr) -> io::Result<std::net::UdpSocket>,
498    ) -> io::Result<Self> {
499        let wildcard = address.is_ipv6() && address.ip().is_unspecified();
500        let mut ipv4_only = false;
501        // Port 0 picks a free TCP port; the same UDP port may be taken, so retry.
502        let mut attempts = if address.port() == 0 { 16 } else { 1 };
503        // RFC 7766 §5: general-purpose DNS implementations support both UDP and TCP.
504        let (tcp, address, sockets) = loop {
505            let tcp = if wildcard && !ipv4_only {
506                match optional_ipv6(bind_ipv6(address)) {
507                    Ok(Some(tcp)) => tcp,
508                    Ok(None) => {
509                        ipv4_only = true;
510                        continue;
511                    }
512                    Err(error) if error.kind() == io::ErrorKind::AddrInUse && attempts > 1 => {
513                        attempts -= 1;
514                        continue;
515                    }
516                    Err(error) => return Err(error),
517                }
518            } else if ipv4_only {
519                TcpListener::bind((Ipv4Addr::UNSPECIFIED, address.port())).await?
520            } else {
521                TcpListener::bind(address).await?
522            };
523            let bound = tcp.local_addr()?;
524            let sockets = (0..workers)
525                .map(|_| bind_udp(bound))
526                .collect::<io::Result<Vec<_>>>();
527            match sockets {
528                Ok(sockets) => break (tcp, bound, sockets),
529                Err(error) if wildcard && bound.is_ipv6() && ipv6_unavailable(&error) => {
530                    ipv4_only = true;
531                }
532                Err(error) if error.kind() == io::ErrorKind::AddrInUse && attempts > 1 => {
533                    attempts -= 1;
534                }
535                Err(error) => return Err(error),
536            }
537        };
538        let (stop, receiver) = watch::channel(false);
539        let (udp_stop, udp_receiver) = watch::channel(false);
540        let handler = Arc::new(handler);
541        let udp_tasks = sockets
542            .into_iter()
543            .map(|socket| {
544                let socket = UdpSocket::from_std(socket)?;
545                Ok(udp::spawn(
546                    socket,
547                    Arc::clone(&handler),
548                    udp_receiver.clone(),
549                ))
550            })
551            .collect::<io::Result<Vec<_>>>()?;
552        let tasks = vec![tcp::spawn(tcp, handler, receiver)];
553        Ok(Self {
554            address,
555            stop,
556            tasks,
557            udp_stop,
558            udp_tasks,
559        })
560    }
561
562    pub fn local_addr(&self) -> SocketAddr {
563        self.address
564    }
565
566    /// Stop accepting Queries and drain for at most five seconds.
567    pub async fn shutdown(mut self) -> io::Result<()> {
568        self.stop.send_replace(true);
569        self.udp_stop.send_replace(true);
570        self.tasks.append(&mut self.udp_tasks);
571        let mut result = Ok(());
572        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
573        while let Some(mut task) = self.tasks.pop() {
574            match tokio::time::timeout_at(deadline, &mut task).await {
575                Ok(next) => {
576                    let next = next.map_err(io::Error::other).and_then(|v| v);
577                    if result.is_ok() {
578                        result = next;
579                    }
580                }
581                Err(_) => {
582                    task.abort();
583                    for task in &self.tasks {
584                        task.abort();
585                    }
586                    let _ = task.await;
587                    for task in self.tasks.drain(..) {
588                        let _ = task.await;
589                    }
590                    break;
591                }
592            }
593        }
594        result
595    }
596}
597
598fn bind_ipv6_dual(address: SocketAddr) -> io::Result<TcpListener> {
599    let socket = socket2::Socket::new(
600        socket2::Domain::IPV6,
601        socket2::Type::STREAM,
602        Some(socket2::Protocol::TCP),
603    )?;
604    socket.set_only_v6(false)?;
605    socket.set_reuse_address(true)?;
606    socket.set_nonblocking(true)?;
607    socket.bind(&address.into())?;
608    socket.listen(1024)?;
609    TcpListener::from_std(socket.into())
610}
611
612fn optional_ipv6<T>(result: io::Result<T>) -> io::Result<Option<T>> {
613    match result {
614        Ok(value) => Ok(Some(value)),
615        Err(error) if ipv6_unavailable(&error) => Ok(None),
616        Err(error) => Err(error),
617    }
618}
619
620fn ipv6_unavailable(error: &io::Error) -> bool {
621    error.kind() == io::ErrorKind::AddrNotAvailable
622        || error.raw_os_error() == Some(libc::EAFNOSUPPORT)
623}
624
625#[cfg(test)]
626mod wildcard_tests {
627    use super::{
628        Action, Handler, Listener, Protocol, bind_ipv6_dual, canonical_peer, optional_ipv6, udp,
629    };
630    use std::{
631        io,
632        net::{Ipv6Addr, SocketAddr, SocketAddrV6},
633    };
634
635    #[test]
636    fn canonical_peer_preserves_native_ipv6_scope() {
637        let scoped = SocketAddr::V6(SocketAddrV6::new(
638            "fe80::1".parse::<Ipv6Addr>().unwrap(),
639            53,
640            7,
641            3,
642        ));
643        assert_eq!(canonical_peer(scoped), scoped);
644        assert_eq!(
645            canonical_peer("[::ffff:127.0.0.1]:53".parse().unwrap()),
646            "127.0.0.1:53".parse().unwrap()
647        );
648    }
649
650    #[test]
651    fn unavailable_ipv6_falls_back_but_other_errors_remain_errors() {
652        assert_eq!(
653            optional_ipv6::<()>(Err(io::Error::from_raw_os_error(libc::EAFNOSUPPORT))).unwrap(),
654            None
655        );
656        assert_eq!(
657            optional_ipv6::<()>(Err(io::Error::from(io::ErrorKind::AddrNotAvailable))).unwrap(),
658            None
659        );
660        assert_eq!(optional_ipv6(Ok(7)).unwrap(), Some(7));
661        assert!(
662            optional_ipv6::<()>(Err(io::Error::new(io::ErrorKind::AddrInUse, "occupied"))).is_err()
663        );
664    }
665
666    struct Echo;
667    impl Handler for Echo {
668        async fn handle(
669            &self,
670            query: &[u8],
671            _: SocketAddr,
672            _: Protocol,
673            response: &mut Vec<u8>,
674        ) -> Action {
675            response.extend_from_slice(query);
676            Action::Respond
677        }
678    }
679
680    #[tokio::test]
681    async fn listener_binds_the_requested_number_of_udp_sockets() {
682        use std::sync::atomic::{AtomicUsize, Ordering};
683        static BINDS: AtomicUsize = AtomicUsize::new(0);
684        fn counting_bind(address: SocketAddr) -> io::Result<std::net::UdpSocket> {
685            BINDS.fetch_add(1, Ordering::SeqCst);
686            udp::bind(address)
687        }
688        let listener = Listener::bind_with_ipv6(
689            "127.0.0.1:0".parse().unwrap(),
690            Echo,
691            3,
692            bind_ipv6_dual,
693            counting_bind,
694        )
695        .await
696        .unwrap();
697        assert_eq!(BINDS.load(Ordering::SeqCst), 3);
698        listener.shutdown().await.unwrap();
699    }
700
701    #[tokio::test]
702    async fn reload_replaces_only_changed_udp_workers_and_keeps_tcp() {
703        use std::collections::HashMap;
704        use tokio::io::{AsyncReadExt, AsyncWriteExt};
705        struct Counted(usize, u8);
706        impl Handler for Counted {
707            fn udp_workers(&self) -> usize {
708                self.0
709            }
710            async fn handle(
711                &self,
712                _: &[u8],
713                _: SocketAddr,
714                _: Protocol,
715                response: &mut Vec<u8>,
716            ) -> Action {
717                response.push(self.1);
718                Action::Respond
719            }
720        }
721        let key = super::ListenerKey {
722            transport: super::Transport::Dns,
723            address: "127.0.0.1:0".parse().unwrap(),
724        };
725        let mut registry = super::ListenerRegistry::default();
726        registry
727            .apply(HashMap::from([(key, Counted(2, 1))]))
728            .await
729            .unwrap();
730        let listener = registry.get(&key).unwrap();
731        let address = listener.local_addr();
732        let ids: Vec<_> = listener
733            .udp_tasks
734            .iter()
735            .map(tokio::task::JoinHandle::id)
736            .collect();
737        assert_eq!(ids.len(), 2);
738        let mut tcp = tokio::net::TcpStream::connect(address).await.unwrap();
739        registry
740            .apply(HashMap::from([(key, Counted(2, 2))]))
741            .await
742            .unwrap();
743        assert_eq!(
744            registry
745                .get(&key)
746                .unwrap()
747                .udp_tasks
748                .iter()
749                .map(tokio::task::JoinHandle::id)
750                .collect::<Vec<_>>(),
751            ids
752        );
753        registry
754            .apply(HashMap::from([(key, Counted(3, 3))]))
755            .await
756            .unwrap();
757        let listener = registry.get(&key).unwrap();
758        assert_eq!(listener.local_addr(), address);
759        let new_ids: Vec<_> = listener
760            .udp_tasks
761            .iter()
762            .map(tokio::task::JoinHandle::id)
763            .collect();
764        assert_eq!(new_ids.len(), 3);
765        assert!(new_ids.iter().all(|id| !ids.contains(id)));
766        tcp.write_all(&[0, 1, 42]).await.unwrap();
767        let mut response = [0; 3];
768        tcp.read_exact(&mut response).await.unwrap();
769        assert_eq!(response, [0, 1, 3]);
770        let udp = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
771        udp.send_to(b"q", address).await.unwrap();
772        let mut byte = [0];
773        tokio::time::timeout(std::time::Duration::from_secs(2), udp.recv(&mut byte))
774            .await
775            .unwrap()
776            .unwrap();
777        assert_eq!(byte, [3]);
778        tokio::time::timeout(std::time::Duration::from_secs(2), async {
779            while registry
780                .get(&key)
781                .unwrap()
782                .tasks
783                .iter()
784                .skip(1)
785                .any(|task| !task.is_finished())
786            {
787                tokio::task::yield_now().await;
788            }
789        })
790        .await
791        .unwrap();
792        registry
793            .apply(HashMap::from([(key, Counted(1, 4))]))
794            .await
795            .unwrap();
796        let listener = registry.get(&key).unwrap();
797        assert_eq!(listener.udp_tasks.len(), 1);
798        assert_eq!(
799            listener.tasks.len(),
800            4,
801            "completed workers should be reaped"
802        );
803        drop(tcp);
804        registry.shutdown().await.unwrap();
805    }
806
807    async fn assert_ipv4_fallback(listener: Listener) {
808        use tokio::io::{AsyncReadExt, AsyncWriteExt};
809        assert!(listener.local_addr().is_ipv4());
810        let socket = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
811        socket
812            .send_to(b"query", ("127.0.0.1", listener.local_addr().port()))
813            .await
814            .unwrap();
815        let mut response = [0; 16];
816        let (len, _) = tokio::time::timeout(
817            std::time::Duration::from_secs(2),
818            socket.recv_from(&mut response),
819        )
820        .await
821        .unwrap()
822        .unwrap();
823        assert_eq!(&response[..len], b"query");
824        let mut tcp = tokio::net::TcpStream::connect(("127.0.0.1", listener.local_addr().port()))
825            .await
826            .unwrap();
827        tcp.write_u16(5).await.unwrap();
828        tcp.write_all(b"query").await.unwrap();
829        let len = tokio::time::timeout(std::time::Duration::from_secs(2), tcp.read_u16())
830            .await
831            .unwrap()
832            .unwrap();
833        assert_eq!(len, 5);
834        let mut response = [0; 5];
835        tokio::time::timeout(
836            std::time::Duration::from_secs(2),
837            tcp.read_exact(&mut response),
838        )
839        .await
840        .unwrap()
841        .unwrap();
842        assert_eq!(&response, b"query");
843        drop(tcp);
844        listener.shutdown().await.unwrap();
845    }
846
847    #[tokio::test]
848    async fn wildcard_falls_back_to_ipv4_when_ipv6_tcp_is_unavailable() {
849        fn unavailable(_: SocketAddr) -> io::Result<tokio::net::TcpListener> {
850            Err(io::Error::from_raw_os_error(libc::EAFNOSUPPORT))
851        }
852        let listener =
853            Listener::bind_with_ipv6("[::]:0".parse().unwrap(), Echo, 1, unavailable, udp::bind)
854                .await
855                .unwrap();
856        assert_ipv4_fallback(listener).await;
857    }
858
859    #[tokio::test]
860    async fn wildcard_falls_back_to_ipv4_when_ipv6_udp_is_unavailable() {
861        use std::sync::atomic::{AtomicBool, Ordering};
862
863        static TRIED_IPV6_UDP: AtomicBool = AtomicBool::new(false);
864        if std::net::TcpListener::bind("[::1]:0").is_err() {
865            return;
866        }
867        fn unavailable(address: SocketAddr) -> io::Result<std::net::UdpSocket> {
868            if address.is_ipv6() {
869                TRIED_IPV6_UDP.store(true, Ordering::Relaxed);
870                Err(io::Error::from(io::ErrorKind::AddrNotAvailable))
871            } else {
872                udp::bind(address)
873            }
874        }
875        let listener = Listener::bind_with_ipv6(
876            "[::]:0".parse().unwrap(),
877            Echo,
878            1,
879            bind_ipv6_dual,
880            unavailable,
881        )
882        .await
883        .unwrap();
884        assert!(TRIED_IPV6_UDP.load(Ordering::Relaxed));
885        assert_ipv4_fallback(listener).await;
886    }
887}
888
889impl Drop for Listener {
890    fn drop(&mut self) {
891        self.stop.send_replace(true);
892        self.udp_stop.send_replace(true);
893    }
894}
895
896#[derive(Default)]
897pub struct ListenerRegistry {
898    listeners: HashMap<ListenerKey, Listener>,
899    handlers: HashMap<ListenerKey, Box<dyn Any + Send + Sync>>,
900}
901
902impl ListenerRegistry {
903    pub fn insert(&mut self, key: ListenerKey, listener: Listener) -> Option<Listener> {
904        self.handlers.remove(&key);
905        self.listeners.insert(key, listener)
906    }
907
908    pub fn remove(&mut self, key: &ListenerKey) -> Option<Listener> {
909        self.handlers.remove(key);
910        self.listeners.remove(key)
911    }
912
913    pub fn get(&self, key: &ListenerKey) -> Option<&Listener> {
914        self.listeners.get(key)
915    }
916
917    /// Stop and drain every Listener.
918    pub async fn shutdown(mut self) -> io::Result<()> {
919        let mut result = Ok(());
920        for (_, listener) in self.listeners.drain() {
921            if let Err(error) = listener.shutdown().await
922                && result.is_ok()
923            {
924                result = Err(error);
925            }
926        }
927        result
928    }
929
930    /// Bind new Listeners before changing handlers; return replaced states for draining.
931    pub async fn apply<H: Handler>(
932        &mut self,
933        mut servers: HashMap<ListenerKey, H>,
934    ) -> io::Result<Vec<Arc<H>>> {
935        let desired: HashSet<_> = servers.keys().copied().collect();
936        if desired.iter().any(|key| key.transport != Transport::Dns) {
937            return Err(io::Error::new(
938                io::ErrorKind::InvalidInput,
939                "unsupported Listener transport",
940            ));
941        }
942        for key in servers
943            .keys()
944            .filter(|key| self.listeners.contains_key(key))
945        {
946            if self
947                .handlers
948                .get(key)
949                .and_then(|handler| handler.downcast_ref::<Arc<ArcSwap<H>>>())
950                .is_none()
951            {
952                return Err(io::Error::new(
953                    io::ErrorKind::InvalidInput,
954                    "Listener handler type mismatch",
955                ));
956            }
957        }
958        let mut added: Vec<(ListenerKey, Listener, Arc<ArcSwap<H>>)> = Vec::new();
959        let new_keys: Vec<_> = servers
960            .keys()
961            .filter(|key| !self.listeners.contains_key(key))
962            .copied()
963            .collect();
964        for key in new_keys {
965            let handler = Arc::new(ArcSwap::from_pointee(
966                servers.remove(&key).expect("new Listener"),
967            ));
968            let listener = match Listener::bind(key.address, Swappable(Arc::clone(&handler))).await
969            {
970                Ok(listener) => listener,
971                Err(error) => {
972                    for (_, listener, _) in added {
973                        let _: io::Result<()> = listener.shutdown().await;
974                    }
975                    return Err(io::Error::new(
976                        error.kind(),
977                        format!("{}: {error}", key.address),
978                    ));
979                }
980            };
981            added.push((key, listener, handler));
982        }
983        let mut replacements = Vec::new();
984        for (key, server) in &servers {
985            let listener = &self.listeners[key];
986            let count = server.udp_workers();
987            if count == listener.udp_tasks.len() {
988                continue;
989            }
990            let sockets = (0..count)
991                .map(|_| udp::bind(listener.address).and_then(UdpSocket::from_std))
992                .collect::<io::Result<Vec<_>>>();
993            match sockets {
994                Ok(sockets) => replacements.push((*key, sockets)),
995                Err(error) => {
996                    for (_, listener, _) in added {
997                        let _: io::Result<()> = listener.shutdown().await;
998                    }
999                    return Err(io::Error::new(
1000                        error.kind(),
1001                        format!("{}: {error}", key.address),
1002                    ));
1003                }
1004            }
1005        }
1006        let mut retired = Vec::new();
1007        for (key, server) in servers {
1008            retired.push(
1009                self.handlers[&key]
1010                    .downcast_ref::<Arc<ArcSwap<H>>>()
1011                    .expect("validated handler")
1012                    .swap(Arc::new(server)),
1013            );
1014        }
1015        for (key, sockets) in replacements {
1016            let handler = self.handlers[&key]
1017                .downcast_ref::<Arc<ArcSwap<H>>>()
1018                .expect("validated handler");
1019            let handler = Arc::new(Swappable(Arc::clone(handler)));
1020            let listener = self.listeners.get_mut(&key).expect("existing Listener");
1021            for task in listener.tasks.split_off(1) {
1022                if task.is_finished() {
1023                    if let Err(error) = task.await.map_err(io::Error::other).and_then(|v| v) {
1024                        logging::log("ERROR", "", &format!("UDP worker shutdown: {error}"));
1025                    }
1026                } else {
1027                    listener.tasks.push(task);
1028                }
1029            }
1030            let (stop, receiver) = watch::channel(false);
1031            let tasks = sockets
1032                .into_iter()
1033                .map(|socket| udp::spawn(socket, Arc::clone(&handler), receiver.clone()))
1034                .collect();
1035            listener.udp_stop.send_replace(true);
1036            listener.tasks.append(&mut listener.udp_tasks);
1037            listener.udp_stop = stop;
1038            listener.udp_tasks = tasks;
1039        }
1040        let removed: Vec<_> = self
1041            .listeners
1042            .keys()
1043            .filter(|key| !desired.contains(key))
1044            .copied()
1045            .collect();
1046        for (key, listener, handler) in added {
1047            self.listeners.insert(key, listener);
1048            self.handlers.insert(key, Box::new(handler));
1049        }
1050        for key in removed {
1051            if let Some(handler) = self
1052                .handlers
1053                .get(&key)
1054                .and_then(|handler| handler.downcast_ref::<Arc<ArcSwap<H>>>())
1055            {
1056                retired.push(handler.load_full());
1057            }
1058            if let Some(listener) = self.remove(&key)
1059                && let Err(error) = listener.shutdown().await
1060            {
1061                logging::log("ERROR", "", &format!("Listener shutdown: {error}"));
1062            }
1063        }
1064        Ok(retired)
1065    }
1066}
1067
1068struct Swappable<H>(Arc<ArcSwap<H>>);
1069
1070impl<H: Handler> Handler for Swappable<H> {
1071    fn udp_workers(&self) -> usize {
1072        self.0.load().udp_workers()
1073    }
1074    fn parallel_query(&self, query: &[u8], peer: SocketAddr) -> Option<watch::Receiver<bool>> {
1075        self.0.load().parallel_query(query, peer)
1076    }
1077    async fn handle(
1078        &self,
1079        query: &[u8],
1080        peer: SocketAddr,
1081        protocol: Protocol,
1082        response: &mut Vec<u8>,
1083    ) -> Action {
1084        let handler = self.0.load();
1085        H::handle_shared(&handler, query, peer, protocol, response).await
1086    }
1087}
pharos_server/lib.rs

pharos_server/
lib.rs
