netget/server/doq/
wire.rs1use anyhow::{bail, ensure, Context, Result};
4use hickory_proto::op::{Message, MessageType};
5use hickory_proto::rr::rdata::opt::EdnsCode;
6use hickory_proto::serialize::binary::{BinDecodable, BinDecoder};
7use std::sync::Arc;
8use std::time::Duration;
9
10pub const MAX_MESSAGE_BYTES: usize = u16::MAX as usize;
11pub const MAX_FRAME_BYTES: usize = MAX_MESSAGE_BYTES + 2;
12pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
13pub const EXCHANGE_TIMEOUT: Duration = Duration::from_secs(30);
14pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
15pub const MAX_CONNECTIONS: usize = 64;
16pub const MAX_STREAMS: usize = 32;
17pub const NO_ERROR: quinn::VarInt = quinn::VarInt::from_u32(0);
18pub const INTERNAL_ERROR: quinn::VarInt = quinn::VarInt::from_u32(1);
19pub const PROTOCOL_ERROR: quinn::VarInt = quinn::VarInt::from_u32(2);
20pub const REQUEST_CANCELLED: quinn::VarInt = quinn::VarInt::from_u32(3);
21pub const EXCESSIVE_LOAD: quinn::VarInt = quinn::VarInt::from_u32(4);
22pub const UNSPECIFIED_ERROR: quinn::VarInt = quinn::VarInt::from_u32(5);
23
24pub fn encode(message: &Message) -> Result<Vec<u8>> {
25 ensure!(message.id() == 0, "DoQ DNS message ID must be zero");
26 let bytes = message.to_vec()?;
27 let length = u16::try_from(bytes.len()).context("DoQ DNS message exceeds 65535 bytes")?;
28 let mut frame = Vec::with_capacity(bytes.len() + 2);
29 frame.extend_from_slice(&length.to_be_bytes());
30 frame.extend_from_slice(&bytes);
31 Ok(frame)
32}
33
34pub fn decode(frame: &[u8], expected: MessageType) -> Result<Message> {
35 ensure!(
36 frame.len() >= 14 && frame.len() <= MAX_FRAME_BYTES,
37 "Invalid DoQ frame size"
38 );
39 let length = u16::from_be_bytes([frame[0], frame[1]]) as usize;
40 ensure!(
41 length + 2 == frame.len(),
42 "DoQ stream must contain exactly one complete DNS message"
43 );
44 let mut decoder = BinDecoder::new(&frame[2..]);
45 let message = Message::read(&mut decoder).context("Invalid DoQ DNS message")?;
46 ensure!(decoder.is_empty(), "Trailing bytes after DoQ DNS message");
47 ensure!(message.id() == 0, "DoQ DNS message ID must be zero");
48 ensure!(
49 message.message_type() == expected,
50 "Unexpected DNS message direction"
51 );
52 ensure!(
54 message
55 .extensions()
56 .as_ref()
57 .is_none_or(|edns| edns.options().get(EdnsCode::Keepalive).is_none()),
58 "EDNS TCP keepalive is forbidden on DoQ"
59 );
60 Ok(message)
61}
62
63pub fn transport(idle: Duration, incoming_bidi: u32) -> Arc<quinn::TransportConfig> {
64 let mut transport = quinn::TransportConfig::default();
65 transport.max_concurrent_bidi_streams(incoming_bidi.into());
66 transport.max_concurrent_uni_streams(0u32.into());
67 transport.max_idle_timeout(Some(idle.try_into().expect("bounded idle timeout")));
68 transport.stream_receive_window((MAX_FRAME_BYTES as u32).into());
69 transport.receive_window((MAX_FRAME_BYTES as u32 * MAX_STREAMS as u32).into());
70 Arc::new(transport)
71}
72
73pub struct EndpointGuard(pub quinn::Endpoint);
76impl Drop for EndpointGuard {
77 fn drop(&mut self) {
78 self.0.close(NO_ERROR, b"endpoint stopped");
79 }
80}
81pub struct ConnectionGuard(pub quinn::Connection);
82impl Drop for ConnectionGuard {
83 fn drop(&mut self) {
84 self.0.close(NO_ERROR, b"connection stopped");
85 }
86}
87
88pub fn parameter(
89 name: &str,
90 description: &str,
91 example: serde_json::Value,
92 default: Option<serde_json::Value>,
93) -> crate::llm::actions::ParameterDefinition {
94 crate::llm::actions::ParameterDefinition {
95 name: name.into(),
96 description: description.into(),
97 type_hint: if example.is_number() {
98 "number"
99 } else {
100 "string"
101 }
102 .into(),
103 required: false,
104 example,
105 default,
106 }
107}
108
109pub fn bounded_parameter(
110 params: Option<&crate::protocol::StartupParams>,
111 name: &str,
112 default: u64,
113 max: u64,
114) -> Result<u64> {
115 let value = params
116 .map(|p| p.get_optional_u64(name))
117 .transpose()?
118 .flatten()
119 .unwrap_or(default);
120 if value == 0 || value > max {
121 bail!("{name} must be between 1 and {max}");
122 }
123 Ok(value)
124}