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