From d970e03723b158a9aef0edd874ac9fbd11a78a5a Mon Sep 17 00:00:00 2001 From: Ximon Eighteen <3304436+ximon18@users.noreply.github.com> Date: Thu, 25 Jul 2024 02:01:15 +0200 Subject: [PATCH] Proof of concept AXFR (for the dig format only) and TSIG support. --- Cargo.lock | 9 ++-- Cargo.toml | 2 +- src/client.rs | 110 +++++++++++++++++++++++++++++++----------- src/commands/query.rs | 53 ++++++++++++++++++-- src/output/dig.rs | 8 +++ 5 files changed, 146 insertions(+), 36 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e0b9564..aa7b494 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -282,14 +282,14 @@ dependencies = [ [[package]] name = "domain" version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7eefe29e8dd614abbee51a1616654cab123c4c56850ab83f5b7f1e1f9977bf7c" +source = "git+https://github.com/NLnetLabs/domain?branch=xfr#c07c769024f466f95125958270ab40a75ca6f7ab" dependencies = [ "bytes", "futures-util", "moka", "octseq", "rand", + "ring", "smallvec", "time", "tokio", @@ -551,9 +551,8 @@ dependencies = [ [[package]] name = "octseq" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ed2eaec452d98ccc1c615dd843fd039d9445f2fb4da114ee7e6af5fcb68be98" +version = "0.5.2-dev" +source = "git+https://github.com/NLnetLabs/octseq.git?rev=3f7797f4274af0a52e66105250ee1186ff2ab6ac#3f7797f4274af0a52e66105250ee1186ff2ab6ac" dependencies = [ "bytes", "smallvec", diff --git a/Cargo.toml b/Cargo.toml index b296663..f9c9869 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ exclude = [ ".github", ".gitignore" ] bytes = "1" clap = { version = "4", features = ["derive", "unstable-doc"] } chrono = { version = "0.4.38", features = [ "alloc", "clock" ] } -domain = { version = "0.10", features = ["resolv", "unstable-client-transport"]} +domain = { version = "0.10.1", features = [ "resolv", "tsig", "unstable-client-transport" ], git = "https://github.com/NLnetLabs/domain", branch = "xfr" } tempfile = "3.1.0" tokio = { version = "1.33", features = ["rt-multi-thread"] } tokio-rustls = { version = "0.26.0", default-features = false, features = [ "ring", "logging", "tls12" ] } diff --git a/src/client.rs b/src/client.rs index f2d1093..22f0bf9 100644 --- a/src/client.rs +++ b/src/client.rs @@ -8,11 +8,15 @@ use domain::base::message_builder::MessageBuilder; use domain::base::name::ToName; use domain::base::question::Question; use domain::net::client::protocol::UdpConnect; -use domain::net::client::request::{RequestMessage, SendRequest}; -use domain::net::client::{dgram, stream}; +use domain::net::client::request::{ + ComposeRequest, GetResponse, RequestMessage, SendRequest, +}; +use domain::net::client::{dgram, stream, tsig, xfr}; use domain::resolv::stub::conf; +use domain::tsig::Key; use std::fmt; use std::net::SocketAddr; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; use tokio::net::TcpStream; @@ -53,6 +57,7 @@ impl Client { pub async fn query>>( &self, question: Q, + tsig_key: Option, ) -> Result { let mut res = MessageBuilder::new_vec(); @@ -62,16 +67,20 @@ impl Client { let mut res = res.question(); res.push(question.into()).unwrap(); - self.request(RequestMessage::new(res)).await + self.request(RequestMessage::new(res), tsig_key).await } pub async fn request( &self, request: RequestMessage>, + tsig_key: Option, ) -> Result { let mut servers = self.servers.as_slice(); while let Some((server, tail)) = servers.split_first() { - match self.request_server(request.clone(), server).await { + match self + .request_server(request.clone(), tsig_key.clone(), server) + .await + { Ok(answer) => return Ok(answer), Err(err) => { if tail.is_empty() { @@ -87,24 +96,56 @@ impl Client { pub async fn request_server( &self, request: RequestMessage>, + tsig_key: Option, server: &Server, ) -> Result { match server.transport { - Transport::Udp => self.request_udp(request, server).await, - Transport::UdpTcp => self.request_udptcp(request, server).await, - Transport::Tcp => self.request_tcp(request, server).await, - Transport::Tls => self.request_tls(request, server).await, + Transport::Udp => { + self.request_udp(request, tsig_key, server).await + } + Transport::UdpTcp => { + self.request_udptcp(request, tsig_key, server).await + } + Transport::Tcp => { + self.request_tcp(request, tsig_key, server).await + } + Transport::Tls => { + self.request_tls(request, tsig_key, server).await + } + } + } + + async fn finalize_request( + mut send_request: Box, + mut stats: Stats, + streaming: bool, + ) -> Result { + let mut msgs = Vec::with_capacity(1); + while !send_request.is_stream_complete() { + msgs.push(send_request.get_response().await?); + if !streaming { + break; + } } + stats.finalize(); + Ok(Answer { + msgs, + stats, + cur_idx: Default::default(), + }) } pub async fn request_udptcp( &self, request: RequestMessage>, + tsig_key: Option, server: &Server, ) -> Result { - let answer = self.request_udp(request.clone(), server).await?; - if answer.message.header().tc() { - self.request_tcp(request, server).await + let answer = self + .request_udp(request.clone(), tsig_key.clone(), server) + .await?; + if answer.message().header().tc() { + self.request_tcp(request, tsig_key, server).await } else { Ok(answer) } @@ -113,38 +154,45 @@ impl Client { pub async fn request_udp( &self, request: RequestMessage>, + tsig_key: Option, server: &Server, ) -> Result { - let mut stats = Stats::new(server.addr, Protocol::Udp); + let stats = Stats::new(server.addr, Protocol::Udp); let conn = dgram::Connection::with_config( UdpConnect::new(server.addr), Self::dgram_config(server), ); - let message = conn.send_request(request).get_response().await?; - stats.finalize(); - Ok(Answer { message, stats }) + let conn = + tsig::Connection::new(tsig_key, xfr::Connection::new(conn)); + let streaming = request.is_streaming(); + let send_request = conn.send_request(request); + Self::finalize_request(send_request, stats, streaming).await } pub async fn request_tcp( &self, request: RequestMessage>, + tsig_key: Option, server: &Server, ) -> Result { - let mut stats = Stats::new(server.addr, Protocol::Tcp); + let stats = Stats::new(server.addr, Protocol::Tcp); let socket = TcpStream::connect(server.addr).await?; let (conn, tran) = stream::Connection::with_config( socket, Self::stream_config(server), ); + let conn = + tsig::Connection::new(tsig_key, xfr::Connection::new(conn)); tokio::spawn(tran.run()); - let message = conn.send_request(request).get_response().await?; - stats.finalize(); - Ok(Answer { message, stats }) + let streaming = request.is_streaming(); + let send_request = conn.send_request(request); + Self::finalize_request(send_request, stats, streaming).await } pub async fn request_tls( &self, request: RequestMessage>, + tsig_key: Option, server: &Server, ) -> Result { let root_store = RootCertStore { @@ -156,7 +204,7 @@ impl Client { .with_no_client_auth(), ); - let mut stats = Stats::new(server.addr, Protocol::Tls); + let stats = Stats::new(server.addr, Protocol::Tls); let tcp_socket = TcpStream::connect(server.addr).await?; let tls_connector = tokio_rustls::TlsConnector::from(client_config); let server_name = server @@ -174,10 +222,12 @@ impl Client { tls_socket, Self::stream_config(server), ); + let conn = + tsig::Connection::new(tsig_key, xfr::Connection::new(conn)); tokio::spawn(tran.run()); - let message = conn.send_request(request).get_response().await?; - stats.finalize(); - Ok(Answer { message, stats }) + let streaming = request.is_streaming(); + let send_request = conn.send_request(request); + Self::finalize_request(send_request, stats, streaming).await } fn dgram_config(server: &Server) -> dgram::Config { @@ -230,8 +280,9 @@ impl From for Transport { /// An answer for a query. pub struct Answer { - message: Message, + msgs: Vec>, stats: Stats, + cur_idx: AtomicUsize, } impl Answer { @@ -240,17 +291,22 @@ impl Answer { } pub fn message(&self) -> &Message { - &self.message + &self.msgs[self.cur_idx.load(Ordering::SeqCst)] } pub fn msg_slice(&self) -> Message<&[u8]> { - self.message.for_slice_ref() + self.msgs[self.cur_idx.load(Ordering::SeqCst)].for_slice_ref() + } + + pub fn has_next(&self) -> bool { + let old_cur_idx = self.cur_idx.fetch_add(1, Ordering::SeqCst); + (old_cur_idx + 1) < self.msgs.len() } } impl AsRef> for Answer { fn as_ref(&self) -> &Message { - &self.message + &self.msgs[self.cur_idx.load(Ordering::SeqCst)] } } diff --git a/src/commands/query.rs b/src/commands/query.rs index 9c09d6a..8bb7211 100644 --- a/src/commands/query.rs +++ b/src/commands/query.rs @@ -13,6 +13,8 @@ use domain::net::client::request::{ComposeRequest, RequestMessage}; use domain::rdata::{AllRecordData, Ns, Soa}; use domain::resolv::stub::conf::ResolvConf; use domain::resolv::stub::StubResolver; +use domain::tsig::{Algorithm, Key, KeyName}; +use domain::utils::base64; use std::collections::HashSet; use std::fmt; use std::net::{IpAddr, SocketAddr}; @@ -28,7 +30,7 @@ pub struct Query { qname: NameOrAddr, /// The record type to look up - #[arg(value_name = "QUERY_TYPE")] + #[arg(value_name = "QUERY_TYPE", default_value = "AAAA or PTR")] qtype: Option, /// The server to send the query to. System servers used if missing @@ -114,6 +116,10 @@ pub struct Query { #[arg(long = "no-rd")] no_rd: bool, + /// TSIG signing key to use: :[]: + #[arg(long = "tsig-key")] + tsig_key: Option, + // No need to set the TC flag in the request. /// Disable all sanity checks. #[arg(long, short = 'f')] @@ -181,7 +187,13 @@ impl Query { } }; - let answer = client.request(self.create_request()).await?; + let tsig_key = if let Some(key_str) = &self.tsig_key { + key_from_str(key_str)? + } else { + None + }; + + let answer = client.request(self.create_request(), tsig_key).await?; self.output.format.print(&answer)?; if self.verify { let auth_answer = self.auth_answer().await?; @@ -202,6 +214,41 @@ impl Query { } } +fn key_from_str(key_str: &str) -> Result, Error> { + let key_parts = key_str + .split(':') + .map(ToString::to_string) + .collect::>(); + if key_parts.len() < 2 { + return Err( + "--tsig-key format error: value should be colon ':' separated" + .into(), + ); + } + let key_name = key_parts[0].trim_matches('"'); + let (alg, base64) = match key_parts.len() { + 2 => (Algorithm::Sha256, key_parts[1].clone()), + 3 => { + let alg = Algorithm::from_str(&key_parts[1]) + .map_err(|_| format!("--tsig-key format error: '{}' is not a valid TSIG algorithm", key_parts[1]))?; + (alg, key_parts[2].clone()) + } + _ => return Err( + "--tsig-key format error: should be :[]:" + .into(), + ), + }; + let key_name = KeyName::from_str(key_name).map_err(|err| { + format!("--tsig-key format error: '{key_name}' is not a valid key name: {err}") + })?; + let secret = base64::decode::>(&base64).map_err(|err| { + format!("--tsig-key format error: base64 decoding error: {err}") + })?; + let key = Key::new(alg, &secret, key_name, None, None) + .map_err(|err| format!("--tsig-key format error: {err}"))?; + Ok(Some(key)) +} + /// # Configuration /// impl Query { @@ -339,7 +386,7 @@ impl Query { self.get_ns_addrs(&ns_set, &resolver).await? }; Client::with_servers(servers) - .query((self.qname.to_name(), self.qtype())) + .query((self.qname.to_name(), self.qtype()), None) .await } diff --git a/src/output/dig.rs b/src/output/dig.rs index 5a23443..4755d54 100644 --- a/src/output/dig.rs +++ b/src/output/dig.rs @@ -109,6 +109,14 @@ pub fn write( for item in section { write_record_item(target, &item?)?; } + + while answer.has_next() { + let msg = &mut answer.msg_slice(); + let section = msg.answer().unwrap(); + for item in section { + write_record_item(target, &item?)?; + } + } } // Authority