Skip to main content

netget/client/dns/
mod.rs

1//! DNS client implementation
2pub mod actions;
3
4pub use actions::DnsClientProtocol;
5
6use anyhow::{Context, Result};
7use std::future::Future;
8use std::net::SocketAddr;
9use std::pin::Pin;
10use std::sync::Arc;
11use tokio::sync::{mpsc, Mutex};
12use tracing::{debug, error, info, trace};
13
14use crate::client::dns::actions::{
15    DNS_CLIENT_CONNECTED_EVENT, DNS_CLIENT_RESPONSE_RECEIVED_EVENT, DNS_CLIENT_ZONE_TRANSFER_EVENT,
16    MAX_REPORTED_RECORDS, MAX_TRANSFER_BYTES, MAX_TRANSFER_RECORDS, TRANSFER_TIMEOUT,
17};
18use crate::client::llm_budget::call_llm_for_client;
19use crate::llm::actions::client_trait::Client;
20use crate::llm::ollama_client::OllamaClient;
21use crate::llm::ClientLlmResult;
22use crate::protocol::{Event, EventType};
23use crate::state::app_state::AppState;
24use crate::state::{ClientId, ClientStatus};
25
26use hickory_client::client::AsyncClient;
27use hickory_client::rr::{Name, RecordType};
28use hickory_client::udp::UdpClientStream;
29use hickory_proto::op::{Edns, Message, MessageType, OpCode, Query, ResponseCode};
30use hickory_proto::rr::dnssec::rdata::DNSSECRData;
31use hickory_proto::rr::rdata::SOA;
32use hickory_proto::rr::{RData, Record};
33use hickory_proto::xfer::{DnsHandle, DnsRequest, DnsRequestOptions, FirstAnswer};
34
35/// DNS client that connects to a DNS server
36pub struct DnsClient;
37
38impl DnsClient {
39    /// Connect to a DNS server with integrated LLM actions
40    pub async fn connect_with_llm_actions(
41        remote_addr: String,
42        llm_client: OllamaClient,
43        app_state: Arc<AppState>,
44        status_tx: mpsc::UnboundedSender<String>,
45        client_id: ClientId,
46    ) -> Result<SocketAddr> {
47        // Parse the DNS server address
48        let dns_server: SocketAddr = remote_addr
49            .parse()
50            .context(format!("Invalid DNS server address: {}", remote_addr))?;
51
52        info!("DNS client {} connecting to {}", client_id, dns_server);
53
54        // Create UDP client stream
55        let stream = UdpClientStream::<tokio::net::UdpSocket>::new(dns_server);
56        let (client, bg) = AsyncClient::connect(stream)
57            .await
58            .context("Failed to create DNS client")?;
59
60        // Spawn the hickory transport driver. Registered so that stopping the client
61        // tears the UDP socket down instead of leaving it running detached.
62        let bg_handle = tokio::spawn(async move {
63            if let Err(e) = bg.await {
64                error!("DNS client {} transport driver stopped: {}", client_id, e);
65            }
66        });
67        app_state.register_client_task(client_id, bg_handle).await;
68
69        // Get local address (best effort)
70        let local_addr: SocketAddr = "0.0.0.0:0".parse()?;
71
72        info!("DNS client {} connected to {}", client_id, dns_server);
73
74        // Update client state
75        app_state
76            .update_client_status(client_id, ClientStatus::Connected)
77            .await;
78        let _ = status_tx.send(format!(
79            "[CLIENT] DNS client {} connected to {}",
80            client_id, dns_server
81        ));
82        let _ = status_tx.send("__UPDATE_UI__".to_string());
83
84        // Wrap the hickory client in Arc<Mutex> for sharing across tasks. It lives out here,
85        // not inside the conversation task, because the command loop needs the same handle:
86        // an injected `send_dns_query` has to go through the one connected resolver, not a
87        // second one of its own.
88        let client_arc = Arc::new(Mutex::new(client));
89        let protocol = Arc::new(DnsClientProtocol::new());
90
91        // Command channel for injected actions (the dashboard's [ send ]).
92        //
93        // Registered BEFORE the connected-event LLM call below: a dashboard-created client
94        // defaults to a `*` -> manual routing rule, so that call can park for minutes waiting
95        // for a human, and [ send ] must work for the whole park.
96        let command_rx =
97            crate::client::command_support::register_command_channel(&app_state, client_id).await;
98
99        let cmd_task = tokio::spawn(Self::command_loop(
100            command_rx,
101            client_arc.clone(),
102            protocol.clone(),
103            client_id,
104            app_state.clone(),
105            llm_client.clone(),
106            status_tx.clone(),
107        ));
108        app_state.register_client_task(client_id, cmd_task).await;
109
110        // Drive the LLM conversation in a registered background task.
111        //
112        // It used to run inline here, so `connect()` did not return until the LLM
113        // stopped asking for queries โ€” and an LLM that never stops asking meant
114        // `connect()` never returned. Running it as a tracked task means the caller
115        // gets its address immediately and `stop_client` can abort the conversation.
116        let conversation_state = app_state.clone();
117        let conversation_llm = llm_client.clone();
118        let conversation_tx = status_tx.clone();
119        let conversation_client = client_arc.clone();
120        let conversation_protocol = protocol.clone();
121        let handle = tokio::spawn(async move {
122            let app_state = conversation_state;
123            let llm_client = conversation_llm;
124            let status_tx = conversation_tx;
125            let client_arc = conversation_client;
126            let protocol = conversation_protocol;
127
128            let Some(instruction) = app_state.get_instruction_for_client(client_id).await else {
129                debug!(
130                    "DNS client {} has no instruction; nothing to drive",
131                    client_id
132                );
133                return;
134            };
135
136            let event = Event::new(
137                &DNS_CLIENT_CONNECTED_EVENT,
138                serde_json::json!({
139                    "remote_addr": dns_server.to_string(),
140                }),
141            );
142
143            let memory = app_state
144                .get_memory_for_client(client_id)
145                .await
146                .unwrap_or_default();
147
148            match call_llm_for_client(
149                &llm_client,
150                &app_state,
151                client_id.to_string(),
152                &instruction,
153                &memory,
154                Some(&event),
155                protocol.as_ref(),
156                &status_tx,
157            )
158            .await
159            {
160                Ok(ClientLlmResult {
161                    actions,
162                    memory_updates,
163                }) => {
164                    // Update memory
165                    if let Some(mem) = memory_updates {
166                        app_state.set_memory_for_client(client_id, mem).await;
167                    }
168
169                    // Drain the initial actions and every follow-up they produce.
170                    // Iterative, not recursive โ€” see run_dns_actions.
171                    Self::run_dns_actions(
172                        &client_arc,
173                        &protocol,
174                        actions,
175                        client_id,
176                        &app_state,
177                        &llm_client,
178                        &status_tx,
179                    )
180                    .await;
181                }
182                Err(e) => {
183                    error!("LLM error for DNS client {}: {}", client_id, e);
184                }
185            }
186
187            debug!("DNS client {} conversation task finished", client_id);
188        });
189        app_state.register_client_task(client_id, handle).await;
190
191        Ok(local_addr)
192    }
193
194    /// Execute a batch of DNS actions from the LLM, plus every follow-up action the
195    /// LLM produces in response.
196    ///
197    /// This used to be self-recursive: `execute_dns_action` awaited the LLM, then
198    /// called itself for each follow-up action. Because each level is a separately
199    /// polled boxed future, stack depth grew with the number of round-trips โ€” a
200    /// non-converging LLM loop (see `tests/client/dns`) overflowed the stack after a
201    /// couple of hundred queries and took the whole process down, rather than just
202    /// stalling the one client.
203    ///
204    /// It is now an explicit work queue: follow-up actions are pushed onto `pending`
205    /// and drained iteratively, so stack depth is constant no matter how many
206    /// query/response rounds occur. Convergence itself is enforced separately by the
207    /// per-client LLM call budget (`crate::client::llm_budget`), which makes
208    /// `call_llm_for_client` start failing once the ceiling is hit and drains this
209    /// queue.
210    async fn run_dns_actions(
211        client: &Arc<Mutex<AsyncClient>>,
212        protocol: &Arc<DnsClientProtocol>,
213        initial_actions: Vec<serde_json::Value>,
214        client_id: ClientId,
215        app_state: &Arc<AppState>,
216        llm_client: &OllamaClient,
217        status_tx: &mpsc::UnboundedSender<String>,
218    ) {
219        let mut pending: std::collections::VecDeque<serde_json::Value> =
220            initial_actions.into_iter().collect();
221
222        while let Some(action) = pending.pop_front() {
223            match Self::execute_dns_action(
224                client, protocol, action, client_id, app_state, llm_client, status_tx,
225            )
226            .await
227            {
228                Ok(follow_ups) => pending.extend(follow_ups),
229                Err(e) => error!("DNS client {} action error: {}", client_id, e),
230            }
231        }
232    }
233
234    /// Execute a single DNS action from the LLM.
235    ///
236    /// Returns any follow-up actions the LLM asked for; the caller
237    /// ([`Self::run_dns_actions`]) queues them. This function never calls itself.
238    #[allow(clippy::too_many_arguments)]
239    fn execute_dns_action<'a>(
240        client: &'a Arc<Mutex<AsyncClient>>,
241        protocol: &'a Arc<DnsClientProtocol>,
242        action: serde_json::Value,
243        client_id: ClientId,
244        app_state: &'a Arc<AppState>,
245        llm_client: &'a OllamaClient,
246        status_tx: &'a mpsc::UnboundedSender<String>,
247    ) -> Pin<Box<dyn Future<Output = Result<Vec<serde_json::Value>>> + Send + 'a>> {
248        Box::pin(async move {
249            match Self::apply_dns_action(client, protocol, action, client_id, app_state, status_tx)
250                .await?
251            {
252                DnsApplied::Queried {
253                    event_data, event, ..
254                } => Ok(Self::report_dns_response(
255                    event_data, event, protocol, client_id, app_state, llm_client, status_tx,
256                )
257                .await),
258                DnsApplied::Other(_) => Ok(Vec::new()),
259            }
260        })
261    }
262
263    /// Run one DNS action against the live hickory client, without involving the LLM.
264    ///
265    /// Split out from [`Self::execute_dns_action`] so the injected-command path can answer
266    /// its caller as soon as the resolver replies, and hand the model its
267    /// `dns_response_received` event afterwards โ€” a manual routing rule parked on that event
268    /// would otherwise hold the dashboard's \[send\] open for its whole timeout.
269    async fn apply_dns_action(
270        client: &Arc<Mutex<AsyncClient>>,
271        protocol: &Arc<DnsClientProtocol>,
272        action: serde_json::Value,
273        client_id: ClientId,
274        app_state: &Arc<AppState>,
275        status_tx: &mpsc::UnboundedSender<String>,
276    ) -> Result<DnsApplied> {
277        match protocol.execute_action(action)? {
278            crate::llm::actions::client_trait::ClientActionResult::Custom { name, data }
279                if name == "dns_query" =>
280            {
281                let domain = data
282                    .get("domain")
283                    .and_then(|v| v.as_str())
284                    .context("Missing domain in query")?;
285                let query_type_str = data
286                    .get("query_type")
287                    .and_then(|v| v.as_str())
288                    .context("Missing query_type")?;
289                let recursion_desired = data
290                    .get("recursion_desired")
291                    .and_then(|v| v.as_bool())
292                    .unwrap_or(true);
293                let dnssec = data
294                    .get("dnssec")
295                    .and_then(|v| v.as_bool())
296                    .unwrap_or(false);
297                let name =
298                    Name::from_utf8(domain).context(format!("Invalid domain name: {}", domain))?;
299                let record_type = Self::parse_record_type(query_type_str)?;
300                debug!(
301                    "DNS client {} querying {} for {} record (rd={recursion_desired}, \
302                     do={dnssec})",
303                    client_id, domain, query_type_str
304                );
305
306                // Built by hand rather than with `ClientHandle::query`, which fixes RD on and
307                // offers no EDNS: `recursion_desired` and `dnssec` both have to reach the wire.
308                let mut message = Message::new();
309                message
310                    .set_id(rand::random::<u16>())
311                    .set_message_type(MessageType::Query)
312                    .set_op_code(OpCode::Query)
313                    .set_recursion_desired(recursion_desired)
314                    .add_query(Query::query(name.clone(), record_type));
315                let mut edns = Edns::new();
316                edns.set_max_payload(1232)
317                    .set_version(0)
318                    .set_dnssec_ok(dnssec);
319                message.set_edns(edns);
320                let request = DnsRequest::new(message, DnsRequestOptions::default());
321                let response = client.lock().await.send(request).first_answer().await;
322
323                match response {
324                    Ok(response) => {
325                        let response_code = response.response_code();
326                        let answers: Vec<serde_json::Value> =
327                            response.answers().iter().filter_map(record_json).collect();
328                        let authority: Vec<serde_json::Value> = response
329                            .name_servers()
330                            .iter()
331                            .filter_map(record_json)
332                            .collect();
333                        trace!(
334                            "DNS client {} received response: {} answers, code: {:?}",
335                            client_id,
336                            answers.len(),
337                            response_code
338                        );
339                        let detail = format!(
340                            "dns_query {domain} {query_type_str} -> {} with {} answer(s) \
341                                 (hickory owns the wire encoding; no byte count is observable)",
342                            Self::response_code_to_string(response_code),
343                            answers.len()
344                        );
345                        let mut event_data = serde_json::json!({
346                            "query_id": response.id(),
347                            "domain": domain,
348                            "query_type": query_type_str,
349                            "answers": answers,
350                            "response_code": Self::response_code_to_string(response_code),
351                            "authenticated_data": response.authentic_data(),
352                            "truncated": response.truncated(),
353                        });
354                        if !authority.is_empty() {
355                            event_data["authority"] = serde_json::json!(authority);
356                        }
357                        Ok(DnsApplied::Queried {
358                            event_data,
359                            detail,
360                            event: &DNS_CLIENT_RESPONSE_RECEIVED_EVENT,
361                        })
362                    }
363                    Err(e) => {
364                        error!("DNS client {} query error: {}", client_id, e);
365                        Err(anyhow::anyhow!("DNS query failed: {}", e))
366                    }
367                }
368            }
369            crate::llm::actions::client_trait::ClientActionResult::Custom { name, data }
370                if name == "dns_zone_transfer" =>
371            {
372                let zone = data
373                    .get("zone")
374                    .and_then(|v| v.as_str())
375                    .context("Missing zone")?;
376                let serial = data
377                    .get("serial")
378                    .and_then(|v| v.as_u64())
379                    .map(|s| s as u32);
380                let server: SocketAddr = app_state
381                    .get_client(client_id)
382                    .await
383                    .context("the DNS client is gone")?
384                    .remote_addr
385                    .parse()
386                    .context("the DNS client's server address")?;
387                let event_data =
388                    Self::zone_transfer(server, zone, serial, client_id, app_state).await?;
389                let detail = format!(
390                    "{} {zone}: {} record(s) in {} message(s)",
391                    event_data["transfer_type"].as_str().unwrap_or("XFR"),
392                    event_data["record_count"],
393                    event_data["messages"]
394                );
395                Ok(DnsApplied::Queried {
396                    event_data,
397                    detail,
398                    event: &DNS_CLIENT_ZONE_TRANSFER_EVENT,
399                })
400            }
401            crate::llm::actions::client_trait::ClientActionResult::Disconnect => {
402                info!("DNS client {} disconnecting", client_id);
403                app_state
404                    .update_client_status(client_id, ClientStatus::Disconnected)
405                    .await;
406                let _ = status_tx.send("__UPDATE_UI__".to_string());
407                Ok(DnsApplied::Other("disconnect".to_string()))
408            }
409            crate::llm::actions::client_trait::ClientActionResult::WaitForMore => {
410                debug!("DNS client {} waiting for more", client_id);
411                Ok(DnsApplied::Other(
412                    "wait_for_more: no query sent".to_string(),
413                ))
414            }
415            other => {
416                // Other action results not applicable to DNS
417                Ok(DnsApplied::Other(format!(
418                    "{other:?} is not a DNS client verb; no query sent"
419                )))
420            }
421        }
422    }
423
424    /// Hand one completed query's answers to the model as a `dns_response_received` event
425    /// and return whatever follow-up actions it asked for.
426    async fn report_dns_response(
427        event_data: serde_json::Value,
428        event_type: &'static EventType,
429        protocol: &Arc<DnsClientProtocol>,
430        client_id: ClientId,
431        app_state: &Arc<AppState>,
432        llm_client: &OllamaClient,
433        status_tx: &mpsc::UnboundedSender<String>,
434    ) -> Vec<serde_json::Value> {
435        let Some(instruction) = app_state.get_instruction_for_client(client_id).await else {
436            return Vec::new();
437        };
438
439        let event = Event::new(event_type, event_data);
440
441        let memory = app_state
442            .get_memory_for_client(client_id)
443            .await
444            .unwrap_or_default();
445
446        match call_llm_for_client(
447            llm_client,
448            app_state,
449            client_id.to_string(),
450            &instruction,
451            &memory,
452            Some(&event),
453            protocol.as_ref(),
454            status_tx,
455        )
456        .await
457        {
458            Ok(ClientLlmResult {
459                actions,
460                memory_updates,
461            }) => {
462                // Update memory
463                if let Some(mem) = memory_updates {
464                    app_state.set_memory_for_client(client_id, mem).await;
465                }
466
467                // Hand follow-up actions back to the caller's work queue rather than
468                // recursing into them here โ€” recursion here is what overflowed the stack.
469                actions
470            }
471            Err(e) => {
472                error!("LLM error for DNS client {}: {}", client_id, e);
473                Vec::new()
474            }
475        }
476    }
477
478    /// Drain injected commands until the channel closes (the client was removed) or an
479    /// injected `disconnect` ends the session.
480    ///
481    /// A query reports `Executed`, never `Sent`: hickory owns the wire encoding and the
482    /// socket, and reports no byte count, so there is no truthful number to hand back. The
483    /// detail string carries the response code and answer count instead, which is what the
484    /// caller actually wants to see.
485    async fn command_loop(
486        mut command_rx: mpsc::Receiver<crate::state::client_handles::ClientCommand>,
487        client: Arc<Mutex<AsyncClient>>,
488        protocol: Arc<DnsClientProtocol>,
489        client_id: ClientId,
490        app_state: Arc<AppState>,
491        llm_client: OllamaClient,
492        status_tx: mpsc::UnboundedSender<String>,
493    ) {
494        use crate::llm::actions::protocol_trait::Protocol;
495        use crate::state::client_handles::ClientSendOutcome;
496        use crate::state::AccessLogOwner;
497
498        while let Some(command) = command_rx.recv().await {
499            let action = command.action.clone();
500            let is_disconnect = action.get("type").and_then(|v| v.as_str()) == Some("disconnect");
501
502            // A completed query is reported to the model only after the caller has been
503            // answered, so `queried` carries the event payload past the reply.
504            let mut queried: Option<(serde_json::Value, &'static EventType)> = None;
505            let outcome = match Self::apply_dns_action(
506                &client,
507                &protocol,
508                action.clone(),
509                client_id,
510                &app_state,
511                &status_tx,
512            )
513            .await
514            {
515                Ok(DnsApplied::Queried {
516                    event_data,
517                    detail,
518                    event,
519                }) => {
520                    queried = Some((event_data, event));
521                    Ok(ClientSendOutcome::Executed { detail })
522                }
523                Ok(DnsApplied::Other(_)) if is_disconnect => Ok(ClientSendOutcome::Disconnected),
524                Ok(DnsApplied::Other(detail)) => Ok(ClientSendOutcome::Executed { detail }),
525                // `execute_action` failing and the query itself failing are different
526                // things: the first is the model naming a verb this client does not have,
527                // the second is a network fault. Only the first is a Rejected.
528                Err(e) if e.to_string().starts_with("Unknown DNS client action") => {
529                    Ok(ClientSendOutcome::Rejected {
530                        error: e.to_string(),
531                    })
532                }
533                Err(e) => Err(e),
534            };
535
536            let outcome_json = match &outcome {
537                Ok(outcome) => serde_json::to_value(outcome).unwrap_or(serde_json::Value::Null),
538                Err(e) => serde_json::json!({"error": e.to_string()}),
539            };
540            app_state
541                .record_access_log(
542                    AccessLogOwner::Client(client_id.as_u32()),
543                    protocol.protocol_name(),
544                    None,
545                    "injected_action",
546                    action,
547                    vec![outcome_json],
548                )
549                .await;
550
551            let disconnect = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
552            if let Err(e) = &outcome {
553                error!("DNS client {} injected action failed: {}", client_id, e);
554                let _ = status_tx.send(format!(
555                    "[WARN] Client {} injected action failed: {}",
556                    client_id, e
557                ));
558            }
559            let _ = status_tx.send("__UPDATE_UI__".to_string());
560            crate::client::command_support::reply(command, outcome);
561
562            if disconnect {
563                break;
564            }
565
566            if let Some((event_data, event)) = queried {
567                let follow_ups = Self::report_dns_response(
568                    event_data,
569                    event,
570                    &protocol,
571                    client_id,
572                    &app_state,
573                    &llm_client,
574                    &status_tx,
575                )
576                .await;
577                Self::run_dns_actions(
578                    &client,
579                    &protocol,
580                    follow_ups,
581                    client_id,
582                    &app_state,
583                    &llm_client,
584                    &status_tx,
585                )
586                .await;
587            }
588        }
589
590        app_state.remove_client_handle(client_id).await;
591        let _ = status_tx.send("__UPDATE_UI__".to_string());
592    }
593
594    /// Parse DNS record type from string
595    fn parse_record_type(type_str: &str) -> Result<RecordType> {
596        match type_str.to_uppercase().as_str() {
597            "A" => Ok(RecordType::A),
598            "AAAA" => Ok(RecordType::AAAA),
599            "ANAME" => Ok(RecordType::ANAME),
600            "CAA" => Ok(RecordType::CAA),
601            "CNAME" => Ok(RecordType::CNAME),
602            "MX" => Ok(RecordType::MX),
603            "NAPTR" => Ok(RecordType::NAPTR),
604            "NS" => Ok(RecordType::NS),
605            "OPENPGPKEY" => Ok(RecordType::OPENPGPKEY),
606            "PTR" => Ok(RecordType::PTR),
607            "SOA" => Ok(RecordType::SOA),
608            "SRV" => Ok(RecordType::SRV),
609            "SSHFP" => Ok(RecordType::SSHFP),
610            "TLSA" => Ok(RecordType::TLSA),
611            "TXT" => Ok(RecordType::TXT),
612            _ => Err(anyhow::anyhow!("Unsupported DNS record type: {}", type_str)),
613        }
614    }
615
616    /// Convert ResponseCode to string
617    fn response_code_to_string(code: ResponseCode) -> String {
618        match code {
619            ResponseCode::NoError => "NOERROR".to_string(),
620            ResponseCode::FormErr => "FORMERR".to_string(),
621            ResponseCode::ServFail => "SERVFAIL".to_string(),
622            ResponseCode::NXDomain => "NXDOMAIN".to_string(),
623            ResponseCode::NotImp => "NOTIMP".to_string(),
624            ResponseCode::Refused => "REFUSED".to_string(),
625            ResponseCode::YXDomain => "YXDOMAIN".to_string(),
626            ResponseCode::YXRRSet => "YXRRSET".to_string(),
627            ResponseCode::NXRRSet => "NXRRSET".to_string(),
628            ResponseCode::NotAuth => "NOTAUTH".to_string(),
629            ResponseCode::NotZone => "NOTZONE".to_string(),
630            _ => format!("{:?}", code),
631        }
632    }
633}
634
635impl DnsClient {
636    /// AXFR, or IXFR from `serial`, over a TCP connection of its own (the query client is UDP).
637    /// The records are counted and bounded at [`MAX_TRANSFER_RECORDS`], the messages' bytes at
638    /// [`MAX_TRANSFER_BYTES`], the whole exchange at [`TRANSFER_TIMEOUT`]. Only the first
639    /// [`MAX_REPORTED_RECORDS`] are kept and reported, so a large zone is a count and a sample
640    /// rather than an event the model cannot read, and memory does not grow with the zone.
641    ///
642    /// The messages are read straight off the socket, one length-prefixed message at a time,
643    /// rather than through hickory's multiplexer, which hands responses over with `try_send`
644    /// into a 32-slot channel and drops what does not fit: a transfer read that way could lose
645    /// messages and still look complete. Read here, a slow reader only slows the server down
646    /// (TCP backpressure), and [`XfrEnd`] decides where the transfer ends as hickory does.
647    async fn zone_transfer(
648        server: SocketAddr,
649        zone: &str,
650        serial: Option<u32>,
651        _client_id: ClientId,
652        _app_state: &Arc<AppState>,
653    ) -> Result<serde_json::Value> {
654        use tokio::io::{AsyncReadExt, AsyncWriteExt};
655
656        let origin = Name::from_utf8(zone).with_context(|| format!("invalid zone {zone:?}"))?;
657        // IXFR names the serial the client holds in an SOA; the other fields are not read.
658        let last_soa = serial.map(|s| SOA::new(origin.clone(), origin.clone(), s, 0, 0, 0, 0));
659        let transfer_type = if last_soa.is_some() { "IXFR" } else { "AXFR" };
660        let mut request =
661            hickory_proto::op::update_message::zone_transfer(origin.clone(), last_soa.clone());
662        let id: u16 = rand::random();
663        request.set_id(id);
664        let wire = request
665            .to_vec()
666            .context("cannot encode the transfer request")?;
667        let request_len =
668            u16::try_from(wire.len()).context("the transfer request does not fit one message")?;
669
670        // The first MAX_REPORTED_RECORDS records; `record_count` counts every one.
671        let mut records: Vec<Record> = Vec::new();
672        let mut record_count = 0usize;
673        let mut bytes_received = 0usize;
674        let mut messages = 0usize;
675        let mut rcode = ResponseCode::NoError;
676        let mut end = XfrEnd::new(last_soa.is_some());
677        let exchange = async {
678            let mut stream = tokio::net::TcpStream::connect(server)
679                .await
680                .with_context(|| {
681                    format!("cannot open a TCP connection to {server} for the transfer")
682                })?;
683            stream.write_all(&request_len.to_be_bytes()).await?;
684            stream.write_all(&wire).await?;
685            loop {
686                let mut len = [0u8; 2];
687                match stream.read_exact(&mut len).await {
688                    Ok(_) => {}
689                    Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => anyhow::bail!(
690                        "{transfer_type} of {zone}: the server closed the connection before the \
691                         transfer ended"
692                    ),
693                    Err(e) => return Err(e.into()),
694                }
695                let len = usize::from(u16::from_be_bytes(len));
696                bytes_received = bytes_received.saturating_add(len + 2);
697                if bytes_received > MAX_TRANSFER_BYTES {
698                    anyhow::bail!(
699                        "{transfer_type} of {zone} passed {MAX_TRANSFER_BYTES} bytes; abandoned"
700                    );
701                }
702                let mut body = vec![0u8; len];
703                stream.read_exact(&mut body).await?;
704                let response = Message::from_vec(&body)
705                    .map_err(|e| anyhow::anyhow!("{transfer_type} failed: {e}"))?;
706                if response.id() != id {
707                    anyhow::bail!(
708                        "{transfer_type} of {zone}: a message answered id {} where {id} was asked",
709                        response.id()
710                    );
711                }
712                messages += 1;
713                rcode = response.response_code();
714                if rcode != ResponseCode::NoError {
715                    break;
716                }
717                let answers = response.answers();
718                record_count = record_count.saturating_add(answers.len());
719                let room = MAX_REPORTED_RECORDS.saturating_sub(records.len());
720                records.extend(answers.iter().take(room).cloned());
721                if record_count > MAX_TRANSFER_RECORDS {
722                    anyhow::bail!(
723                        "{transfer_type} of {zone} passed {MAX_TRANSFER_RECORDS} records; abandoned"
724                    );
725                }
726                end.process(answers)
727                    .map_err(|e| anyhow::anyhow!("{transfer_type} of {zone}: {e}"))?;
728                if end.ended() {
729                    break;
730                }
731            }
732            Ok::<(), anyhow::Error>(())
733        };
734        tokio::time::timeout(TRANSFER_TIMEOUT, exchange)
735            .await
736            .map_err(|_| {
737                anyhow::anyhow!(
738                    "{transfer_type} of {zone} did not finish within {}s",
739                    TRANSFER_TIMEOUT.as_secs()
740                )
741            })??;
742
743        let serial_now = records.first().and_then(|r| match r.data() {
744            Some(RData::SOA(s)) => Some(s.serial()),
745            _ => None,
746        });
747        // RFC 1995 ยง4: an incremental answer's second record is an SOA (the old serial); a
748        // full one's is not. A single SOA says the client is current.
749        let form = match (transfer_type, record_count) {
750            (_, 0) => "none",
751            ("IXFR", 1) => "current",
752            ("IXFR", _) if records.get(1).map(|r| r.record_type()) == Some(RecordType::SOA) => {
753                "incremental"
754            }
755            _ => "full",
756        };
757        let reported: Vec<serde_json::Value> = records.iter().filter_map(record_json).collect();
758        Ok(serde_json::json!({
759            "zone": origin.to_string(),
760            "transfer_type": transfer_type,
761            "form": form,
762            "response_code": Self::response_code_to_string(rcode),
763            "serial": serial_now,
764            "record_count": record_count,
765            "messages": messages,
766            "records": reported,
767            "truncated": record_count > MAX_REPORTED_RECORDS,
768        }))
769    }
770}
771
772/// Where an AXFR or IXFR ends, decided from the answers in order: hickory-client 0.24's
773/// `ClientStreamXfrState`, without the stream it wraps. The first record's SOA serial is the
774/// one the transfer closes on. An AXFR ends at the next SOA; a second SOA with the same serial
775/// is an empty zone. An IXFR (RFC 1995) alternates old and new SOAs and ends on an even count
776/// of them whose last carries the expected serial.
777#[derive(Debug)]
778pub enum XfrEnd {
779    Start { maybe_incr: bool },
780    Second { maybe_incr: bool, expected: u32 },
781    Axfr { expected: u32 },
782    Ixfr { even: bool, expected: u32 },
783    Ended,
784}
785
786impl XfrEnd {
787    pub fn new(maybe_incr: bool) -> Self {
788        XfrEnd::Start { maybe_incr }
789    }
790
791    pub fn ended(&self) -> bool {
792        matches!(self, XfrEnd::Ended)
793    }
794
795    /// Takes one message's answers in order; `Err` for a transfer that cannot be valid.
796    pub fn process(&mut self, answers: &[Record]) -> Result<(), &'static str> {
797        fn serial(r: &Record) -> Option<u32> {
798            r.data().and_then(RData::as_soa).map(SOA::serial)
799        }
800        if answers.is_empty() {
801            return Ok(());
802        }
803        match std::mem::replace(self, XfrEnd::Ended) {
804            XfrEnd::Start { maybe_incr } => match serial(&answers[0]) {
805                Some(expected) => {
806                    *self = XfrEnd::Second {
807                        maybe_incr,
808                        expected,
809                    };
810                    self.process(&answers[1..])
811                }
812                None => Err("the transfer does not start with the zone's SOA"),
813            },
814            XfrEnd::Second {
815                maybe_incr,
816                expected,
817            } => match serial(&answers[0]) {
818                Some(s) if s == expected => {
819                    if answers.len() == 1 {
820                        Ok(())
821                    } else {
822                        Err("invalid zone transfer, contains trailing records")
823                    }
824                }
825                Some(_) if maybe_incr => {
826                    *self = XfrEnd::Ixfr {
827                        even: true,
828                        expected,
829                    };
830                    self.process(&answers[1..])
831                }
832                Some(_) => Err("invalid zone transfer, expected AXFR, got IXFR"),
833                None => {
834                    *self = XfrEnd::Axfr { expected };
835                    self.process(&answers[1..])
836                }
837            },
838            XfrEnd::Axfr { expected } => {
839                let soas = answers
840                    .iter()
841                    .filter(|a| a.record_type() == RecordType::SOA)
842                    .count();
843                match soas {
844                    0 => {
845                        *self = XfrEnd::Axfr { expected };
846                        Ok(())
847                    }
848                    1 if answers.last().map(|r| r.record_type()) == Some(RecordType::SOA) => Ok(()),
849                    _ => Err("invalid zone transfer, contains trailing records"),
850                }
851            }
852            XfrEnd::Ixfr { even, expected } => {
853                let even = answers
854                    .iter()
855                    .fold(even, |even, a| even ^ (a.record_type() == RecordType::SOA));
856                let last_is_expected = answers.last().and_then(serial) == Some(expected);
857                if !(even && last_is_expected) {
858                    *self = XfrEnd::Ixfr { even, expected };
859                }
860                Ok(())
861            }
862            XfrEnd::Ended => Ok(()),
863        }
864    }
865}
866
867/// A record type's mnemonic, or RFC 3597's `TYPE<n>` for one hickory has no name for.
868fn type_name(t: RecordType) -> String {
869    match t {
870        RecordType::Unknown(n) => format!("TYPE{n}"),
871        other => other.to_string(),
872    }
873}
874
875/// One record as the model reads it. DNSSEC records are reported structurally โ€” what a
876/// signature covers and which key made it โ€” never as base64.
877fn record_json(record: &Record) -> Option<serde_json::Value> {
878    let data = record.data()?;
879    let mut v = match data {
880        RData::A(a) => serde_json::json!({"type": "A", "ip": a.to_string()}),
881        RData::AAAA(a) => serde_json::json!({"type": "AAAA", "ip": a.to_string()}),
882        RData::CNAME(c) => serde_json::json!({"type": "CNAME", "target": c.to_string()}),
883        RData::MX(mx) => serde_json::json!({
884            "type": "MX",
885            "exchange": mx.exchange().to_string(),
886            "preference": mx.preference(),
887        }),
888        RData::TXT(txt) => {
889            let text: Vec<String> = txt
890                .iter()
891                .map(|bytes| String::from_utf8_lossy(bytes).to_string())
892                .collect();
893            serde_json::json!({"type": "TXT", "text": text.join("")})
894        }
895        RData::NS(ns) => serde_json::json!({"type": "NS", "nameserver": ns.to_string()}),
896        RData::SOA(soa) => serde_json::json!({
897            "type": "SOA",
898            "mname": soa.mname().to_string(),
899            "rname": soa.rname().to_string(),
900            "serial": soa.serial(),
901            "refresh": soa.refresh(),
902            "retry": soa.retry(),
903            "expire": soa.expire(),
904            "minimum": soa.minimum(),
905        }),
906        RData::PTR(ptr) => serde_json::json!({"type": "PTR", "domain": ptr.to_string()}),
907        RData::SRV(srv) => serde_json::json!({
908            "type": "SRV",
909            "priority": srv.priority(),
910            "weight": srv.weight(),
911            "port": srv.port(),
912            "target": srv.target().to_string(),
913        }),
914        RData::DNSSEC(DNSSECRData::RRSIG(sig)) => serde_json::json!({
915            "type": "RRSIG",
916            "type_covered": type_name(sig.type_covered()),
917            "algorithm": u8::from(sig.algorithm()),
918            "key_tag": sig.key_tag(),
919            "signer": sig.signer_name().to_string(),
920            "inception": sig.sig_inception(),
921            "expiration": sig.sig_expiration(),
922        }),
923        RData::DNSSEC(DNSSECRData::DNSKEY(key)) => serde_json::json!({
924            "type": "DNSKEY",
925            "flags": key.flags(),
926            "algorithm": u8::from(key.algorithm()),
927            "key_tag": key.calculate_key_tag().ok(),
928            "secure_entry_point": key.secure_entry_point(),
929        }),
930        RData::DNSSEC(DNSSECRData::NSEC(nsec)) => serde_json::json!({
931            "type": "NSEC",
932            "next": nsec.next_domain_name().to_string(),
933            "types": nsec.type_bit_maps().iter().map(|t| type_name(*t)).collect::<Vec<_>>(),
934        }),
935        RData::DNSSEC(DNSSECRData::DS(ds)) => serde_json::json!({
936            "type": "DS",
937            "key_tag": ds.key_tag(),
938            "algorithm": u8::from(ds.algorithm()),
939            "digest_type": u8::from(ds.digest_type()),
940        }),
941        // A type this client does not read: named, sized, never shown as encoded bytes.
942        _ => serde_json::json!({
943            "type": type_name(record.record_type()),
944            "rdata_length": hickory_proto::serialize::binary::BinEncodable::to_bytes(data)
945                .map(|b| b.len())
946                .unwrap_or_default(),
947        }),
948    };
949    v["name"] = serde_json::json!(record.name().to_string());
950    v["ttl"] = serde_json::json!(record.ttl());
951    Some(v)
952}
953
954/// What [`DnsClient::apply_dns_action`] did with one action.
955enum DnsApplied {
956    /// A query completed. `event_data` is the `dns_response_received` payload the model
957    /// still has to be shown; `detail` summarises it for the injecting caller.
958    Queried {
959        event_data: serde_json::Value,
960        detail: String,
961        /// `dns_response_received` for a query, `dns_zone_transfer_received` for a transfer.
962        event: &'static EventType,
963    },
964    /// The action ran but sent no query; the string says why.
965    Other(String),
966}