Skip to main content

pharos_plugins/
loop.rs

1//! Detect forwarding loops during startup.
2
3use pharos_proto::{Class, Message, MessageBuilder, Name, NameRef, Type};
4use pharos_server::{Context, Lifecycle, Next, Outcome, Phase, Plugin};
5use std::{
6    io::{self, Write},
7    net::{IpAddr, Ipv4Addr, SocketAddr},
8    sync::{
9        Arc, Mutex,
10        atomic::{AtomicBool, AtomicUsize, Ordering},
11    },
12    time::Duration,
13};
14use tokio::{
15    net::UdpSocket,
16    sync::watch,
17    task::JoinHandle,
18    time::{Instant, timeout_at},
19};
20
21#[derive(Debug)]
22pub struct Loop {
23    pub domain: String,
24    state: Arc<State>,
25}
26
27#[derive(Debug)]
28struct State {
29    domain: String,
30    qname: String,
31    probe: Option<Name>,
32    seen: AtomicUsize,
33    disabled: AtomicBool,
34    complete: AtomicBool,
35    resume: AtomicBool,
36    address: Mutex<SocketAddr>,
37    task: Mutex<Option<JoinHandle<()>>>,
38    cancel: watch::Sender<bool>,
39}
40
41impl Loop {
42    pub(crate) fn parallel_query(&self, query: &Message<'_>) -> Option<watch::Receiver<bool>> {
43        self.is_probe(query).then(|| self.state.cancel.subscribe())
44    }
45    pub(crate) fn is_probe(&self, query: &Message<'_>) -> bool {
46        !self.state.disabled.load(Ordering::Acquire)
47            && query.questions().next().is_some_and(|question| {
48                question.qtype == Type::HINFO
49                    && self
50                        .state
51                        .probe
52                        .as_ref()
53                        .is_some_and(|name| question.name == *name)
54            })
55    }
56    pub fn new(domain: &str) -> Self {
57        let random = || {
58            let mut bytes = [0; 8];
59            if getrandom::getrandom(&mut bytes).is_err() {
60                return std::time::SystemTime::now()
61                    .duration_since(std::time::UNIX_EPOCH)
62                    .unwrap_or_default()
63                    .as_nanos() as u64;
64            }
65            u64::from_ne_bytes(bytes) & ((1 << 62) - 1)
66        };
67        let qname = format!(
68            "{}.{}.{}",
69            random(),
70            random(),
71            domain.trim_start_matches('.')
72        );
73        Self {
74            domain: domain.into(),
75            state: Arc::new(State {
76                domain: domain.into(),
77                probe: qname.parse().ok(),
78                qname,
79                seen: AtomicUsize::new(0),
80                disabled: AtomicBool::new(false),
81                complete: AtomicBool::new(false),
82                resume: AtomicBool::new(false),
83                address: Mutex::new(SocketAddr::from(([127, 0, 0, 1], 0))),
84                task: Mutex::new(None),
85                cancel: watch::channel(false).0,
86            }),
87        }
88    }
89
90    pub(crate) fn register(&self, lifecycle: &mut Lifecycle, mut address: SocketAddr) {
91        if address.ip().is_unspecified() {
92            address.set_ip(IpAddr::V4(Ipv4Addr::LOCALHOST));
93        }
94        *self.state.address.lock().unwrap() = address;
95        let state = Arc::clone(&self.state);
96        lifecycle.register_sync(Phase::Started, move || {
97            state.start();
98            Ok(())
99        });
100        let state = Arc::clone(&self.state);
101        lifecycle.register_sync(Phase::RestartFailed, move || {
102            if state.resume.swap(false, Ordering::AcqRel) {
103                state.start();
104            }
105            Ok(())
106        });
107        for phase in [Phase::Restart, Phase::Shutdown] {
108            let state = Arc::clone(&self.state);
109            lifecycle.register_sync(phase, move || {
110                state.stop();
111                Ok(())
112            });
113        }
114    }
115}
116
117impl State {
118    fn start(self: &Arc<Self>) {
119        self.cancel.send_replace(false);
120        self.disabled.store(false, Ordering::Release);
121        self.seen.store(0, Ordering::Relaxed);
122        let state = Arc::clone(self);
123        *self.task.lock().unwrap() = Some(tokio::spawn(async move {
124            state.probe().await;
125        }));
126    }
127
128    fn stop(&self) {
129        self.cancel.send_replace(true);
130        let active = !self.disabled.swap(true, Ordering::AcqRel);
131        if let Some(task) = self.task.lock().unwrap().take() {
132            self.resume.store(
133                active && !self.complete.load(Ordering::Acquire),
134                Ordering::Release,
135            );
136            task.abort();
137        }
138    }
139
140    async fn probe(&self) {
141        let deadline = Instant::now() + Duration::from_secs(30);
142        let address = *self.address.lock().unwrap();
143        while Instant::now() < deadline {
144            if timeout_at(deadline, self.exchange(address))
145                .await
146                .is_ok_and(|result| result.is_ok())
147            {
148                break;
149            }
150            self.seen.store(0, Ordering::Relaxed);
151            tokio::time::sleep_until((Instant::now() + Duration::from_secs(1)).min(deadline)).await;
152        }
153        self.complete.store(true, Ordering::Release);
154        self.disabled.store(true, Ordering::Release);
155    }
156
157    async fn exchange(&self, address: SocketAddr) -> io::Result<()> {
158        let name = self
159            .probe
160            .as_ref()
161            .ok_or_else(|| io::Error::other("invalid probe name"))?;
162        let mut id = [0; 2];
163        getrandom::getrandom(&mut id).map_err(|_| io::Error::other("random ID unavailable"))?;
164        let id = u16::from_ne_bytes(id);
165        let mut wire = Vec::new();
166        let mut builder =
167            MessageBuilder::new(&mut wire, 512, id, 0x0100).map_err(io::Error::other)?;
168        builder
169            .question(
170                NameRef::parse(name.as_wire(), 0).map_err(io::Error::other)?,
171                Type::HINFO,
172                Class::IN,
173            )
174            .map_err(io::Error::other)?;
175        let _ = builder.finish();
176        let bind = if address.is_ipv4() {
177            "0.0.0.0:0"
178        } else {
179            "[::]:0"
180        };
181        let socket = UdpSocket::bind(bind).await?;
182        socket.connect(address).await?;
183        socket.send(&wire).await?;
184        let mut response = [0; 65535];
185        let len =
186            tokio::time::timeout(Duration::from_secs(2), socket.recv(&mut response)).await??;
187        let message = Message::parse(&response[..len]).map_err(io::Error::other)?;
188        // RFC 1035 ยง4.1.1: the response ID matches the query ID.
189        if message.id() != id {
190            return Err(io::Error::other("probe response ID mismatch"));
191        }
192        Ok(())
193    }
194}
195
196impl Plugin for Loop {
197    async fn serve<N: Next + Send>(&self, cx: &mut Context<'_>, next: N) -> Outcome {
198        let state = &self.state;
199        if self.is_probe(&cx.query) && state.seen.fetch_add(1, Ordering::Relaxed) >= 2 {
200            let address = *state.address.lock().unwrap();
201            pharos_server::logging::log(
202                "FATAL",
203                "loop",
204                &format!(
205                    "Loop ({} -> {}) detected for zone {:?}, see https://coredns.io/plugins/loop#troubleshooting. Query: \"HINFO {}\"",
206                    cx.client, address, state.domain, state.qname
207                ),
208            );
209            let _ = io::stdout().flush();
210            // CoreDNS log.Fatalf exits immediately, skipping all shutdown hooks.
211            std::process::exit(1);
212        }
213        next.run(cx).await
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use pharos_proto::{Class, Message, MessageBuilder, Name, NameRef, Type};
221    use pharos_server::{Metadata, Protocol, Transport};
222    use std::sync::atomic::Ordering;
223
224    struct Pass;
225    impl Next for Pass {
226        async fn run(self, _: &mut Context<'_>) -> Outcome {
227            Outcome::Unanswered(pharos_proto::Rcode::NOERROR)
228        }
229    }
230
231    async fn query(plugin: &Loop, name: &str, kind: Type) -> Outcome {
232        let name: Name = name.parse().unwrap();
233        let mut wire = Vec::new();
234        let mut builder = MessageBuilder::new(&mut wire, 512, 1, 0).unwrap();
235        builder
236            .question(NameRef::parse(name.as_wire(), 0).unwrap(), kind, Class::IN)
237            .unwrap();
238        let _ = builder.finish();
239        let domain: Name = ".".parse().unwrap();
240        let mut response = Vec::new();
241        let mut slots = [];
242        let mut cx = Context {
243            query: Message::parse(&wire).unwrap(),
244            client: "127.0.0.1:1000".parse().unwrap(),
245            local: "127.0.0.1:53".parse().unwrap(),
246            transport: Transport::Dns,
247            protocol: Protocol::Udp,
248            domain: &domain,
249            metadata: Metadata::new(&mut slots),
250            response: &mut response,
251            max_response_size: 512,
252        };
253        plugin.serve(&mut cx, Pass).await
254    }
255
256    // Derived from CoreDNS v1.15.0 plugin/loop/loop_test.go.
257    #[tokio::test]
258    async fn counts_only_enabled_matching_hinfo_queries() {
259        let plugin = Loop::new("example.org.");
260        let name = plugin.state.qname.to_string();
261        for (query_name, kind) in [
262            ("outside.example.", Type::HINFO),
263            (name.as_str(), Type::A),
264            ("other.example.org.", Type::HINFO),
265        ] {
266            assert_eq!(
267                query(&plugin, query_name, kind).await,
268                Outcome::Unanswered(pharos_proto::Rcode::NOERROR)
269            );
270        }
271        assert_eq!(plugin.state.seen.load(Ordering::Relaxed), 0);
272        for count in 1..=2 {
273            assert_eq!(
274                query(&plugin, &name, Type::HINFO).await,
275                Outcome::Unanswered(pharos_proto::Rcode::NOERROR)
276            );
277            assert_eq!(plugin.state.seen.load(Ordering::Relaxed), count);
278        }
279        plugin.state.disabled.store(true, Ordering::Relaxed);
280        assert_eq!(
281            query(&plugin, &name, Type::HINFO).await,
282            Outcome::Unanswered(pharos_proto::Rcode::NOERROR)
283        );
284        assert_eq!(plugin.state.seen.load(Ordering::Relaxed), 2);
285    }
286
287    #[tokio::test]
288    async fn startup_probe_stops_on_restart_and_disables_after_success() {
289        let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
290        let plugin = Loop::new(".");
291        let mut lifecycle = Lifecycle::default();
292        plugin.register(&mut lifecycle, socket.local_addr().unwrap());
293        lifecycle.run(Phase::Started).await.unwrap();
294        let mut wire = [0; 512];
295        let (len, _) = tokio::time::timeout(Duration::from_secs(2), socket.recv_from(&mut wire))
296            .await
297            .unwrap()
298            .unwrap();
299        let query = Message::parse(&wire[..len]).unwrap();
300        assert_eq!(query.questions().next().unwrap().qtype, Type::HINFO);
301        assert_eq!(
302            query.questions().next().unwrap().name.to_string(),
303            plugin.state.qname
304        );
305        lifecycle.run(Phase::Restart).await.unwrap();
306        assert!(plugin.state.disabled.load(Ordering::Acquire));
307        lifecycle.run(Phase::RestartFailed).await.unwrap();
308        let (len, peer) = tokio::time::timeout(Duration::from_secs(2), socket.recv_from(&mut wire))
309            .await
310            .unwrap()
311            .unwrap();
312        wire[2] |= 0x80;
313        socket.send_to(&wire[..len], peer).await.unwrap();
314        tokio::time::sleep(Duration::from_millis(200)).await;
315        assert!(plugin.state.disabled.load(Ordering::Acquire));
316        assert!(plugin.state.complete.load(Ordering::Acquire));
317        lifecycle.run(Phase::Shutdown).await.unwrap();
318    }
319
320    #[tokio::test]
321    async fn failed_reload_does_not_repeat_a_successful_probe() {
322        let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
323        let plugin = Loop::new(".");
324        let mut lifecycle = Lifecycle::default();
325        plugin.register(&mut lifecycle, socket.local_addr().unwrap());
326        lifecycle.run(Phase::Started).await.unwrap();
327        let mut wire = [0; 512];
328        let (len, peer) = socket.recv_from(&mut wire).await.unwrap();
329        wire[2] |= 0x80;
330        socket.send_to(&wire[..len], peer).await.unwrap();
331        tokio::time::timeout(Duration::from_secs(1), async {
332            while !plugin.state.complete.load(Ordering::Acquire) {
333                tokio::task::yield_now().await;
334            }
335        })
336        .await
337        .unwrap();
338        lifecycle.run(Phase::Restart).await.unwrap();
339        lifecycle.run(Phase::RestartFailed).await.unwrap();
340        assert!(
341            tokio::time::timeout(Duration::from_millis(100), socket.recv_from(&mut wire))
342                .await
343                .is_err()
344        );
345        assert!(plugin.state.disabled.load(Ordering::Acquire));
346        lifecycle.run(Phase::Shutdown).await.unwrap();
347    }
348}