Skip to main content

pharos_server/
outcome.rs

1use pharos_proto::{EdnsOptions, Message, Opt, Rcode, Type, normalized_udp_size};
2
3use crate::{Action, Error, Outcome, Protocol, admission::write_error};
4
5/// Maximum Response size requested by the client for this protocol.
6pub fn response_limit(query: Message<'_>, protocol: Protocol) -> usize {
7    // RFC 6891 §6.2.3: EDNS advertises a UDP payload size, with a 512-octet minimum.
8    usize::from(normalized_udp_size(
9        query.opt().map_or(512, |opt| opt.udp_payload_size()),
10        protocol == Protocol::Tcp,
11    ))
12}
13
14/// Finish a Plugin outcome, reporting Plugin and malformed Response errors.
15pub 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        // shortcut: Zone Transfer streams need a streaming Listener.
25        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    // RFC 6891 §6.1.1: an EDNS request requires an OPT record in the response.
81    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    // RFC 6891 §7: a truncated EDNS response must retain the question and OPT record.
84    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                // shortcut: records following OPT are dropped to preserve compression offsets.
112                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    // RFC 1035 §4.1.1: header counts describe the records in each response section.
143    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}