1pub 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
27const REMOVAL_CHECK_INTERVAL: Duration = Duration::from_secs(5);
30
31enum Applied {
34 Queried {
38 detail: String,
39 event_data: serde_json::Value,
40 },
41 Disconnect,
43 WaitForMore,
45 Ignored(String),
47}
48
49pub struct DohClient;
51
52impl DohClient {
53 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 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 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 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 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 let connected_event = Event::new(
134 &DOH_CLIENT_CONNECTED_EVENT,
135 serde_json::json!({
136 "server_url": server_url,
137 }),
138 );
139
140 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 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 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 Ok("0.0.0.0:0".parse().unwrap())
191 }
192
193 #[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; 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 #[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 Ok(Applied::Queried { detail, event_data }) => {
268 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 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 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 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 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 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 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 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 .timeout(std::time::Duration::from_secs(10));
485
486 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 .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 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 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 let name =
607 Name::from_str(domain).with_context(|| format!("Invalid domain name: {}", domain))?;
608
609 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 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 let query_bytes = query_msg.to_vec().context("Failed to encode DNS query")?;
630
631 let http_client = Self::build_http_client(client_id, server_url, app_state).await?;
639 let response = if use_get {
640 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 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 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 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}