1use pharos_proto::{EdnsOptions, Message, Opt, Rcode, Type, normalized_udp_size};
2
3use crate::{Action, Error, Outcome, Protocol, admission::write_error};
4
5pub fn response_limit(query: Message<'_>, protocol: Protocol) -> usize {
7 usize::from(normalized_udp_size(
9 query.opt().map_or(512, |opt| opt.udp_payload_size()),
10 protocol == Protocol::Tcp,
11 ))
12}
13
14pub fn finish_outcome(
16 query: Message<'_>,
17 protocol: Protocol,
18 outcome: Outcome,
19 response: &mut Vec<u8>,
20 mut log: impl FnMut(Error),
21) -> Action {
22 match outcome {
23 Outcome::Drop => return Action::Drop,
24 Outcome::Stream(_) => return Action::Drop,
26 Outcome::Unanswered(code) => {
27 if !matches!(
28 code,
29 Rcode::SERVFAIL | Rcode::REFUSED | Rcode::FORMERR | Rcode::NOTIMP
30 ) {
31 return Action::Drop;
32 }
33 write_error(query, response, code);
34 }
35 Outcome::Error(error) => {
36 log(error);
37 write_error(query, response, Rcode::SERVFAIL);
38 }
39 Outcome::Response => {}
40 }
41 if let Err(error) = scrub(query, response_limit(query, protocol), protocol, response) {
42 log(error);
43 write_error(query, response, Rcode::SERVFAIL);
44 }
45 Action::Respond
46}
47
48fn scrub(
49 query: Message<'_>,
50 limit: usize,
51 protocol: Protocol,
52 wire: &mut Vec<u8>,
53) -> Result<(), Error> {
54 let message = Message::parse(wire).map_err(|_| Error::InvalidQuery)?;
55 if query.opt().is_none() && message.opt().is_none() && wire.len() <= limit {
56 return Ok(());
57 }
58 let mut options = [0_u8; 65535];
59 let existing = message.opt().map(|opt| {
60 let bytes = opt.options().as_bytes();
61 options[..bytes.len()].copy_from_slice(bytes);
62 (
63 opt.udp_payload_size(),
64 opt.extended_rcode(),
65 opt.version(),
66 opt.flags(),
67 bytes.len(),
68 )
69 });
70 let existing = existing.map(|(size, code, version, flags, len)| {
71 Opt::new(
72 size,
73 code,
74 version,
75 flags,
76 EdnsOptions::parse(&options[..len]).expect("validated options"),
77 )
78 });
79 let rcode = message.rcode();
80 let response_opt = query.opt().map(|opt| opt.response(existing, &[], rcode));
82 let opt_len = response_opt.map_or(0, |opt| opt.wire_len());
83 let budget = limit
85 .max(512)
86 .checked_sub(opt_len)
87 .ok_or(Error::ResponseBufferTooSmall)?;
88 let mut end = 12
89 + message
90 .questions()
91 .map(|q| q.name.wire_len() + 4)
92 .sum::<usize>();
93 if end > budget {
94 return Err(Error::ResponseBufferTooSmall);
95 }
96 let mut counts = [message.question_count(), 0, 0, 0];
97 let original = [
98 message.question_count(),
99 message.answer_count(),
100 message.authority_count(),
101 message.additional_count() - u16::from(existing.is_some()),
102 ];
103 let mut stopped = false;
104 for (section, records) in [
105 (1, message.answers()),
106 (2, message.authorities()),
107 (3, message.additionals()),
108 ] {
109 for record in records {
110 if record.rtype == Type::OPT {
111 stopped = true;
113 break;
114 }
115 let next = record.rdata_offset + record.rdata.len();
116 if next > budget {
117 stopped = true;
118 break;
119 }
120 end = next;
121 counts[section] += 1;
122 }
123 if stopped {
124 break;
125 }
126 }
127 let truncated = counts != original;
128 if end + opt_len > wire.capacity() {
129 if protocol == Protocol::Udp {
130 return Err(Error::ResponseBufferTooSmall);
131 }
132 wire.reserve((end + opt_len).saturating_sub(wire.len()));
133 }
134 wire.truncate(end);
135 if truncated {
136 wire[2] |= 0x02;
137 }
138 if let Some(opt) = response_opt {
139 opt.encode(wire).map_err(Error::Build)?;
140 counts[3] += 1;
141 }
142 for (i, count) in counts.into_iter().enumerate() {
144 wire[4 + 2 * i..6 + 2 * i].copy_from_slice(&count.to_be_bytes());
145 }
146 Ok(())
147}