Skip to main content

pharos_plugins/
header.rs

1//! Modify DNS header flags in Queries and Responses.
2
3use pharos_corefile::{Cursor, Error};
4use pharos_server::{Context, Next, Outcome, Plugin};
5
6#[derive(Debug, Default)]
7struct Rules {
8    set: u16,
9    clear: u16,
10}
11
12impl Rules {
13    fn add(&mut self, flag: u16, set: bool) {
14        if set {
15            self.set |= flag;
16            self.clear &= !flag;
17        } else {
18            self.clear |= flag;
19            self.set &= !flag;
20        }
21    }
22
23    fn apply(&self, flags: u16) -> u16 {
24        (flags | self.set) & !self.clear
25    }
26}
27
28#[derive(Debug, Default)]
29pub struct Header {
30    query: Rules,
31    response: Rules,
32}
33
34impl Header {
35    pub(crate) fn parse(cursor: &mut Cursor<'_>) -> Result<Self, Error> {
36        while cursor.next() {
37            let mut header = Self::default();
38            let _ = cursor.remaining_args();
39            let mut any = false;
40            while cursor.next_block() {
41                let rules = if cursor.val().eq_ignore_ascii_case("query") {
42                    &mut header.query
43                } else if cursor.val().eq_ignore_ascii_case("response") {
44                    &mut header.response
45                } else {
46                    return Err(cursor.error("invalid header selector"));
47                };
48                let args = cursor.remaining_args();
49                if args.len() < 2 {
50                    return Err(cursor.arg_err());
51                }
52                let set = if args[0].eq_ignore_ascii_case("set") {
53                    true
54                } else if args[0].eq_ignore_ascii_case("clear") {
55                    false
56                } else {
57                    return Err(cursor.error("invalid header action"));
58                };
59                for flag in &args[1..] {
60                    // RFC 1035 ยง4.1.1: AA, RD and RA occupy these header bits.
61                    let bit = if flag.eq_ignore_ascii_case("aa") {
62                        0x0400
63                    } else if flag.eq_ignore_ascii_case("rd") {
64                        0x0100
65                    } else if flag.eq_ignore_ascii_case("ra") {
66                        0x0080
67                    } else {
68                        return Err(cursor.error("unsupported header flag"));
69                    };
70                    rules.add(bit, set);
71                    any = true;
72                }
73            }
74            if any {
75                return Ok(header);
76            }
77        }
78        Err(cursor.arg_err())
79    }
80}
81
82impl Plugin for Header {
83    async fn serve<N: Next + Send>(&self, cx: &mut Context<'_>, next: N) -> Outcome {
84        cx.query = cx.query.with_flags(self.query.apply(cx.query.flags()));
85        let outcome = next.run(cx).await;
86        if outcome == Outcome::Response && cx.response.len() >= 4 {
87            let flags = u16::from_be_bytes([cx.response[2], cx.response[3]]);
88            cx.response[2..4].copy_from_slice(&self.response.apply(flags).to_be_bytes());
89        }
90        outcome
91    }
92}