Skip to main content

netget/client/doh/
mod.rs

1//! DNS-over-HTTPS (DoH) client implementation
2pub mod actions;
3
4pub use actions::DohClientProtocol;
5
6use anyhow::{Context, Result};
7use hickory_proto::op::{Message, Query};
8use hickory_proto::rr::{Name, RecordType};
9use std::net::SocketAddr;
10use std::str::FromStr;
11use std::sync::Arc;
12use std::time::Duration;
13use tokio::sync::mpsc;
14use tracing::{debug, error, info, trace, warn};
15
16use crate::client::doh::actions::{DOH_CLIENT_CONNECTED_EVENT, DOH_CLIENT_RESPONSE_RECEIVED_EVENT};
17use crate::client::llm_budget::call_llm_for_client;
18use crate::llm::actions::client_trait::{Client, ClientActionResult};
19use crate::llm::actions::protocol_trait::Protocol;
20use crate::llm::ollama_client::OllamaClient;
21use crate::llm::ClientLlmResult;
22use crate::protocol::Event;
23use crate::state::app_state::AppState;
24use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
25use crate::state::{AccessLogOwner, ClientId, ClientStatus};
26
27/// How often the command loop re-checks that its client still exists. DoH is
28/// request/response over HTTPS with no read loop, so nothing else would notice.
29const REMOVAL_CHECK_INTERVAL: Duration = Duration::from_secs(5);
30
31/// What one executed action did. Shared vocabulary between the LLM action path and
32/// the injected-command loop.
33enum Applied {
34    /// A DNS query completed. `detail` describes it; `event_data` is the payload for
35    /// the `doh_response_received` event, which the caller delivers on its own terms
36    /// (inline for the LLM chain, from a separate task for an injected command).
37    Queried {
38        detail: String,
39        event_data: serde_json::Value,
40    },
41    /// The action asked to end the session.
42    Disconnect,
43    /// The model asked to wait for more.
44    WaitForMore,
45    /// The action ran but this client has nothing to do with the result.
46    Ignored(String),
47}
48
49/// DoH client that makes DNS queries over HTTPS
50pub struct DohClient;
51
52impl DohClient {
53    /// Connect to a DoH server with integrated LLM actions
54    pub async fn connect_with_llm_actions(
55        server_url: String,
56        llm_client: OllamaClient,
57        app_state: Arc<AppState>,
58        status_tx: mpsc::UnboundedSender<String>,
59        client_id: ClientId,
60        startup_params: Option<crate::protocol::StartupParams>,
61    ) -> Result<SocketAddr> {
62        info!("DoH client {} initializing for {}", client_id, server_url);
63
64        // Trust settings, read once and stored for the per-query HTTP client.
65        let ca_cert_pem = startup_params
66            .as_ref()
67            .and_then(|p| p.get_optional_string("ca_cert_pem").ok().flatten());
68        let insecure_skip_verify = startup_params
69            .as_ref()
70            .and_then(|p| p.get_optional_bool("insecure_skip_verify").ok().flatten())
71            .unwrap_or(false);
72
73        // Store server URL in client state
74        app_state
75            .with_client_mut(client_id, |client| {
76                client.set_protocol_field(
77                    "server_url".to_string(),
78                    serde_json::json!(server_url.clone()),
79                );
80                if let Some(pem) = &ca_cert_pem {
81                    client.set_protocol_field(
82                        "ca_cert_pem".to_string(),
83                        serde_json::json!(pem.clone()),
84                    );
85                }
86                client.set_protocol_field(
87                    "insecure_skip_verify".to_string(),
88                    serde_json::json!(insecure_skip_verify),
89                );
90            })
91            .await;
92
93        // Update status to connected
94        app_state
95            .update_client_status(client_id, ClientStatus::Connected)
96            .await;
97        let _ = status_tx.send(format!(
98            "[CLIENT] DoH client {} connected to {}",
99            client_id, server_url
100        ));
101        let _ = status_tx.send("__UPDATE_UI__".to_string());
102
103        info!("DoH client {} connected to {}", client_id, server_url);
104
105        let protocol = Arc::new(DohClientProtocol::new());
106
107        // Injected commands (the dashboard's [ send ]). Registered BEFORE the
108        // connected-event LLM call below: a dashboard-created client defaults to a
109        // `*` manual rule, so that call can park for minutes waiting for a human and
110        // [ send ] has to work for the whole park.
111        let command_rx =
112            crate::client::command_support::register_command_channel(&app_state, client_id).await;
113        let cmd_state = app_state.clone();
114        let cmd_llm = llm_client.clone();
115        let cmd_tx = status_tx.clone();
116        let cmd_protocol = protocol.clone();
117        let cmd_url = server_url.clone();
118        let cmd_task = tokio::spawn(async move {
119            Self::command_loop(
120                command_rx,
121                client_id,
122                cmd_state,
123                cmd_llm,
124                cmd_tx,
125                cmd_protocol,
126                cmd_url,
127            )
128            .await;
129        });
130        app_state.register_client_task(client_id, cmd_task).await;
131
132        // Send connected event to LLM
133        let connected_event = Event::new(
134            &DOH_CLIENT_CONNECTED_EVENT,
135            serde_json::json!({
136                "server_url": server_url,
137            }),
138        );
139
140        // Get instruction and memory for LLM call
141        let instruction = app_state
142            .get_instruction_for_client(client_id)
143            .await
144            .unwrap_or_default();
145        let memory = app_state
146            .get_memory_for_client(client_id)
147            .await
148            .unwrap_or_default();
149
150        // Call LLM with connected event
151        match call_llm_for_client(
152            &llm_client,
153            &app_state,
154            client_id.to_string(),
155            &instruction,
156            &memory,
157            Some(&connected_event),
158            protocol.as_ref(),
159            &status_tx,
160        )
161        .await
162        {
163            Ok(llm_result) => {
164                debug!(
165                    "DoH client {} received {} actions from LLM",
166                    client_id,
167                    llm_result.actions.len()
168                );
169
170                // Execute any immediate actions. Registered with AppState so
171                // stop_client can abort this task — dropping a JoinHandle only
172                // detaches it in Tokio.
173                let task_handle = tokio::spawn(Self::execute_llm_actions(
174                    client_id,
175                    llm_result,
176                    app_state.clone(),
177                    llm_client.clone(),
178                    status_tx.clone(),
179                    protocol.clone(),
180                    server_url.clone(),
181                ));
182                app_state.register_client_task(client_id, task_handle).await;
183            }
184            Err(e) => {
185                error!("DoH client {} LLM call failed: {}", client_id, e);
186            }
187        }
188
189        // Return dummy local address (DoH is over HTTPS, connectionless at app level)
190        Ok("0.0.0.0:0".parse().unwrap())
191    }
192
193    /// Drain injected commands until the channel closes (client removed) or an
194    /// injected `disconnect` ends the session.
195    ///
196    /// `command_support::handle_stream_client_command` cannot serve this client:
197    /// there is no write half, and `query_dns` yields `ClientActionResult::Custom`.
198    /// So the action goes through [`Self::apply_action`] — the same function the LLM
199    /// path uses — and the outcome is recorded and replied the way the generic arm
200    /// does it.
201    #[allow(clippy::too_many_arguments)]
202    async fn command_loop(
203        mut command_rx: mpsc::Receiver<ClientCommand>,
204        client_id: ClientId,
205        app_state: Arc<AppState>,
206        llm_client: OllamaClient,
207        status_tx: mpsc::UnboundedSender<String>,
208        protocol: Arc<DohClientProtocol>,
209        server_url: String,
210    ) {
211        let mut removal_check = tokio::time::interval(REMOVAL_CHECK_INTERVAL);
212        removal_check.tick().await; // the first tick completes immediately
213
214        loop {
215            tokio::select! {
216                received = command_rx.recv() => {
217                    let Some(command) = received else { break };
218                    if Self::handle_command(
219                        command,
220                        client_id,
221                        &app_state,
222                        &llm_client,
223                        &status_tx,
224                        &protocol,
225                        &server_url,
226                    )
227                    .await
228                    {
229                        break;
230                    }
231                }
232                _ = removal_check.tick() => {
233                    if app_state.get_client(client_id).await.is_none() {
234                        info!("DoH client {} stopped", client_id);
235                        break;
236                    }
237                }
238            }
239        }
240
241        app_state.remove_client_handle(client_id).await;
242        let _ = status_tx.send("__UPDATE_UI__".to_string());
243    }
244
245    /// Execute one injected action, record it, and reply. Returns `true` when the
246    /// command loop should stop.
247    #[allow(clippy::too_many_arguments)]
248    async fn handle_command(
249        command: ClientCommand,
250        client_id: ClientId,
251        app_state: &Arc<AppState>,
252        llm_client: &OllamaClient,
253        status_tx: &mpsc::UnboundedSender<String>,
254        protocol: &Arc<DohClientProtocol>,
255        server_url: &str,
256    ) -> bool {
257        let action = command.action.clone();
258        let outcome = match protocol.as_ref().execute_action(action.clone()) {
259            Err(e) => Ok(ClientSendOutcome::Rejected {
260                error: e.to_string(),
261            }),
262            Ok(action_result) => {
263                match Self::apply_action(action_result, client_id, server_url, &app_state).await {
264                    // Never `Sent`: reqwest owns the socket and reports no wire byte count
265                    // for the query, so a number here would be invented. `Executed`
266                    // carries the answer count and DNS rcode instead.
267                    Ok(Applied::Queried { detail, event_data }) => {
268                        // Deliver the response event from its own registered task: a
269                        // dashboard client's manual rule can park that LLM call for
270                        // minutes, and the command loop must stay responsive.
271                        let notify = tokio::spawn(Self::notify_response(
272                            client_id,
273                            event_data,
274                            app_state.clone(),
275                            llm_client.clone(),
276                            status_tx.clone(),
277                            protocol.clone(),
278                            server_url.to_string(),
279                        ));
280                        app_state.register_client_task(client_id, notify).await;
281                        Ok(ClientSendOutcome::Executed { detail })
282                    }
283                    Ok(Applied::Disconnect) => Ok(ClientSendOutcome::Disconnected),
284                    Ok(Applied::WaitForMore) => Ok(ClientSendOutcome::Executed {
285                        detail: "wait_for_more".to_string(),
286                    }),
287                    Ok(Applied::Ignored(detail)) => Ok(ClientSendOutcome::Executed { detail }),
288                    Err(e) => Err(e),
289                }
290            }
291        };
292
293        let outcome_json = match &outcome {
294            Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
295            Err(e) => serde_json::json!({"error": e.to_string()}),
296        };
297        app_state
298            .record_access_log(
299                AccessLogOwner::Client(client_id.as_u32()),
300                protocol.protocol_name(),
301                None,
302                "injected_action",
303                action,
304                vec![outcome_json],
305            )
306            .await;
307
308        let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
309        if let Err(e) = &outcome {
310            error!("DoH client {} injected action failed: {}", client_id, e);
311            let _ = status_tx.send(format!(
312                "[WARN] Client {} injected action failed: {}",
313                client_id, e
314            ));
315        }
316        let _ = status_tx.send("__UPDATE_UI__".to_string());
317        crate::client::command_support::reply(command, outcome);
318
319        if disconnect {
320            app_state
321                .update_client_status(client_id, ClientStatus::Disconnected)
322                .await;
323        }
324        disconnect
325    }
326
327    /// Execute actions returned by LLM
328    async fn execute_llm_actions(
329        client_id: ClientId,
330        llm_result: ClientLlmResult,
331        app_state: Arc<AppState>,
332        llm_client: OllamaClient,
333        status_tx: mpsc::UnboundedSender<String>,
334        protocol: Arc<DohClientProtocol>,
335        server_url: String,
336    ) {
337        // Update memory if provided
338        if let Some(memory) = llm_result.memory_updates {
339            app_state.set_memory_for_client(client_id, memory).await;
340        }
341
342        for action_value in llm_result.actions {
343            match protocol.as_ref().execute_action(action_value.clone()) {
344                Ok(action_result) => {
345                    match Self::apply_action(action_result, client_id, &server_url, &app_state)
346                        .await
347                    {
348                        Ok(Applied::Queried { detail, event_data }) => {
349                            debug!("DoH client {}: {}", client_id, detail);
350                            // Inline, not spawned: this is the model's own chain and
351                            // the next round of actions continues from here.
352                            Self::notify_response(
353                                client_id,
354                                event_data,
355                                app_state.clone(),
356                                llm_client.clone(),
357                                status_tx.clone(),
358                                protocol.clone(),
359                                server_url.clone(),
360                            )
361                            .await;
362                        }
363                        Ok(Applied::Disconnect) => {
364                            info!("DoH client {} disconnecting", client_id);
365                            app_state.remove_client_handle(client_id).await;
366                            app_state
367                                .update_client_status(client_id, ClientStatus::Disconnected)
368                                .await;
369                            let _ = status_tx
370                                .send(format!("[CLIENT] DoH client {} disconnected", client_id));
371                            let _ = status_tx.send("__UPDATE_UI__".to_string());
372                            break;
373                        }
374                        Ok(Applied::WaitForMore) => {
375                            debug!("DoH client {} waiting for more queries", client_id);
376                            break;
377                        }
378                        Ok(Applied::Ignored(detail)) => {
379                            error!("DoH client {} {}", client_id, detail);
380                        }
381                        Err(e) => {
382                            // `{:#}` rather than `{}`: reqwest's real cause (connection
383                            // refused, TLS rejected, bad status) lives in the source chain,
384                            // and `{}` prints only our own "DoH POST request failed"
385                            // context -- which names the operation and says nothing about
386                            // why it failed.
387                            error!("DoH client {} query failed: {:#}", client_id, e);
388                            let _ = status_tx.send(format!("[CLIENT] DoH query failed: {:#}", e));
389                        }
390                    }
391                }
392                Err(e) => {
393                    error!("DoH client {} action execution failed: {}", client_id, e);
394                }
395            }
396        }
397    }
398
399    /// Turn one executed action into a DNS-over-HTTPS query (or a session end).
400    ///
401    /// Shared by the LLM action path and the injected-command loop so the query
402    /// encoding and the DoH HTTP call exist exactly once.
403    /// Build an HTTPS client honouring this client's declared trust settings.
404    ///
405    /// `ca_cert_pem` adds one certificate to the system roots -- the right answer for a
406    /// private CA or a known self-signed server. `insecure_skip_verify` accepts any
407    /// certificate at all and is off unless explicitly asked for; it is named for what it
408    /// does because a DoH connection made with it is unauthenticated, which defeats most
409    /// of the point of doing DNS over HTTPS.
410    async fn build_http_client(
411        client_id: ClientId,
412        server_url: &str,
413        app_state: &Arc<AppState>,
414    ) -> Result<reqwest::Client> {
415        let (ca_pem, insecure) = app_state
416            .with_client_mut(client_id, |c| {
417                (
418                    c.get_protocol_field("ca_cert_pem")
419                        .and_then(|v| v.as_str())
420                        .map(|s| s.to_string()),
421                    c.get_protocol_field("insecure_skip_verify")
422                        .and_then(|v| v.as_bool())
423                        .unwrap_or(false),
424                )
425            })
426            .await
427            .unwrap_or((None, false));
428
429        // rustls explicitly, not reqwest's default TLS backend. RFC 8484 DoH runs over
430        // HTTP/2, and the peer is reached only if ALPN offers `h2` -- NetGet's own DoH
431        // server advertises `h2` alone and answers with HTTP/2 frames regardless, so a
432        // client that negotiated no protocol tries to parse them as HTTP/1.1 and fails with
433        // "invalid HTTP version parsed" before a single query gets through.
434        // One HTTP client per distinct trust configuration, kept for the process.
435        //
436        // This was rebuilt for every query. Constructing a rustls client loads the
437        // platform root store -- on macOS that means reading the keychain -- so every
438        // single DNS query paid for a fresh TLS stack and a fresh handshake, and none of
439        // the connection pooling reqwest provides was ever used. Under load it is
440        // seconds per query.
441        //
442        // Keyed by host as well as by trust settings, because the resolver override below
443        // is per-host: a client cached under one host and reused for another would carry
444        // the wrong override.
445        static CLIENTS: std::sync::OnceLock<
446            std::sync::Mutex<
447                std::collections::HashMap<(String, Option<String>, bool), reqwest::Client>,
448            >,
449        > = std::sync::OnceLock::new();
450        let cache = CLIENTS.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()));
451        let host = crate::llm::ollama_client::host_of(server_url).to_string();
452        let key = (host.clone(), ca_pem.clone(), insecure);
453        if let Ok(guard) = cache.lock() {
454            if let Some(client) = guard.get(&key) {
455                return Ok(client.clone());
456            }
457        }
458
459        if insecure {
460            warn!(
461                "DoH client {} is skipping certificate verification (insecure_skip_verify); \
462                 this connection is not authenticated",
463                client_id
464            );
465        }
466
467        // Built on a blocking thread, not here.
468        //
469        // `Client::builder().build()` sets up the rustls stack, which loads the platform
470        // root store -- on macOS that reads the keychain. That is a synchronous,
471        // syscall-heavy operation, and calling it from async context parks a tokio worker
472        // thread for as long as it takes. Under load it took long enough to stall this
473        // client's runtime entirely: the query logged "querying example.com" and then
474        // nothing, no response and no timeout, because the request future had not been
475        // created yet and there was no timeout to fire.
476        let pem_for_build = ca_pem.clone();
477        let host_for_build = host.clone();
478        let client = tokio::task::spawn_blocking(move || {
479            let mut builder = reqwest::Client::builder()
480                .use_rustls_tls()
481                // There was no timeout at all, so a DoH server that accepted the
482                // connection and then said nothing hung the query forever with nothing in
483                // the log. A DNS query that takes longer than this is useless anyway.
484                .timeout(std::time::Duration::from_secs(10));
485
486            // Never resolve a literal IP through the system resolver.
487            //
488            // `reqwest` hands the URL host to its resolver unconditionally and
489            // `hyper-util`'s `GaiResolver` does not special-case a dotted quad, so
490            // `https://127.0.0.1:8443/dns-query` performs a real `getaddrinfo("127.0.0.1")`.
491            // On macOS that goes through libinfo to mDNSResponder — one system-wide daemon,
492            // and a serialisation point — measured at 8.25 seconds with ~100 processes
493            // asking at once. That is longer than the 10s budget above leaves for the rest
494            // of the query, against a server that is up and idle.
495            //
496            // This is the same override `crate::llm::ollama_client::client_for_endpoint`
497            // applies; it cannot be reused directly because this builder needs rustls, the
498            // caller's trust settings and a timeout. `host_of` does the parsing, including
499            // the trap it exists for: stripping only the scheme leaves `127.0.0.1:8443`,
500            // which does not parse as an `IpAddr`, so the bypass silently would not engage.
501            // A hostname is left alone — resolving it is the resolver's job.
502            if let Ok(ip) = host_for_build.parse::<std::net::IpAddr>() {
503                builder = builder.resolve(&host_for_build, std::net::SocketAddr::new(ip, 0));
504            }
505
506            if let Some(pem) = pem_for_build {
507                let cert = reqwest::Certificate::from_pem(pem.as_bytes())
508                    .context("ca_cert_pem is not a valid PEM certificate")?;
509                builder = builder.add_root_certificate(cert);
510            }
511            if insecure {
512                builder = builder
513                    .danger_accept_invalid_certs(true)
514                    // Do not load the platform root store when no certificate is going to
515                    // be checked against it. On macOS that load reads the system keychain
516                    // through Security.framework, which serialises across processes: with
517                    // a hundred netget processes starting at once it dominates the time to
518                    // build the client, and it buys exactly nothing here.
519                    .tls_built_in_root_certs(false);
520            }
521            builder.build().context("Failed to build DoH HTTP client")
522        })
523        .await
524        .context("DoH HTTP client build task panicked")??;
525        if let Ok(mut guard) = cache.lock() {
526            guard.insert(key, client.clone());
527        }
528        Ok(client)
529    }
530
531    async fn apply_action(
532        action_result: ClientActionResult,
533        client_id: ClientId,
534        server_url: &str,
535        app_state: &Arc<AppState>,
536    ) -> Result<Applied> {
537        match action_result {
538            ClientActionResult::Custom { name, data } if name == "dns_query" => {
539                let domain = data["domain"].as_str().unwrap_or("example.com").to_string();
540                let record_type = data["record_type"].as_str().unwrap_or("A").to_string();
541                let use_get = data["use_get"].as_bool().unwrap_or(false);
542
543                info!(
544                    "DoH client {} querying {} (type: {})",
545                    client_id, domain, record_type
546                );
547
548                let event_data = Self::perform_query(
549                    server_url,
550                    &domain,
551                    &record_type,
552                    use_get,
553                    client_id,
554                    app_state,
555                )
556                .await?;
557
558                let answer_count = event_data
559                    .get("answers")
560                    .and_then(|a| a.as_array())
561                    .map(|a| a.len())
562                    .unwrap_or(0);
563                let status = event_data
564                    .get("status")
565                    .and_then(|s| s.as_str())
566                    .unwrap_or("?");
567
568                Ok(Applied::Queried {
569                    detail: format!(
570                        "query_dns {} {} -> {} answers ({})",
571                        domain, record_type, answer_count, status
572                    ),
573                    event_data,
574                })
575            }
576            ClientActionResult::Disconnect => Ok(Applied::Disconnect),
577            ClientActionResult::WaitForMore => Ok(Applied::WaitForMore),
578            ClientActionResult::NoAction => Ok(Applied::Ignored("no_action".to_string())),
579            // Not swallowed: an action this client cannot carry out says so, rather
580            // than looking identical to success.
581            ClientActionResult::Custom { name, .. } => Ok(Applied::Ignored(format!(
582                "custom result '{name}' is not handled by the DoH client"
583            ))),
584            ClientActionResult::SendData(_) => Ok(Applied::Ignored(
585                "send_data has no meaning for a DoH client".to_string(),
586            )),
587            ClientActionResult::Multiple(_) => Ok(Applied::Ignored(
588                "multiple results are not produced by the DoH client".to_string(),
589            )),
590        }
591    }
592
593    /// Run one DNS query over HTTPS and return the `doh_response_received` payload.
594    ///
595    /// No LLM involvement, so a caller can await this and know exactly what the
596    /// resolver answered.
597    pub async fn perform_query(
598        server_url: &str,
599        domain: &str,
600        record_type: &str,
601        use_get: bool,
602        client_id: ClientId,
603        app_state: &Arc<AppState>,
604    ) -> Result<serde_json::Value> {
605        // Parse domain name
606        let name =
607            Name::from_str(domain).with_context(|| format!("Invalid domain name: {}", domain))?;
608
609        // Parse record type. An unknown type falls back to A rather than failing the
610        // whole query, which is what the LLM path has always done.
611        let rtype = match RecordType::from_str(record_type) {
612            Ok(rt) => rt,
613            Err(_) => {
614                warn!(
615                    "DoH: invalid record type {}, falling back to A",
616                    record_type
617                );
618                RecordType::A
619            }
620        };
621
622        // Build DNS query message
623        let mut query_msg = Message::new();
624        query_msg.set_id(rand::random());
625        query_msg.set_recursion_desired(true);
626        query_msg.add_query(Query::query(name, rtype));
627
628        // Encode query to DNS wire format
629        let query_bytes = query_msg.to_vec().context("Failed to encode DNS query")?;
630
631        // Make HTTPS request to DoH server.
632        //
633        // Built per query from the client's stored trust settings rather than with
634        // `Client::new()`. With the default system roots a DoH client cannot reach a
635        // server presenting a private-CA or self-signed certificate -- including NetGet's
636        // own DoH server, which is TLS-only and serves a self-signed cert, so the two
637        // halves of this codebase could not talk to each other at all.
638        let http_client = Self::build_http_client(client_id, server_url, app_state).await?;
639        let response = if use_get {
640            // GET method with base64url-encoded query
641            use base64::Engine as _;
642            let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&query_bytes);
643            http_client
644                .get(format!("{}?dns={}", server_url, encoded))
645                .header("Accept", "application/dns-message")
646                .send()
647                .await
648                .context("DoH GET request failed")?
649        } else {
650            // POST method with DNS query in body
651            http_client
652                .post(server_url)
653                .header("Content-Type", "application/dns-message")
654                .header("Accept", "application/dns-message")
655                .body(query_bytes)
656                .send()
657                .await
658                .context("DoH POST request failed")?
659        };
660
661        if !response.status().is_success() {
662            return Err(anyhow::anyhow!(
663                "DoH server answered with HTTP {}",
664                response.status()
665            ));
666        }
667
668        let response_bytes = response
669            .bytes()
670            .await
671            .context("Failed to read DoH response body")?;
672
673        let response_msg =
674            Message::from_vec(&response_bytes).context("Failed to parse DNS response")?;
675        trace!("DoH response: {:?}", response_msg);
676
677        let answers: Vec<serde_json::Value> = response_msg
678            .answers()
679            .iter()
680            .map(|record| {
681                serde_json::json!({
682                    "name": record.name().to_string(),
683                    "type": record.record_type().to_string(),
684                    "ttl": record.ttl(),
685                    "data": record.data().map(|d| d.to_string()).unwrap_or_default(),
686                })
687            })
688            .collect();
689
690        let status = format!("{:?}", response_msg.response_code());
691
692        Ok(serde_json::json!({
693            "query_id": response_msg.id(),
694            "domain": domain,
695            "query_type": record_type,
696            "answers": answers,
697            "status": status,
698        }))
699    }
700
701    /// Hand a completed query to the LLM as a `doh_response_received` event, then run
702    /// whatever it asks for next.
703    async fn notify_response(
704        client_id: ClientId,
705        event_data: serde_json::Value,
706        app_state: Arc<AppState>,
707        llm_client: OllamaClient,
708        status_tx: mpsc::UnboundedSender<String>,
709        protocol: Arc<DohClientProtocol>,
710        server_url: String,
711    ) {
712        let response_event = Event::new(&DOH_CLIENT_RESPONSE_RECEIVED_EVENT, event_data);
713
714        let instruction = app_state
715            .get_instruction_for_client(client_id)
716            .await
717            .unwrap_or_default();
718        let memory = app_state
719            .get_memory_for_client(client_id)
720            .await
721            .unwrap_or_default();
722
723        match call_llm_for_client(
724            &llm_client,
725            &app_state,
726            client_id.to_string(),
727            &instruction,
728            &memory,
729            Some(&response_event),
730            protocol.as_ref(),
731            &status_tx,
732        )
733        .await
734        {
735            Ok(next_llm_result) => {
736                // Boxed: execute_llm_actions calls back into this function, so the
737                // two futures are mutually recursive and need an indirection.
738                Box::pin(Self::execute_llm_actions(
739                    client_id,
740                    next_llm_result,
741                    app_state.clone(),
742                    llm_client.clone(),
743                    status_tx.clone(),
744                    protocol.clone(),
745                    server_url,
746                ))
747                .await;
748            }
749            Err(e) => {
750                error!(
751                    "DoH client {} LLM call after response failed: {}",
752                    client_id, e
753                );
754            }
755        }
756    }
757}