Skip to main content

pharos_plugins/
whoami.rs

1use pharos_proto::{Class, MessageBuilder, NameRef, Opcode, Section, Type};
2use pharos_server::{Context, Error, Outcome, Protocol};
3use std::net::IpAddr;
4
5/// Returns the client's address and protocol in the Additional section.
6#[derive(Debug)]
7pub struct Whoami;
8
9impl pharos_server::Plugin for Whoami {
10    async fn serve<N: pharos_server::Next + Send>(
11        &self,
12        context: &mut Context<'_>,
13        _next: N,
14    ) -> Outcome {
15        match response(context) {
16            Ok(()) => Outcome::Response,
17            Err(error) => Outcome::Error(error),
18        }
19    }
20}
21
22fn response(context: &mut Context<'_>) -> Result<(), Error> {
23    let question = context.query.questions().next();
24    let name = question.map_or_else(
25        || NameRef::parse(&[0], 0).expect("root name"),
26        |question| question.name,
27    );
28    let class = question.map_or(Class(0), |question| question.qclass);
29    let max_size = context.max_response_size.max(512);
30    if context.response.capacity() < max_size {
31        if context.protocol == Protocol::Tcp {
32            context
33                .response
34                .reserve(max_size.saturating_sub(context.response.len()));
35        } else {
36            return Err(Error::ResponseBufferTooSmall);
37        }
38    }
39    // RFC 1035 ยง3.1: encoded domain names are limited to 255 octets.
40    let mut srv_name = [0_u8; 255];
41    let label = match context.protocol {
42        Protocol::Udp => b"_udp",
43        Protocol::Tcp => b"_tcp",
44    };
45    srv_name[0] = 4;
46    srv_name[1..5].copy_from_slice(label);
47    let mut position = 5;
48    for label in name.labels() {
49        if position + 1 + label.len() >= srv_name.len() {
50            return Err(Error::InvalidQuery);
51        }
52        srv_name[position] = label.len() as u8;
53        position += 1;
54        srv_name[position..position + label.len()].copy_from_slice(label);
55        position += label.len();
56    }
57    let srv_name = NameRef::parse(&srv_name[..=position], 0).map_err(|_| Error::InvalidQuery)?;
58    let mut flags = 0x8400 | (context.query.flags() & 0x7800);
59    if context.query.opcode() == Opcode::QUERY {
60        flags |= context.query.flags() & 0x0110;
61    }
62    let mut builder = MessageBuilder::new(context.response, max_size, context.query.id(), flags)
63        .map_err(Error::Build)?;
64    if let Some(question) = question {
65        builder
66            .question(question.name, question.qtype, question.qclass)
67            .map_err(Error::Build)?;
68    }
69    match context.client.ip() {
70        IpAddr::V4(ip) => {
71            builder.record(Section::Additional, name, Type::A, class, 0, &ip.octets())
72        }
73        IpAddr::V6(ip) => builder.record(
74            Section::Additional,
75            name,
76            Type::AAAA,
77            class,
78            0,
79            &ip.octets(),
80        ),
81    }
82    .map_err(Error::Build)?;
83    let port = context.client.port().to_be_bytes();
84    builder
85        .record(
86            Section::Additional,
87            srv_name,
88            Type::SRV,
89            class,
90            0,
91            &[0, 0, 0, 0, port[0], port[1], 0],
92        )
93        .map_err(Error::Build)?;
94    let _ = builder.finish();
95    Ok(())
96}