1use 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 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 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 #[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}