1use pharos_proto::{Class, MessageBuilder, NameRef, Opcode, Section, Type};
2use pharos_server::{Context, Error, Outcome, Protocol};
3use std::net::IpAddr;
4
5#[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 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}