Skip to main content

netget/server/doq/
wire.rs

1//! RFC 9250 framing shared by the two DoQ endpoints. This is deliberately not DoT:
2//! each stream carries one transaction, IDs are zero, and FIN is part of the frame.
3use 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    // EDNS TCP keepalive is a connection-level error on a QUIC transport (RFC 9250 4.3.3).
53    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
73/// Closing on Drop matters when the owning registered task is aborted during removal.
74/// Quinn's driver otherwise retains the endpoint until all connection handles disappear.
75pub 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}