pharos_plugins/forward/
mod.rs

1//! Forward Queries to configured Upstreams.
2
3mod resolve;
4mod setup;
5
6use std::sync::{
7    Mutex,
8    atomic::{AtomicUsize, Ordering},
9};
10use std::time::Duration;
11
12use pharos_proto::{Message, Name, NameRef, Rcode};
13use pharos_server::{Context, Next, Outcome, Plugin, Protocol, metrics::ForwardMetrics};
14use tokio::time::{Instant, timeout_at};
15
16use crate::proxy::{self, Options, Proxy};
17use setup::Policy;
18
19struct Upstream {
20    proxy: Option<Proxy>,
21    address: String,
22}
23
24struct Stanza {
25    from: Name,
26    ignored: Vec<Name>,
27    proxies: Vec<Upstream>,
28    options: Options,
29    policy: Policy,
30    robin: AtomicUsize,
31    concurrent: AtomicUsize,
32    max_concurrent: usize,
33    max_connect_attempts: Option<u32>,
34    failfast: bool,
35    failover: Vec<Rcode>,
36    next: Vec<Rcode>,
37    next_on_nodata: bool,
38}
39
40/// Multiple forward Directives form one contiguous Plugin Chain segment.
41#[derive(Default)]
42pub struct Forward {
43    stanzas: Vec<Stanza>,
44    metrics: Option<ForwardMetrics>,
45    buffers: BufferPool,
46}
47
48#[derive(Default)]
49// shortcut: the pool retains peak concurrency, cap it if idle memory becomes costly.
50struct BufferPool(Mutex<Vec<Scratch>>);
51
52struct Scratch {
53    bytes: Vec<u8>,
54    order: Vec<usize>,
55    query: Vec<u8>,
56}
57
58struct Buffer<'a> {
59    pool: &'a BufferPool,
60    scratch: Option<Scratch>,
61}
62
63impl BufferPool {
64    fn get(&self) -> Buffer<'_> {
65        let scratch = self.0.lock().unwrap().pop().unwrap_or_else(|| Scratch {
66            bytes: vec![0; 65535],
67            order: Vec::new(),
68            query: Vec::new(),
69        });
70        Buffer {
71            pool: self,
72            scratch: Some(scratch),
73        }
74    }
75}
76
77impl Buffer<'_> {
78    fn parts(&mut self) -> (&mut Vec<usize>, &mut [u8], &mut Vec<u8>) {
79        let scratch = self.scratch.as_mut().expect("buffer returned only on drop");
80        (&mut scratch.order, &mut scratch.bytes, &mut scratch.query)
81    }
82}
83
84impl Drop for Buffer<'_> {
85    fn drop(&mut self) {
86        self.pool
87            .0
88            .lock()
89            .unwrap()
90            .push(self.scratch.take().expect("buffer exists"));
91    }
92}
93
94impl std::fmt::Debug for Forward {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        f.debug_struct("Forward")
97            .field("stanzas", &self.stanzas.len())
98            .finish()
99    }
100}
101
102// RFC 1034 §3.1: a subdomain's name ends with its containing domain's name.
103fn contains(name: NameRef<'_>, domain: &Name) -> bool {
104    let domain = NameRef::parse(domain.as_wire(), 0).expect("valid configured Domain");
105    let count = name.labels().count();
106    let suffix = domain.labels().count();
107    count >= suffix
108        && name
109            .labels()
110            .skip(count - suffix)
111            .zip(domain.labels())
112            .all(|(a, b)| a.eq_ignore_ascii_case(b))
113}
114
115fn random(max: usize) -> usize {
116    let mut bytes = [0_u8; 8];
117    if getrandom::getrandom(&mut bytes).is_err() {
118        return 0;
119    }
120    (u64::from_ne_bytes(bytes) % max as u64) as usize
121}
122
123struct InFlight<'a>(&'a AtomicUsize);
124impl Drop for InFlight<'_> {
125    fn drop(&mut self) {
126        self.0.fetch_sub(1, Ordering::Relaxed);
127    }
128}
129
130impl Forward {
131    pub(crate) fn parse(
132        cursor: &mut pharos_corefile::Cursor<'_>,
133    ) -> Result<Self, pharos_corefile::Error> {
134        let mut result = Self::default();
135        while cursor.has_next() {
136            let config = setup::parse(cursor)?;
137            let parse_name = |name: &str| {
138                name.parse::<Name>()
139                    .map_err(|_| cursor.error(format!("invalid Domain {name}")))
140            };
141            let from = parse_name(&config.from)?;
142            let ignored = config
143                .except
144                .iter()
145                .map(|name| parse_name(name))
146                .collect::<Result<Vec<_>, _>>()?;
147            let domain = parse_name(&config.health_domain)?;
148            let mut proxies = Vec::new();
149            for target in config.upstreams {
150                let address = match target {
151                    setup::Target::Address(address) => address,
152                    setup::Target::Invalid(address) => {
153                        proxies.push(Upstream {
154                            proxy: None,
155                            address,
156                        });
157                        continue;
158                    }
159                };
160                let mut proxy = Proxy::new(address);
161                proxy.set_expire(config.expire);
162                proxy.set_max_age(config.max_age);
163                proxy.set_max_idle_conns(config.max_idle_conns);
164                proxy.set_read_timeout(config.read_timeout);
165                proxy.set_source_address(config.source_address);
166                let mut health = proxy::HealthConfig::new(address, config.max_fails);
167                health.domain = domain.clone();
168                health.recursive = config.health_recursion_desired;
169                health.force_tcp = config.force_tcp;
170                health.interval = config.health_interval;
171                health.source_address = config.source_address;
172                proxy.set_health_config(health);
173                proxies.push(Upstream {
174                    proxy: Some(proxy),
175                    address: address.to_string(),
176                });
177            }
178            result.stanzas.push(Stanza {
179                from,
180                ignored,
181                proxies,
182                options: Options {
183                    force_tcp: config.force_tcp,
184                    prefer_udp: config.prefer_udp,
185                },
186                policy: config.policy,
187                robin: AtomicUsize::new(0),
188                concurrent: AtomicUsize::new(0),
189                max_concurrent: config.max_concurrent,
190                max_connect_attempts: config.max_connect_attempts,
191                failfast: config.failfast_all_unhealthy_upstreams,
192                failover: config.failover.into_iter().map(Rcode).collect(),
193                next: config.next.into_iter().map(Rcode).collect(),
194                next_on_nodata: config.next_on_nodata,
195            });
196        }
197        Ok(result)
198    }
199
200    pub(crate) fn set_metrics(&mut self, registry: &pharos_server::metrics::Registry) {
201        self.metrics = Some(registry.forward());
202        for stanza in &mut self.stanzas {
203            for upstream in &mut stanza.proxies {
204                if let Some(proxy) = &mut upstream.proxy {
205                    proxy.set_metrics(registry, "forward");
206                }
207            }
208        }
209    }
210
211    pub(crate) fn stop(&self) {
212        for stanza in &self.stanzas {
213            for upstream in &stanza.proxies {
214                if let Some(proxy) = &upstream.proxy {
215                    proxy.stop();
216                }
217            }
218        }
219    }
220}
221
222impl Plugin for Forward {
223    async fn serve<N: Next + Send>(&self, cx: &mut Context<'_>, next: N) -> Outcome {
224        let root = [0];
225        let name = cx.query.questions().next().map_or_else(
226            || NameRef::parse(&root, 0).expect("root"),
227            |question| question.name,
228        );
229        for (index, stanza) in self.stanzas.iter().enumerate() {
230            if !contains(name, &stanza.from)
231                || (name != stanza.from
232                    && stanza.ignored.iter().any(|ignored| contains(name, ignored)))
233            {
234                continue;
235            }
236            let _guard = if stanza.max_concurrent > 0 {
237                let old = stanza.concurrent.fetch_add(1, Ordering::Relaxed);
238                let guard = InFlight(&stanza.concurrent);
239                if old >= stanza.max_concurrent {
240                    if let Some(metrics) = &self.metrics {
241                        metrics.concurrent_reject();
242                    }
243                    return Outcome::Unanswered(Rcode::REFUSED);
244                }
245                Some(guard)
246            } else {
247                None
248            };
249
250            let count = stanza.proxies.len();
251            let mut buffer = self.buffers.get();
252            let (order, wire, query_buffer) = buffer.parts();
253            let original_query = cx.query.as_wire();
254            let query_wire = if cx.query.wire_changed() {
255                if cx.query.flags() != u16::from_be_bytes([original_query[2], original_query[3]])
256                    && !cx.query.can_change_flags_in_wire()
257                {
258                    return Outcome::Unanswered(Rcode::FORMERR);
259                }
260                cx.query.copy_wire(query_buffer);
261                query_buffer.as_slice()
262            } else {
263                original_query
264            };
265            order.clear();
266            order.extend(0..count);
267            match stanza.policy {
268                Policy::Random => {
269                    for i in (1..count).rev() {
270                        order.swap(i, random(i + 1));
271                    }
272                }
273                Policy::RoundRobin => {
274                    let first = (stanza.robin.fetch_add(1, Ordering::Relaxed) + 1) % count;
275                    order[0] = first;
276                    let mut j = 1;
277                    for i in 0..count {
278                        if i != first {
279                            order[j] = i;
280                            j += 1;
281                        }
282                    }
283                }
284                Policy::Sequential => {}
285            }
286            let deadline = Instant::now() + Duration::from_secs(5);
287            let max_attempts = stanza
288                .max_connect_attempts
289                .map_or((count * 2) as u64, u64::from);
290            let mut attempts = 0_u64;
291            let mut i = 0;
292            let mut fails = 0;
293            let mut failovers = 0;
294            let mut try_next = false;
295            let protocol = if cx.protocol == Protocol::Tcp {
296                proxy::Transport::Tcp
297            } else {
298                proxy::Transport::Udp
299            };
300            while Instant::now() < deadline && (max_attempts == 0 || attempts < max_attempts) {
301                if i == count {
302                    i = 0;
303                    fails = 0;
304                }
305                let mut upstream = &stanza.proxies[order[i]];
306                i += 1;
307                if upstream.proxy.as_ref().is_some_and(Proxy::down) {
308                    fails += 1;
309                    if fails < count {
310                        continue;
311                    }
312                    if let Some(metrics) = &self.metrics {
313                        metrics.healthcheck_broken();
314                    }
315                    if stanza.failfast {
316                        break;
317                    }
318                    upstream = &stanza.proxies[random(count)];
319                }
320                if cx.metadata.enabled() {
321                    cx.metadata
322                        .insert_inline("forward/upstream", &upstream.address);
323                }
324                let Some(proxy) = &upstream.proxy else {
325                    attempts += 1;
326                    if fails >= count {
327                        break;
328                    }
329                    tokio::time::sleep(Duration::from_millis(1)).await;
330                    continue;
331                };
332                let result = loop {
333                    let result = timeout_at(
334                        deadline,
335                        proxy.exchange(query_wire, wire, protocol, stanza.options),
336                    )
337                    .await;
338                    if matches!(result, Ok(Err(proxy::Error::CachedClosed))) {
339                        continue;
340                    }
341                    break result;
342                };
343                let exchange = match result {
344                    Ok(Ok(exchange)) => exchange,
345                    Ok(Err(proxy::Error::Mismatch | proxy::Error::InvalidRequest)) => {
346                        return Outcome::Unanswered(Rcode::FORMERR);
347                    }
348                    Ok(Err(proxy::Error::UnsupportedRequest)) => {
349                        return Outcome::Unanswered(Rcode::NOTIMP);
350                    }
351                    _ => {
352                        attempts += 1;
353                        if fails >= count {
354                            break;
355                        }
356                        continue;
357                    }
358                };
359                let answer = &wire[..exchange.len];
360                let response = Message::parse(answer).expect("proxy validated Response");
361                if stanza.failover.contains(&response.rcode()) {
362                    failovers += 1;
363                    if failovers < count {
364                        continue;
365                    }
366                }
367                if index + 1 < self.stanzas.len()
368                    && (stanza.next.contains(&response.rcode())
369                        || (stanza.next_on_nodata
370                            && response.rcode() == Rcode::NOERROR
371                            && response.answer_count() == 0))
372                {
373                    try_next = true;
374                    break;
375                }
376                cx.response.clear();
377                cx.response.extend_from_slice(answer);
378                return Outcome::Response;
379            }
380            if !try_next {
381                return Outcome::Unanswered(Rcode::SERVFAIL);
382            }
383        }
384        next.run(cx).await
385    }
386}
387
388#[cfg(test)]
389mod tests {
390    use super::BufferPool;
391
392    #[test]
393    fn returned_buffers_keep_their_contents_and_are_reused() {
394        let pool = BufferPool::default();
395        {
396            let mut buffer = pool.get();
397            buffer.parts().1[65534] = 123;
398        }
399        let mut buffer = pool.get();
400        assert_eq!(buffer.parts().1[65534], 123);
401    }
402}
pharos_plugins/forward/mod.rs

pharos_plugins/forward/
mod.rs
