netget/client/doq/
mod.rs

1//! An authenticated reusable QUIC connection, with independent bounded streams.
2//! The registered owner polls all query and handler futures; no detached tasks
3//! survive client removal and a parked handler never prevents command injection.
4pub mod actions;
5use crate::client::{command_support, llm_budget::call_llm_for_client};
6use crate::llm::actions::{
7    client_trait::{Client, ClientActionResult},
8    protocol_trait::Protocol,
9};
10use crate::logging::emit::Log;
11use crate::protocol::{ConnectContext, Event};
12use crate::server::doq::wire::*;
13use crate::state::{
14    client_handles::{ClientCommand, ClientSendOutcome},
15    AccessLogOwner, ClientStatus,
16};
17use actions::{DoqClientProtocol, DOQ_CONNECTED_EVENT, DOQ_ERROR_EVENT, DOQ_RESPONSE_EVENT};
18use anyhow::{bail, ensure, Context, Result};
19use futures::{future::BoxFuture, stream::FuturesUnordered, FutureExt, StreamExt};
20use hickory_proto::{
21    op::{Message, MessageType, Query},
22    rr::{Name, RecordType},
23};
24use serde_json::{json, Value};
25use std::{net::SocketAddr, str::FromStr, sync::Arc, time::Duration};
26
27const MAX_FOLLOWUP_DEPTH: u8 = 4;
28type HandlerFuture = BoxFuture<'static, (u8, Result<Vec<Value>>)>;
29type QueryFuture = BoxFuture<'static, (u8, Event)>;
30
31pub struct DoqClient;
32impl DoqClient {
33    pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
34        let params = ctx.startup_params.as_ref();
35        let handshake = Duration::from_secs(bounded_parameter(
36            params,
37            "handshake_timeout_secs",
38            HANDSHAKE_TIMEOUT.as_secs(),
39            60,
40        )?);
41        let exchange = Duration::from_secs(bounded_parameter(
42            params,
43            "exchange_timeout_secs",
44            EXCHANGE_TIMEOUT.as_secs(),
45            300,
46        )?);
47        let idle = Duration::from_secs(bounded_parameter(
48            params,
49            "idle_timeout_secs",
50            IDLE_TIMEOUT.as_secs(),
51            3600,
52        )?);
53        let auth_name = params
54            .map(|p| p.get_optional_string("server_name"))
55            .transpose()?
56            .flatten();
57        let ca = params
58            .map(|p| p.get_optional_string("ca_cert_path"))
59            .transpose()?
60            .flatten();
61        let (_, port) = ctx
62            .remote_addr
63            .rsplit_once(':')
64            .context("DoQ remote address must be hostname:port or [IPv6]:port")?;
65        ensure!(
66            port.parse::<u16>()? != 53,
67            "RFC 9250 forbids DoQ on UDP port 53"
68        );
69        let default_name = ctx
70            .remote_addr
71            .rsplit_once(':')
72            .unwrap()
73            .0
74            .trim_start_matches('[')
75            .trim_end_matches(']');
76        let server_name = auth_name.unwrap_or_else(|| default_name.to_string());
77        let mut roots = rustls::RootCertStore::empty();
78        roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
79        if let Some(path) = ca {
80            let pem = std::fs::read(path).context("Read DoQ trust certificate")?;
81            let certs = rustls_pemfile::certs(&mut pem.as_slice())
82                .collect::<std::result::Result<Vec<_>, _>>()?;
83            ensure!(!certs.is_empty(), "DoQ trust file contains no certificates");
84            for cert in certs {
85                roots.add(cert)?;
86            }
87        }
88        let _ = rustls::crypto::ring::default_provider().install_default();
89        let mut tls = rustls::ClientConfig::builder()
90            .with_root_certificates(roots)
91            .with_no_client_auth();
92        tls.alpn_protocols = vec![b"doq".to_vec()];
93        tls.enable_early_data = false;
94        let mut config = quinn::ClientConfig::new(Arc::new(
95            quinn::crypto::rustls::QuicClientConfig::try_from(tls)?,
96        ));
97        // A DoQ server never opens streams. Quinn enforces zero credit in both directions.
98        config.transport_config(transport(idle, 0));
99        let (endpoint, connection) = tokio::time::timeout(handshake, async {
100            let remote = tokio::net::lookup_host(&ctx.remote_addr)
101                .await?
102                .next()
103                .context("No addresses for DoQ server")?;
104            let bind = if remote.is_ipv6() {
105                "[::]:0"
106            } else {
107                "0.0.0.0:0"
108            };
109            let mut endpoint = quinn::Endpoint::client(bind.parse()?)?;
110            endpoint.set_default_client_config(config);
111            let endpoint = EndpointGuard(endpoint);
112            let connection = endpoint
113                .0
114                .connect(remote, &server_name)?
115                .await
116                .context("DoQ TLS handshake (check ca_cert_path and server_name)")?;
117            Ok::<_, anyhow::Error>((endpoint, ConnectionGuard(connection)))
118        })
119        .await
120        .context("DoQ connection deadline exceeded")??;
121        let local_addr = endpoint.0.local_addr()?;
122        ctx.state
123            .update_client_status(ctx.client_id, ClientStatus::Connected)
124            .await;
125        let command_rx = command_support::register_command_channel(&ctx.state, ctx.client_id).await;
126        let state = ctx.state.clone();
127        let id = ctx.client_id;
128        let task = tokio::spawn(Self::run(ctx, endpoint, connection, command_rx, exchange));
129        state.register_client_task(id, task).await;
130        Ok(local_addr)
131    }
132
133    async fn run(
134        ctx: ConnectContext,
135        _endpoint: EndpointGuard,
136        connection: ConnectionGuard,
137        mut commands: tokio::sync::mpsc::Receiver<ClientCommand>,
138        exchange: Duration,
139    ) {
140        let mut handlers: FuturesUnordered<HandlerFuture> = FuturesUnordered::new();
141        let mut queries: FuturesUnordered<QueryFuture> = FuturesUnordered::new();
142        handlers.push(handle_event(
143            ctx.clone(),
144            Event::new(&DOQ_CONNECTED_EVENT, json!({"remote_addr":ctx.remote_addr})),
145            0,
146        ));
147        let mut disconnect = false;
148        while !disconnect {
149            tokio::select! {
150                command = commands.recv() => {
151                    let Some(command) = command else { break; };
152                    disconnect = enqueue(&ctx,&connection.0,exchange,&mut queries,command.action.clone(),Some(command),0).await;
153                }
154                Some((depth,event)) = queries.next(), if !queries.is_empty() => {
155                    // The final response still reaches the handler; only subsequent actions
156                    // are stopped at the depth boundary. A returned disconnect is honoured.
157                    if handlers.len() < MAX_STREAMS {
158                        handlers.push(handle_event(ctx.clone(),event,depth));
159                    } else { Log::new(Some(&ctx.status_tx)).warn("DoQ response handler limit reached"); }
160                }
161                Some((depth,result)) = handlers.next(), if !handlers.is_empty() => {
162                    match result {
163                        Ok(actions) if actions.len() <= MAX_STREAMS => {
164                            for action in actions {
165                                if depth >= MAX_FOLLOWUP_DEPTH && action["type"] != "disconnect" {
166                                    if action["type"] == "send_dns_query" { Log::new(Some(&ctx.status_tx)).warn("DoQ follow-up limit reached"); }
167                                    continue;
168                                }
169                                if enqueue(&ctx,&connection.0,exchange,&mut queries,action,None,depth).await { disconnect = true; break; }
170                            }
171                        }
172                        Ok(_) => Log::new(Some(&ctx.status_tx)).warn("DoQ handler returned more than 32 actions"),
173                        Err(e) => Log::new(Some(&ctx.status_tx)).warn(format!("DoQ event handler failed: {e}")),
174                    }
175                }
176                _ = connection.0.closed() => break,
177            }
178        }
179        drop(queries);
180        drop(handlers);
181        connection.0.close(NO_ERROR, b"client disconnected");
182        ctx.state.remove_client_handle(ctx.client_id).await;
183        ctx.state
184            .update_client_status(ctx.client_id, ClientStatus::Disconnected)
185            .await;
186        let _ = ctx.status_tx.send("__UPDATE_UI__".into());
187    }
188}
189
190pub fn build_query(data: &Value) -> Result<Message> {
191    let domain = data["domain"].as_str().context("Missing domain")?;
192    let kind = RecordType::from_str(data["query_type"].as_str().context("Missing query_type")?)?;
193    ensure!(
194        !matches!(kind, RecordType::AXFR | RecordType::IXFR),
195        "DoQ zone transfers are not implemented"
196    );
197    let mut message = Message::new();
198    message
199        .set_id(0)
200        .set_recursion_desired(data["recursion_desired"].as_bool().unwrap_or(true))
201        .add_query(Query::query(Name::from_str(domain)?, kind));
202    Ok(message)
203}
204
205fn handle_event(ctx: ConnectContext, event: Event, depth: u8) -> HandlerFuture {
206    async move {
207        let result = async {
208            let instruction = ctx
209                .state
210                .get_instruction_for_client(ctx.client_id)
211                .await
212                .unwrap_or_default();
213            let memory = ctx
214                .state
215                .get_memory_for_client(ctx.client_id)
216                .await
217                .unwrap_or_default();
218            let result = call_llm_for_client(
219                &ctx.llm_client,
220                &ctx.state,
221                ctx.client_id.to_string(),
222                &instruction,
223                &memory,
224                Some(&event),
225                &DoqClientProtocol::new(),
226                &ctx.status_tx,
227            )
228            .await?;
229            if let Some(memory) = result.memory_updates {
230                ctx.state.set_memory_for_client(ctx.client_id, memory).await;
231            }
232            Ok(result.actions)
233        }
234        .await;
235        (depth, result)
236    }
237    .boxed()
238}
239
240async fn enqueue(
241    ctx: &ConnectContext,
242    connection: &quinn::Connection,
243    deadline: Duration,
244    queries: &mut FuturesUnordered<QueryFuture>,
245    action: Value,
246    command: Option<ClientCommand>,
247    depth: u8,
248) -> bool {
249    let outcome = match DoqClientProtocol::new().execute_action(action.clone()) {
250        Ok(ClientActionResult::Disconnect) => Ok(ClientSendOutcome::Disconnected),
251        Ok(ClientActionResult::WaitForMore | ClientActionResult::NoAction) => {
252            Ok(ClientSendOutcome::Executed {
253                detail: "Waiting for more queries".into(),
254            })
255        }
256        Ok(ClientActionResult::Custom { name, data }) if name == "dns_query" => {
257            if queries.len() >= MAX_STREAMS {
258                Err(anyhow::anyhow!("DoQ client busy: 32 active transactions"))
259            } else {
260                queries.push(query_future(
261                    ctx.clone(),
262                    connection.clone(),
263                    deadline,
264                    data,
265                    action,
266                    command,
267                    depth + 1,
268                ));
269                return false;
270            }
271        }
272        Ok(_) => Err(anyhow::anyhow!("Unsupported DoQ client action")),
273        Err(e) => Err(e),
274    };
275    let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
276    record_action(ctx, &action, &outcome).await;
277    if let Some(command) = command {
278        command_support::reply(command, outcome);
279    }
280    disconnect
281}
282
283async fn record_action(ctx: &ConnectContext, action: &Value, outcome: &Result<ClientSendOutcome>) {
284    let result = match outcome {
285        Ok(v) => serde_json::to_value(v).unwrap_or(Value::Null),
286        Err(e) => json!({"error":e.to_string()}),
287    };
288    ctx.state
289        .record_access_log(
290            AccessLogOwner::Client(ctx.client_id.as_u32()),
291            DoqClientProtocol::new().protocol_name(),
292            None,
293            "client_action",
294            action.clone(),
295            vec![result],
296        )
297        .await;
298}
299
300fn query_future(
301    ctx: ConnectContext,
302    connection: quinn::Connection,
303    deadline: Duration,
304    data: Value,
305    action: Value,
306    command: Option<ClientCommand>,
307    depth: u8,
308) -> QueryFuture {
309    async move {
310        let result = async {
311            let query = build_query(&data)?;
312            exchange(&connection,&query,deadline).await
313        }.await;
314        let (outcome,event) = match result {
315            Ok((message,stream)) => {
316                let records = |records: &[hickory_proto::rr::Record]| records.iter().map(|r|json!({"name":r.name().to_utf8(),"type":r.record_type().to_string(),"class":r.dns_class().to_string(),"ttl":r.ttl(),"data":r.data().map(ToString::to_string)})).collect::<Vec<_>>();
317                let event = Event::new(&DOQ_RESPONSE_EVENT,json!({"query_id":0,"stream_id":stream,"domain":data["domain"],"query_type":data["query_type"],"response_code":message.response_code().to_string(),"answers":records(message.answers()),"authorities":records(message.name_servers()),"additionals":records(message.additionals())}));
318                (Ok(ClientSendOutcome::Executed{detail:format!("DNS {} ({} answers)",message.response_code(),message.answers().len())}),event)
319            }
320            Err(e) => {
321                Log::new(Some(&ctx.status_tx)).warn(format!("DoQ query failed: {e:#}"));
322                let event = Event::new(&DOQ_ERROR_EVENT,json!({"domain":data["domain"],"query_type":data["query_type"],"error":format!("{e:#}")}));
323                (Err(e),event)
324            }
325        };
326        record_action(&ctx,&action,&outcome).await;
327        if let Some(command) = command { command_support::reply(command,outcome); }
328        (depth,event)
329    }.boxed()
330}
331
332/// Execute one ordinary DNS exchange. Deadline and cancellation cover open/write/read/FIN.
333/// Public for transport validation and embedding; CLI actions use the same path.
334pub async fn exchange(
335    connection: &quinn::Connection,
336    query: &Message,
337    deadline: Duration,
338) -> Result<(Message, u64)> {
339    let frame = encode(query)?;
340    let end = tokio::time::Instant::now() + deadline;
341    let (send, recv) = tokio::time::timeout_at(end, connection.open_bi())
342        .await
343        .context("DoQ stream-open timeout")??;
344    let mut streams = CancelStreams {
345        send,
346        recv,
347        done: false,
348    };
349    let result = tokio::time::timeout_at(end, async {
350        streams.send.write_all(&frame).await?;
351        streams.send.finish()?;
352        let frame = match streams.recv.read_to_end(MAX_FRAME_BYTES).await {
353            Ok(frame) => frame,
354            Err(quinn::ReadToEndError::TooLong) => {
355                connection.close(PROTOCOL_ERROR, b"DoQ response too large");
356                bail!("DoQ response exceeds the DNS message limit");
357            }
358            Err(e) => return Err(e.into()),
359        };
360        let response = match decode(&frame, MessageType::Response) {
361            Ok(response) => response,
362            Err(e) => {
363                connection.close(PROTOCOL_ERROR, b"invalid DoQ response");
364                return Err(e);
365            }
366        };
367        if response.queries() != query.queries() || response.op_code() != query.op_code() {
368            connection.close(PROTOCOL_ERROR, b"DoQ question or opcode mismatch");
369            bail!("DoQ response question or opcode does not match the query");
370        }
371        Ok((response, u64::from(streams.recv.id())))
372    })
373    .await
374    .context("DoQ query deadline exceeded")?;
375    if result.is_ok() {
376        streams.done = true;
377    }
378    result
379}
380struct CancelStreams {
381    send: quinn::SendStream,
382    recv: quinn::RecvStream,
383    done: bool,
384}
385impl Drop for CancelStreams {
386    fn drop(&mut self) {
387        if !self.done {
388            let _ = self.recv.stop(REQUEST_CANCELLED);
389            let _ = self.send.reset(REQUEST_CANCELLED);
390        }
391    }
392}
netget/client/doq/mod.rs

netget/client/doq/
mod.rs
