netget/server/doq/
mod.rs

1//! DNS over QUIC. All connection and stream futures are owned by the registered
2//! accept task; aborting it drops them and closes the endpoint, with no detached I/O.
3pub mod actions;
4pub mod wire;
5
6use crate::llm::action_helper::call_llm;
7use crate::llm::ActionResult;
8use crate::logging::emit::Log;
9use crate::protocol::{Event, SpawnContext};
10use crate::server::connection::ConnectionId;
11use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
12use actions::{DoqProtocol, DOQ_QUERY_EVENT};
13use anyhow::{bail, Context, Result};
14use futures::{stream::FuturesUnordered, StreamExt};
15use hickory_proto::op::{Message, MessageType, OpCode, ResponseCode};
16use hickory_proto::rr::{DNSClass, RecordType};
17use serde_json::json;
18use std::{net::SocketAddr, sync::Arc, time::Duration};
19use wire::*;
20
21#[derive(Clone)]
22struct Settings {
23    handshake: Duration,
24    exchange: Duration,
25    idle: Duration,
26    connections: usize,
27    streams: usize,
28}
29impl Settings {
30    fn read(params: Option<&crate::protocol::StartupParams>) -> Result<Self> {
31        Ok(Self {
32            handshake: Duration::from_secs(bounded_parameter(
33                params,
34                "handshake_timeout_secs",
35                HANDSHAKE_TIMEOUT.as_secs(),
36                60,
37            )?),
38            exchange: Duration::from_secs(bounded_parameter(
39                params,
40                "exchange_timeout_secs",
41                EXCHANGE_TIMEOUT.as_secs(),
42                300,
43            )?),
44            idle: Duration::from_secs(bounded_parameter(
45                params,
46                "idle_timeout_secs",
47                IDLE_TIMEOUT.as_secs(),
48                3600,
49            )?),
50            connections: bounded_parameter(params, "max_connections", MAX_CONNECTIONS as u64, 256)?
51                as usize,
52            streams: bounded_parameter(
53                params,
54                "max_streams",
55                MAX_STREAMS as u64,
56                MAX_STREAMS as u64,
57            )? as usize,
58        })
59    }
60}
61
62pub struct DoqServer;
63impl DoqServer {
64    pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
65        let bind = ctx.legacy_listen_addr();
66        if bind.port() == 53 {
67            bail!("RFC 9250 forbids DoQ on UDP port 53");
68        }
69        let settings = Settings::read(ctx.startup_params.as_ref())?;
70        let cert = ctx
71            .startup_params
72            .as_ref()
73            .map(|p| p.get_optional_string("cert_path"))
74            .transpose()?
75            .flatten();
76        let key = ctx
77            .startup_params
78            .as_ref()
79            .map(|p| p.get_optional_string("key_path"))
80            .transpose()?
81            .flatten();
82        let _ = rustls::crypto::ring::default_provider().install_default();
83        let tls = match (cert, key) {
84            (Some(cert), Some(key)) => {
85                crate::server::tls_cert_manager::load_tls_config_from_files(&cert, &key)?
86            }
87            (None, None) => crate::server::tls_cert_manager::generate_default_tls_config()?,
88            _ => bail!("DoQ cert_path and key_path must be supplied together"),
89        };
90        let mut crypto = (*tls).clone();
91        crypto.alpn_protocols = vec![b"doq".to_vec()];
92        crypto.max_early_data_size = 0;
93        let mut config = quinn::ServerConfig::with_crypto(Arc::new(
94            quinn::crypto::rustls::QuicServerConfig::try_from(crypto)?,
95        ));
96        config.transport_config(transport(settings.idle, settings.streams as u32));
97        // Disable migration: this endpoint is a simulator, not a mobile resolver session.
98        config.migration(false);
99        let endpoint =
100            EndpointGuard(quinn::Endpoint::server(config, bind).context("Bind DoQ UDP endpoint")?);
101        let local = endpoint.0.local_addr()?;
102        Log::new(Some(&ctx.status_tx)).info(format!("DoQ listening on {local}"));
103        let state = ctx.state.clone();
104        let id = ctx.server_id;
105        let task = tokio::spawn(async move {
106            let mut sessions = FuturesUnordered::new();
107            loop {
108                tokio::select! {
109                    incoming = endpoint.0.accept() => {
110                        let Some(incoming) = incoming else { break; };
111                        if sessions.len() >= settings.connections {
112                            incoming.refuse();
113                            Log::new(Some(&ctx.status_tx)).warn("DoQ connection limit reached");
114                            continue;
115                        }
116                        sessions.push(Self::session(incoming, ctx.clone(), settings.clone(), local));
117                    }
118                    _ = sessions.next(), if !sessions.is_empty() => {}
119                }
120            }
121        });
122        state.register_server_task(id, task).await;
123        Ok(local)
124    }
125
126    async fn session(
127        incoming: quinn::Incoming,
128        ctx: SpawnContext,
129        settings: Settings,
130        local: SocketAddr,
131    ) {
132        let connection = match tokio::time::timeout(settings.handshake, incoming).await {
133            Ok(Ok(c)) => ConnectionGuard(c),
134            _ => return,
135        };
136        let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
137        let now = crate::utils::clock::Instant::now();
138        ctx.state
139            .add_connection_to_server(
140                ctx.server_id,
141                ConnectionState {
142                    id,
143                    remote_addr: connection.0.remote_address(),
144                    local_addr: local,
145                    bytes_sent: 0,
146                    bytes_received: 0,
147                    packets_sent: 0,
148                    packets_received: 0,
149                    last_activity: now,
150                    status: ConnectionStatus::Active,
151                    status_changed_at: now,
152                    protocol_info: ProtocolConnectionInfo::new(json!({"alpn":"doq"})),
153                },
154            )
155            .await;
156        let _ = ctx.status_tx.send("__UPDATE_UI__".into());
157        let mut streams = FuturesUnordered::new();
158        loop {
159            tokio::select! {
160                result = connection.0.accept_bi(), if streams.len() < settings.streams => {
161                    match result {
162                        Ok((send,recv)) => streams.push(Self::transaction(send,recv,&connection.0,&ctx,id,settings.exchange)),
163                        Err(_) => break,
164                    }
165                }
166                _ = connection.0.closed() => break,
167                _ = streams.next(), if !streams.is_empty() => {}
168            }
169        }
170        // Dropping these futures cancels handlers and resets any unfinished responses.
171        drop(streams);
172        ctx.state
173            .close_connection_on_server(ctx.server_id, id)
174            .await;
175        let _ = ctx.status_tx.send("__UPDATE_UI__".into());
176    }
177
178    async fn transaction(
179        mut send: quinn::SendStream,
180        mut recv: quinn::RecvStream,
181        connection: &quinn::Connection,
182        ctx: &SpawnContext,
183        id: ConnectionId,
184        deadline: Duration,
185    ) {
186        let end = tokio::time::Instant::now() + deadline;
187        let exchange = async {
188            let frame = match recv.read_to_end(MAX_FRAME_BYTES).await {
189                Ok(frame) => frame,
190                // Reset by the client is cancellation, not malformed input. Other
191                // streams continue on the same connection.
192                Err(quinn::ReadToEndError::Read(quinn::ReadError::Reset(_))) => return Ok(None),
193                Err(quinn::ReadToEndError::TooLong) => {
194                    connection.close(PROTOCOL_ERROR, b"DoQ frame too large");
195                    return Ok(None);
196                }
197                Err(e) => return Err(anyhow::Error::from(e)),
198            };
199            let query = match decode(&frame, MessageType::Query) {
200                Ok(message) => message,
201                Err(_) => {
202                    connection.close(PROTOCOL_ERROR, b"invalid DoQ query");
203                    return Ok(None);
204                }
205            };
206            ctx.state
207                .update_connection_stats(
208                    ctx.server_id,
209                    id,
210                    Some(frame.len() as u64),
211                    None,
212                    Some(1),
213                    None,
214                )
215                .await;
216            let response = if query.queries().len() != 1 {
217                error_response(&query, ResponseCode::FormErr)
218            } else if query.op_code() != OpCode::Query
219                || query.queries()[0].query_class() != DNSClass::IN
220                || !matches!(
221                    query.queries()[0].query_type(),
222                    RecordType::A
223                        | RecordType::AAAA
224                        | RecordType::MX
225                        | RecordType::TXT
226                        | RecordType::CNAME
227                        | RecordType::ANY
228                )
229            {
230                error_response(&query, ResponseCode::NotImp)
231            } else {
232                let q = &query.queries()[0];
233                let event = Event::new(
234                    &DOQ_QUERY_EVENT,
235                    json!({
236                        "query_id":0,"domain":q.name().to_utf8(),"query_type":q.query_type().to_string(),
237                        "query_class":q.query_class().to_string(),"recursion_desired":query.recursion_desired(),
238                        "stream_id":u64::from(recv.id()),"peer_addr":connection.remote_address().to_string(),
239                    }),
240                );
241                Log::new(Some(&ctx.status_tx)).debug(format!(
242                    "DoQ {} {} stream {}",
243                    q.query_type(),
244                    q.name(),
245                    recv.id()
246                ));
247                Log::new(Some(&ctx.status_tx)).trace(format!("DoQ query: {}", event.data));
248                match call_llm(
249                    &ctx.llm_client,
250                    &ctx.state,
251                    ctx.server_id,
252                    Some(id),
253                    &event,
254                    &DoqProtocol::new(),
255                )
256                .await
257                {
258                    Ok(result) => {
259                        if !result.failures.is_empty() {
260                            Log::new(Some(&ctx.status_tx)).error(format!(
261                                "DoQ stream {} decision=fail_closed_action_error: {} action(s) failed; answering SERVFAIL",
262                                recv.id(), result.failures.len()
263                            ));
264                            return Ok(Some(encode(&error_response(
265                                &query,
266                                ResponseCode::ServFail,
267                            ))?));
268                        }
269                        for message in result.messages {
270                            Log::new(Some(&ctx.status_tx)).info(message);
271                        }
272                        let mut messages = Vec::new();
273                        for result in result.protocol_results {
274                            if let ActionResult::Output(bytes) = result {
275                                messages.push(Message::from_vec(&bytes)?);
276                            }
277                        }
278                        if messages.is_empty() {
279                            Log::new(Some(&ctx.status_tx)).info(format!(
280                                "DoQ stream {} decision=model_silent; cancelling query with DOQ_NO_ERROR",
281                                recv.id()
282                            ));
283                            // A dropped query is an explicit transaction cancellation, not an
284                            // empty DNS stream or an indefinitely dangling request.
285                            return Ok(None);
286                        }
287                        let mut response = messages.remove(0);
288                        for message in messages {
289                            response.add_answers(message.answers().iter().cloned());
290                        }
291                        response
292                            .set_id(0)
293                            .set_message_type(MessageType::Response)
294                            .set_op_code(query.op_code())
295                            .set_recursion_desired(query.recursion_desired());
296                        *response.queries_mut() = query.queries().to_vec();
297                        let decision = if matches!(
298                            response.response_code(),
299                            ResponseCode::NXDomain | ResponseCode::Refused
300                        ) {
301                            "model_reject"
302                        } else {
303                            "model_answer"
304                        };
305                        Log::new(Some(&ctx.status_tx)).info(format!(
306                            "DoQ stream {} decision={decision} rcode={}",
307                            recv.id(),
308                            response.response_code()
309                        ));
310                        response
311                    }
312                    Err(e) => {
313                        let category = if crate::utils::WireFailure::classify(&e).is_overloaded() {
314                            "overloaded"
315                        } else {
316                            "unavailable"
317                        };
318                        Log::new(Some(&ctx.status_tx)).error(format!(
319                            "DoQ stream {} decision=fail_closed_llm_error category={category}; answering SERVFAIL: {e}", recv.id()
320                        ));
321                        error_response(&query, ResponseCode::ServFail)
322                    }
323                }
324            };
325            Ok::<_, anyhow::Error>(Some(encode(&response)?))
326        };
327        // While a handler runs, STOP_SENDING promptly drops the work rather than
328        // retaining its LLM request until the exchange deadline.
329        let result = tokio::select! {
330            _ = send.stopped() => {
331                Log::new(Some(&ctx.status_tx)).info(format!("DoQ decision=peer_cancelled stream {}", send.id()));
332                let _ = recv.stop(REQUEST_CANCELLED); return;
333            }
334            result = tokio::time::timeout_at(end, exchange) => result,
335        };
336        match result {
337            Ok(Ok(Some(frame))) => {
338                let sent = tokio::time::timeout_at(end, async {
339                    send.write_all(&frame).await?;
340                    send.finish()?;
341                    Ok::<_, anyhow::Error>(())
342                })
343                .await;
344                if matches!(sent, Ok(Ok(()))) {
345                    ctx.state
346                        .update_connection_stats(
347                            ctx.server_id,
348                            id,
349                            None,
350                            Some(frame.len() as u64),
351                            None,
352                            Some(1),
353                        )
354                        .await;
355                } else {
356                    Log::new(Some(&ctx.status_tx)).error(format!(
357                        "DoQ stream {} decision=fail_closed_send_error; resetting stream",
358                        send.id()
359                    ));
360                    let _ = send.reset(INTERNAL_ERROR);
361                }
362            }
363            Ok(Ok(None)) => {
364                let _ = send.reset(NO_ERROR);
365            }
366            Ok(Err(e)) => {
367                Log::new(Some(&ctx.status_tx)).error(format!(
368                    "DoQ stream {} decision=fail_closed_response_error: {e}",
369                    send.id()
370                ));
371                let _ = send.reset(INTERNAL_ERROR);
372            }
373            Err(_) => {
374                Log::new(Some(&ctx.status_tx)).error(format!(
375                    "DoQ stream {} decision=fail_closed_timeout; resetting stream",
376                    send.id()
377                ));
378                let _ = recv.stop(UNSPECIFIED_ERROR);
379                let _ = send.reset(UNSPECIFIED_ERROR);
380            }
381        }
382    }
383}
384
385fn error_response(query: &Message, code: ResponseCode) -> Message {
386    let mut response = Message::new();
387    response
388        .set_id(0)
389        .set_message_type(MessageType::Response)
390        .set_op_code(query.op_code())
391        .set_recursion_desired(query.recursion_desired())
392        .set_response_code(code)
393        .add_queries(query.queries().iter().cloned());
394    response
395}
netget/server/doq/mod.rs

netget/server/doq/
mod.rs
