diff --git a/examples/new_client.rs b/examples/new_client.rs index 3c4cc04a..52a3827c 100644 --- a/examples/new_client.rs +++ b/examples/new_client.rs @@ -1,12 +1,18 @@ use std::net::SocketAddr; use std::time::Duration; -use domain::new_base::name::RevName; -use domain::new_base::{QClass, QType, Question}; -use domain::new_client::exchange::{Allocator, Exchange, ParsedMessage}; +use domain::new_base::build::{BuilderContext, MessageBuilder}; +use domain::new_base::name::{RevName, RevNameBuf}; +use domain::new_base::parse::SplitMessageBytes; +use domain::new_base::wire::U16; +use domain::new_base::{ + Header, HeaderFlags, QClass, QType, Question, Record, +}; +// use domain::new_client::redundant::RedundantClient; use domain::new_client::tcp::{TcpClient, TcpConfig}; use domain::new_client::udp::{UdpClient, UdpConfig}; -use domain::new_client::Client; +use domain::new_client::{Client, ExtendedMessageBuilder}; +use domain::new_rdata::RecordData; use tokio::join; use tokio::net::TcpStream; @@ -15,76 +21,141 @@ async fn main() { env_logger::init(); let metrics = tokio::runtime::Handle::current().metrics(); + let mut args = std::env::args(); + let _ = args.next().unwrap(); + let protocol = args.next().unwrap_or("all".into()).to_lowercase(); + let example = b"\x00\x03org\x07example"; let nlnetlabs = b"\x00\x02nl\x09nlnetlabs"; let google = b"\x00\x03com\x06google"; let addr: SocketAddr = "1.1.1.1:53".parse().unwrap(); - println!("\n=== UDP ==="); - let client = UdpClient::new(addr, UdpConfig::default()); - let res = send_request(example, &client).await; - println!("{}", res); + if protocol == "all" || protocol == "udp" { + println!("\n=== UDP ==="); + let client = UdpClient::new(addr, UdpConfig::default()); + let res = send_request(example, &client).await; + println!("{}", res); + } - println!("\n=== TCP ==="); - let stream = TcpStream::connect(addr).await.unwrap(); - let client = TcpClient::new(stream, TcpConfig::default()); + if protocol == "all" || protocol == "tcp" { + println!("\n=== TCP ==="); + let stream = TcpStream::connect(addr).await.unwrap(); + let client = TcpClient::new(stream); - let res = join!( - send_request(example, &client), - send_request(nlnetlabs, &client), - send_request(google, &client), - ); - println!("{}\n", res.0); - println!("{}\n", res.1); - println!("{}\n", res.2); + let res = join!( + send_request(example, &client), + send_request(nlnetlabs, &client), + send_request(google, &client), + ); + println!("{}\n", res.0); + println!("{}\n", res.1); + println!("{}\n", res.2); - drop(client); + drop(client); - println!("Waiting to see whether tokio will stop the task"); - // Give tokio a bit of time to exit the background task - tokio::time::sleep(Duration::from_secs(1)).await; - let n = metrics.num_alive_tasks(); - println!("Runtime has {} alive tasks", n); + println!("Waiting to see whether tokio will stop the task"); + // Give tokio a bit of time to exit the background task + tokio::time::sleep(Duration::from_secs(1)).await; + let n = metrics.num_alive_tasks(); + println!("Runtime has {} alive tasks", n); + } - println!("\n=== TCP WITH TIMEOUT ==="); - let stream = TcpStream::connect(addr).await.unwrap(); - let client = TcpClient::new( - stream, - TcpConfig { - idle_timeout: Some(Duration::from_millis(500)), - ..Default::default() - }, - ); + if protocol == "all" || protocol == "tcp-timeout" { + println!("\n=== TCP WITH TIMEOUT ==="); + let stream = TcpStream::connect(addr).await.unwrap(); + let client = TcpClient::with_config( + stream, + TcpConfig { + idle_timeout: Some(Duration::from_millis(500)), + ..Default::default() + }, + ); - let res = join!( - send_request(example, &client), - send_request(nlnetlabs, &client), - ); - println!("{}\n", res.0); - println!("{}\n", res.1); - tokio::time::sleep(Duration::from_secs(1)).await; + let res = join!( + send_request(example, &client), + send_request(nlnetlabs, &client), + ); + println!("{}\n", res.0); + println!("{}\n", res.1); + tokio::time::sleep(Duration::from_secs(1)).await; - let res = send_request(google, &client).await; - println!("{res}\n"); + let res = send_request(google, &client).await; + println!("{res}\n"); - drop(client); + drop(client); + } + + // if protocol == "all" || protocol == "redundant" { + // let client = RedundantClient::new(); + // client.add_client(UdpClient::new(addr, UdpConfig::default())); + // client.add_client(UdpClient::new( + // "9.9.9.9:53".parse().unwrap(), + // UdpConfig::default(), + // )); + + // let res = join!( + // send_request(example, &client), + // send_request(nlnetlabs, &client), + // ); + // println!("{}\n", res.0); + // println!("{}\n", res.1); + // } } async fn send_request(name: &[u8], client: &impl Client) -> String { - let mut request = ParsedMessage::default(); - request.flags.request_recursion(true); + let mut buffer = vec![0u8; 65536]; + let mut context = BuilderContext::default(); + + let mut builder = MessageBuilder::new(&mut buffer, &mut context); + *builder.header_mut() = Header { + id: U16::new(0), + flags: *HeaderFlags::default().request_recursion(true), + counts: Default::default(), + }; + let name = unsafe { RevName::from_bytes_unchecked(name) }; - request - .questions - .push(Question::new(name, QType::A, QClass::IN)); + builder + .build_question(&Question::new(name, QType::A, QClass::IN)) + .unwrap() + .unwrap() + .commit(); - let mut bump = bumpalo::Bump::new(); - let mut exchange = Exchange::new(&mut bump); - exchange.request = request; + let request = ExtendedMessageBuilder { + builder, + edns_record: None, + }; - match client.request(&mut exchange).await { - Ok(()) => format!("{:?}", exchange.response), + match client.request(request).await { + Ok(msg) => { + let mut s = msg.header.to_string(); + + let mut offset = 0; + + for _ in 0..msg.header.counts.questions.get() { + let (_question, rest) = + Question::::split_message_bytes( + &msg.contents, + offset, + ) + .unwrap(); + offset = rest; + } + + for _ in 0..msg.header.counts.answers.get() { + let (answer, rest) = Record::< + RevNameBuf, + RecordData<'_, RevNameBuf>, + >::split_message_bytes( + &msg.contents, offset + ) + .unwrap(); + s.push('\n'); + s.push_str(&format!("{:?}", answer)); + offset = rest; + } + s + } Err(err) => format!("Error: {:?}", err), } } diff --git a/src/new_client/exchange.rs b/src/new_client/exchange.rs deleted file mode 100644 index 3321ba1d..00000000 --- a/src/new_client/exchange.rs +++ /dev/null @@ -1,599 +0,0 @@ -//! Request-response exchanges for DNS servers. -//! -//! This module provides a number of utility types for the DNS service layer -//! architecture. In particular, an [`Exchange`] represents a DNS request as -//! it is being passed along a server pipeline, and an [`OutgoingResponse`] is -//! the corresponding response as it is passed back through. - -use core::{ - alloc::Layout, - any::{Any, TypeId}, -}; -use std::{boxed::Box, vec::Vec}; - -use bumpalo::Bump; - -use crate::{ - new_base::{ - build::{BuilderContext, MessageBuilder}, - name::{RevName, RevNameBuf}, - parse::SplitMessageBytes, - wire::{BuildBytes, ParseError, SizePrefixed, TruncationError, U16}, - HeaderFlags, Message, Question, RType, Record, SectionCounts, - }, - new_edns::{EdnsFlags, EdnsOption, EdnsRecord}, - new_rdata::{Opt, RecordData}, - utils::UnsizedClone, -}; - -//----------- Exchange ------------------------------------------------------- - -/// A DNS request-response exchange. -/// -/// An [`Exchange`] represents a request sent to a DNS server and the server's -/// response (as it is being built). It tracks basic information about the -/// request, such as when it was sent and the connection it originates from, -/// as well as metadata stored by layers in the DNS server. -pub struct Exchange<'a> { - /// An allocator for storing parts of the message. - pub alloc: Allocator<'a>, - - /// The request message. - pub request: ParsedMessage<'a>, - - /// The response message being built. - pub response: ParsedMessage<'a>, - - /// Dynamic metadata stored by the DNS server. - pub metadata: Vec, -} - -impl<'a> Exchange<'a> { - pub fn new(bump: &'a mut Bump) -> Self { - Self { - alloc: Allocator::new(bump), - request: ParsedMessage::default(), - response: ParsedMessage::default(), - metadata: Vec::new(), - } - } -} - -//----------- OutgoingResponse ----------------------------------------------- - -/// An [`Exchange`] with an initialized response message. -pub struct OutgoingResponse<'e, 'a> { - /// An allocator for storing parts of the message. - pub alloc: &'e mut Allocator<'a>, - - /// The response message being built. - pub response: &'e mut ParsedMessage<'a>, - - /// Dynamic metadata stored by the DNS server. - pub metadata: &'e mut Vec, -} - -impl<'e, 'a> OutgoingResponse<'e, 'a> { - /// Construct an [`OutgoingResponse`] on an [`Exchange`]. - pub fn new(exchange: &'e mut Exchange<'a>) -> Self { - Self { - alloc: &mut exchange.alloc, - response: &mut exchange.response, - metadata: &mut exchange.metadata, - } - } - - /// Reborrow this response for a shorter lifetime. - pub fn reborrow(&mut self) -> OutgoingResponse<'_, 'a> { - OutgoingResponse { - alloc: self.alloc, - response: self.response, - metadata: self.metadata, - } - } -} - -//----------- ParsedMessage -------------------------------------------------- - -/// A pre-parsed DNS message. -/// -/// This is a simple representation of DNS messages outside the wire format, -/// making it easy to inspect and modify them efficiently. -#[derive(Clone, Default, Debug)] -pub struct ParsedMessage<'a> { - /// The message ID. - pub id: U16, - - /// The message flags. - pub flags: HeaderFlags, - - /// Questions in the message. - pub questions: Vec>, - - /// Answer records in the message. - pub answers: Vec>>, - - /// Authority records in the message. - pub authorities: Vec>>, - - /// Additional records in the message. - /// - /// If there is an EDNS record, it will be included here, but its record - /// data (which contains the EDNS options) will be empty. The options are - /// stored in the `options` field for easier access. - pub additional: Vec>>, - - /// EDNS options in the message. - /// - /// These options will be appended to the EDNS record in the additional - /// section (there must be one for any options to exist). The order of - /// the options is meaningless. - pub options: Vec>, -} - -impl<'a> ParsedMessage<'a> { - /// Parse an existing [`Message`]. - /// - /// Decompressed domain names are allocated using the given [`Bump`]. - pub fn parse( - message: &Message, - alloc: &mut Allocator<'a>, - ) -> Result { - type ParsedQuestion = Question; - type ParsedRecord<'a> = - Record>; - - let mut this = ParsedMessage::<'a>::default(); - let mut offset = 0; - - // Parse the message header. - this.id = message.header.id; - this.flags = message.header.flags; - let counts = message.header.counts; - - // Parse the question section. - this.questions - .reserve(counts.questions.get().max(256) as usize); - for _ in 0..counts.questions.get() { - let (question, rest) = ParsedQuestion::split_message_bytes( - &message.contents, - offset, - )?; - - this.questions - .push(question.map_name(|n| &*alloc.alloc_unsized(&*n))); - offset = rest; - } - - // Parse the answer section. - this.answers.reserve(counts.answers.get().max(256) as usize); - for _ in 0..counts.answers.get() { - let (answer, rest) = - ParsedRecord::split_message_bytes(&message.contents, offset)?; - - this.answers.push(Record { - rname: alloc.alloc_unsized(&*answer.rname), - rtype: answer.rtype, - rclass: answer.rclass, - ttl: answer.ttl, - rdata: answer - .rdata - .map_names(|n| &*alloc.alloc_unsized(&*n)) - .clone_to_bump(alloc.inner), - }); - offset = rest; - } - - // Parse the authority section. - this.authorities - .reserve(counts.authorities.get().max(256) as usize); - for _ in 0..counts.authorities.get() { - let (authority, rest) = - ParsedRecord::split_message_bytes(&message.contents, offset)?; - - this.authorities.push(Record { - rname: alloc.alloc_unsized(&*authority.rname), - rtype: authority.rtype, - rclass: authority.rclass, - ttl: authority.ttl, - rdata: authority - .rdata - .map_names(|n| &*alloc.alloc_unsized(&*n)) - .clone_to_bump(alloc.inner), - }); - offset = rest; - } - - // The EDNS record data. - let mut edns_data = None; - - // Parse the additional section. - this.additional - .reserve(counts.additional.get().max(256) as usize); - for _ in 0..counts.additional.get() { - let (additional, rest) = - ParsedRecord::split_message_bytes(&message.contents, offset)?; - - if let RecordData::Opt(opt) = additional.rdata { - if edns_data.is_some() { - // A message cannot contain two distinct EDNS records. - return Err(ParseError); - } - - edns_data = Some(opt); - - // XXX: Deduplicate the EDNS data. - // additional.rdata = RecordData::Opt(Opt::EMPTY); - } - - this.additional.push(Record { - rname: alloc.alloc_unsized(&*additional.rname), - rtype: additional.rtype, - rclass: additional.rclass, - ttl: additional.ttl, - rdata: additional - .rdata - .map_names(|n| &*alloc.alloc_unsized(&*n)) - .clone_to_bump(alloc.inner), - }); - offset = rest; - } - - // Ensure there's no other content in the message. - if offset != message.contents.len() { - return Err(ParseError); - } - - // Parse EDNS options. - if let Some(edns_data) = edns_data { - for option in edns_data.options() { - this.options.push(option?.clone_to_bump(alloc.inner)); - } - } - - Ok(this) - } - - /// Build this message into the given buffer. - /// - /// If the message could not fit in the given buffer, a - /// [`TruncationError`] is returned. - pub fn build<'b, 'c>( - &self, - context: &'c mut BuilderContext, - buffer: &'b mut [u8], - ) -> Result, TruncationError> { - // Construct a 'MessageBuilder'. - if buffer.len() < 12 { - return Err(TruncationError); - } - let mut builder = MessageBuilder::new(buffer, context); - - // Build the message header. - let header = builder.header_mut(); - header.id = self.id; - header.flags = self.flags; - header.counts = SectionCounts::default(); - - // Build the question section. - for question in &self.questions { - builder - .build_question(question)? - .expect("No answers, authorities, or additionals are built") - .commit(); - } - - // Build the answer section. - for answer in &self.answers { - builder - .build_answer(answer)? - .expect("No authorities, or additionals are built") - .commit(); - } - - // Build the authority section. - for authority in &self.authorities { - builder - .build_authority(authority)? - .expect("No additionals are built") - .commit(); - } - - // Build the additional section. - let mut edns_built = false; - for additional in &self.additional { - if additional.rtype == RType::OPT { - // Technically, multiple OPT records are an error. But this - // isn't the right place to report that. - debug_assert!(!edns_built, "Multiple EDNS records found"); - - let mut builder = builder.build_additional(additional)?; - let mut delegate = builder.delegate(); - let mut uninit = delegate.uninitialized(); - for option in &self.options { - uninit = option.build_bytes(uninit)?; - } - let uninit_len = uninit.len(); - let appended = delegate.uninitialized().len() - uninit_len; - delegate.mark_appended(appended); - delegate.commit(); - builder.commit(); - - edns_built = true; - continue; - } - - builder.build_additional(additional)?.commit(); - } - - debug_assert!( - self.options.is_empty() || edns_built, - "EDNS options found, but no OPT record", - ); - - Ok(builder) - } -} - -impl ParsedMessage<'_> { - /// Whether this message has an EDNS record. - pub fn has_edns(&self) -> bool { - self.additional.iter().any(|r| r.rtype == RType::OPT) - } - - pub fn set_max_udp_payload_size(&mut self, max_payload_size: u16) { - self.additional - .iter_mut() - .find_map(|r| match r.rtype { - RType::OPT => { - r.rclass.code.set(max_payload_size); - Some(()) - } - _ => None, - }) - .unwrap_or_else(|| { - self.additional.push( - EdnsRecord { - max_udp_payload: max_payload_size.into(), - ext_rcode: 0, - version: 0, - flags: EdnsFlags::default(), - options: SizePrefixed::new(Opt::EMPTY), - } - .into(), - ) - }); - } -} - -impl ParsedMessage<'_> { - /// Reset this object to a blank message. - /// - /// This is helpful in order to reuse the underlying allocations. - pub fn reset(&mut self) { - self.id = U16::new(0); - self.flags = HeaderFlags::default(); - self.questions.clear(); - self.answers.clear(); - self.authorities.clear(); - self.additional.clear(); - self.options.clear(); - } -} - -//----------- ResponseCode --------------------------------------------------- - -/// A (possibly extended) DNS response code. -#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] -pub enum ResponseCode { - /// The request was answered successfully. - Success, - - /// The request was misformatted. - FormatError, - - /// The server encountered an internal error. - ServerFailure, - - /// The queried domain name does not exist. - NonExistentDomain, - - /// The server does not support the requested kind of query. - NotImplemented, - - /// Policy prevents the server from answering the query. - Refused, - - /// The TSIG record in the request was invalid. - InvalidTSIG, - - /// The server does not support the request's OPT record version. - UnsupportedOptVersion, - - /// The request did not contain a valid EDNS server cookie. - BadCookie, -} - -impl ResponseCode { - /// This code's representation in the DNS message header. - pub const fn header_bits(&self) -> u8 { - match self { - Self::Success => 0, - Self::FormatError => 1, - Self::ServerFailure => 2, - Self::NonExistentDomain => 3, - Self::NotImplemented => 4, - Self::Refused => 5, - Self::InvalidTSIG => 9, - Self::UnsupportedOptVersion => 0, - Self::BadCookie => 7, - } - } - - /// This code's representation in the EDNS record header. - pub const fn edns_bits(&self) -> u8 { - match self { - Self::Success => 0, - Self::FormatError => 0, - Self::ServerFailure => 0, - Self::NonExistentDomain => 0, - Self::NotImplemented => 0, - Self::Refused => 0, - Self::InvalidTSIG => 0, - Self::UnsupportedOptVersion => 1, - Self::BadCookie => 1, - } - } -} - -//----------- Metadata ------------------------------------------------------- - -/// Arbitrary metadata about a DNS exchange. -/// -/// This should be used by [`ServiceLayer`](super::ServiceLayer)s for storing -/// information they have extracted from an incoming DNS request message. The -/// metadata may be relevant to future layers: for example, some may wish to -/// handle TSIG-signed requests differently from others. The metadata is also -/// relevant to the original layer in [`process_outgoing()`], as it does not -/// have access to the original request. -/// -/// [`process_outgoing()`]: super::ServiceLayer::process_outgoing() -/// -/// # Implementation -/// -/// This is an enhanced version of `Box` that can -/// perform downcasting more efficiently. It stores the [`TypeId`] of the -/// object inline, allowing it to skip a vtable lookup. -pub struct Metadata { - /// The type ID of the object. - type_id: TypeId, - - /// The underlying object. - object: Box, -} - -impl Metadata { - /// Wrap an object in [`Metadata`]. - pub fn new(object: T) -> Self { - let type_id = TypeId::of::(); - let object = Box::new(object) as Box; - Self { type_id, object } - } - - /// Check whether this is metadata of a certain type. - pub fn is(&self) -> bool { - self.type_id == TypeId::of::() - } - - /// Try downcasting to a reference of a particular type. - pub fn try_as(&self) -> Option<&T> { - if !self.is::() { - return None; - } - - let pointer: *const (dyn Any + Send + 'static) = &*self.object; - // SAFETY: 'pointer' was created by 'Box::into_raw()', and thus is - // safe to dereference (the pointer will only be dropped when 'self' - // is, but that cannot happen during the current lifetime). - Some(unsafe { &*pointer.cast::() }) - } - - /// Try downcasting to a mutable reference of a particular type. - pub fn try_as_mut(&mut self) -> Option<&mut T> { - if !self.is::() { - return None; - } - - let pointer: *mut (dyn Any + Send + 'static) = &mut *self.object; - // SAFETY: 'pointer' was created by 'Box::into_raw()', and thus is - // safe to dereference (the pointer will only be dropped when 'self' - // is, but that cannot happen during the current lifetime). - Some(unsafe { &mut *pointer.cast::() }) - } - - /// Try moving this object out of the [`Metadata`]. - pub fn try_into(self) -> Result { - if !self.is::() { - return Err(self); - } - - let pointer: *mut _ = Box::into_raw(self.object); - // SAFETY: 'pointer' was created by 'Box::into_raw()', and thus is - // safe to move into the same 'Box'. - Ok(*unsafe { Box::from_raw(pointer.cast::()) }) - } -} - -//----------- Allocator ------------------------------------------------------ - -/// A bump allocator with a fixed lifetime. -/// -/// This is a wrapper around [`bumpalo::Bump`] that guarantees thread safety. -/// It is equivalent to `&'a mut Bump`, but `&mut &'a mut Bump` does not work -/// (allocated objects only last for the shorter lifetime, not for `'a`). -/// `&mut Allocator<'a>` does work, giving objects of lifetime `'a`. -/// -/// # Thread Safety -/// -/// [`Bump`] is not thread safe; using it from multiple threads simultaneously -/// would cause undefined behaviour. [`Allocator`] implements [`Send`], and -/// so it cannot directly expose shared references to the underlying [`Bump`]; -/// a user could get `&Bump` on one thread, send the [`Allocator`] to another -/// thread, then get `&Bump` over there. This is why [`Allocator`] copies -/// [`Bump`]'s methods instead of implementing [`Deref`] to [`Bump`]. -/// -/// [`Deref`]: core::ops::Deref -#[derive(Debug)] -#[repr(transparent)] -pub struct Allocator<'a> { - /// The underlying allocator. - /// - /// In order to share access to a [`Bump`], even on a single thread, it - /// must be a shared reference (`&'a Bump`). That is how we store it - /// here. However, we guarantee that the [`Allocator`] is constructed - /// from a mutable reference -- thus that this is the only reference to - /// the bump allocator. It is never exposed publicly, so it cannot be - /// copied and used from multiple threads. - inner: &'a Bump, -} - -impl<'a> Allocator<'a> { - /// Construct a new [`Allocator`]. - pub fn new(inner: &'a mut Bump) -> Self { - // NOTE: The 'Bump' is mutably borrowed for lifetime 'a; the reference - // we store is thus guaranteed to be unique. - Self { inner } - } - - /// Allocate an object. - pub fn alloc(&mut self, val: T) -> &'a mut T { - self.inner.alloc(val) - } - - /// Allocate a slice and copy the given contents into it. - pub fn alloc_slice_copy(&mut self, src: &[T]) -> &'a mut [T] { - self.inner.alloc_slice_copy(src) - } - - pub fn alloc_unsized( - &mut self, - val: &T, - ) -> &'a mut T { - let layout = Layout::for_value(val); - let ptr = self.inner.alloc_layout(layout).as_ptr().cast::<()>(); - unsafe { - val.unsized_clone(ptr); - }; - let ptr = val.ptr_with_address(ptr); - unsafe { &mut *ptr } - } -} - -// SAFETY: An 'Allocator' contains '&Bump', which is '!Send' because 'Bump' is -// '!Sync'. However, we guarantee that there are no other references to the -// 'Bump' -- that this is really '&mut Bump' (which is 'Send'). -unsafe impl Send for Allocator<'_> {} - -// NOTE: 'Allocator' acts a bit like the nightly-only 'std::sync::Exclusive', -// since it doesn't provide any shared access to the underlying 'Bump'. It is -// sound for it to implement 'Sync', but we defer this until necessary. diff --git a/src/new_client/load_balancer.rs b/src/new_client/load_balancer.rs new file mode 100644 index 00000000..0e2a4b86 --- /dev/null +++ b/src/new_client/load_balancer.rs @@ -0,0 +1,268 @@ +use core::time::Duration; +use std::{boxed::Box, sync::Arc, vec::Vec}; + +use futures_util::{stream::FuturesUnordered, StreamExt}; +use parking_lot::RwLock; +use tokio::time::{timeout_at, Instant}; + +use crate::new_base::Message; + +use super::{BoxClient, Client, ClientError}; + +#[derive(Clone, Debug)] +pub struct LoadBalancerConfig { + /// Defer transport errors. + pub defer_transport_error: bool, + + /// Defer replies that report Refused. + pub defer_refused: bool, + + /// Defer replies that report ServFail. + pub defer_servfail: bool, + + /// Cut-off for slow upstreams as a factor of the fastest upstream. + pub slow_rt_factor: f64, +} + +impl Default for LoadBalancerConfig { + fn default() -> Self { + Self { + defer_transport_error: false, + defer_refused: false, + defer_servfail: false, + slow_rt_factor: 5.0, + } + } +} + +/// Configuration variables for each upstream. +#[derive(Clone, Copy, Debug, Default)] +pub struct SubClientConfig { + /// Maximum burst of upstream queries. + max_burst: Option, + + /// Interval over which the burst is counted. + burst_interval: Duration, +} + +pub struct SubClient { + config: SubClientConfig, + client: Box, + stats: Mutex, +} + +pub struct LoadBalancerClient { + config: LoadBalancerConfig, + clients: RwLock>>, +} + +impl From for SubClient { + fn from(value: C) -> Self { + SubClient { + config: Default::default(), + client: Box::new(value), + } + } +} + +impl LoadBalancerClient { + pub fn new() -> Self { + Self::with_config(Default::default()) + } + + pub fn with_config(config: LoadBalancerConfig) -> Self { + Self { + config, + clients: Default::default(), + } + } + + pub fn add_client(&mut self, client: impl BoxClient) { + self.add_client_with_config(client, Default::default()) + } + + pub fn add_client_with_config( + &mut self, + client: impl BoxClient, + config: SubClientConfig, + ) { + self.clients.push(client.into()); + } + + fn sort_clients(&self) -> Vec { + let mut clients: Vec = self + .clients + .read() + .iter() + .cloned() + .map(Into::into) + .collect(); + + todo!() + } + + /// Determine whether a successful response should be skipped. + /// + /// We skip a response if the `RCODE` is `SERVFAIL` or `REFUSED` + fn skip(&self, msg: &Vec) -> bool { + let Ok(msg) = Message::parse_bytes_by_ref(msg) else { + return false; + }; + + // We match on SERVFAIL and REFUSED. If the normal rcode matches that + // we have to ensure that the extended rcode is 0. + match msg.header.flags.rcode() { + 2 /* SERVFAIL */ if self.config.defer_servfail => { + matches!(find_opt_rcode(msg), Some(0) | None) + } + 5 /* REFUSED */ if self.config.defer_refused => { + matches!(find_opt_rcode(msg), Some(0) | None) + } + _ => false + } + } +} + +impl Client for LoadBalancerClient { + async fn request( + &self, + request: &Message, + ) -> Result, ClientError> { + // This will be our view of the clients for this request. + // We sort them based on their timeout and iterate over them in that + // order. + let clients = self.sort_clients(); + if clients.is_empty() { + return Err(ClientError::NoTransportAvailable); + } + let mut clients = clients.into_iter(); + + let mut futs = FuturesUnordered::new(); + + // The time at which the next request should be sent out. + let mut next_request_time = Instant::now(); + + // This will hold the result of requests that fail or that we skip, + // so we can return them later when the subsequent requests also + // fail. + let mut deferred_result = None; + + loop { + match timeout_at(next_request_time, futs.next()).await { + Ok(Some(res)) => { + // got some response, so we decide whether to return it, + // store it into the deferred_result or discard it. + match res { + Ok(msg) if self.skip(&msg) => { + if let Some(Err(_)) | None = deferred_result { + deferred_result = Some(Ok(msg)); + } + } + Err(err) if self.config.defer_transport_error => { + if deferred_result.is_none() { + deferred_result = Some(Err(err)); + } + } + // It's not one of the cases we defer or skip, so we + // return it! + result => { + return result; + } + } + } + // On a timeout or empty set of futures we start a new + // request. + // + // An empty set of futures happens in two cases: + // 1. We haven't send out any requests yet + // 2. All sent requests have been resolved, which means + // we can send more requests immediately. + Ok(None) | Err(_) => { + if let Some(RequestClient { client, timeout }) = + clients.next() + { + futs.push(client.request(message)); + next_request_time = Instant::now() + timeout; + } else { + return deferred_result + .unwrap_or(Err(ClientError::Bug)); + } + } + } + } + } +} + +struct RequestClient { + client: Arc, + timeout: Duration, +} + +impl RequestClient { + fn new(client: &Arc) -> Self { + let timeout = client.stats.lock().timeout; + RequestClient { + timeout, + client: client.clone(), + } + } +} + +/// Find the extended RCODE in the message. +/// +/// Note that the returned `u8` only contains the upper 8 bits of the +/// `RCODE`, i.e. the part stored in the `OPT` record. +/// +/// `None` is returned on parse errors. +/// +/// We have to write this here because new_base is lacking some proper +/// handling of the (extended) `RCODE`. +fn find_opt_rcode(msg: &Message) -> Option { + let counts = msg.header.counts; + + let mut offset = 0; + for _ in 0..counts.questions.get() { + let (_, rest) = Question::<&UnparsedName>::split_message_bytes( + &msg.contents, + offset, + ) + .ok()?; + offset = rest; + } + + for _ in 0..counts.answers.get() { + let (_, rest) = Record::<&UnparsedName, &UnparsedRecordData>::split_message_bytes( + &msg.contents, + offset, + ).ok()?; + offset = rest; + } + + for _ in 0..counts.authorities.get() { + let (_, rest) = Record::<&UnparsedName, &UnparsedRecordData>::split_message_bytes( + &msg.contents, + offset, + ).ok()?; + offset = rest; + } + + for _ in 0..counts.additional.get() { + let (r, rest) = + Record::<&UnparsedName, &UnparsedRecordData>::split_message_bytes( + &msg.contents, + offset, + ) + .ok()?; + + if let RType::OPT = r.rtype { + // The extension of the rcode is specified as the first 8 bits + // of the TTL field in RFC 6891. + let ttl_bytes: [u8; 4] = r.ttl.value.get().to_be_bytes(); + return Some(ttl_bytes[0]); + } + + offset = rest; + } + + None +} diff --git a/src/new_client/mod.rs b/src/new_client/mod.rs index b22bc5f9..ce1e2c58 100644 --- a/src/new_client/mod.rs +++ b/src/new_client/mod.rs @@ -1,17 +1,102 @@ +use core::{future::Future, pin::Pin}; +use std::boxed::Box; use std::io; +use std::vec::Vec; -use exchange::Exchange; +use crate::new_base::build::MessageBuilder; +use crate::new_base::name::RevName; +use crate::new_base::wire::{SizePrefixed, TruncationError, U16}; +use crate::new_base::{Message, Record}; +use crate::new_edns::{EdnsFlags, EdnsOption, EdnsRecord}; +use crate::new_rdata::{Opt, RecordData}; -pub mod exchange; +// pub mod redundant; +pub mod multi_tcp; pub mod tcp; pub mod udp; pub trait Client { #[allow(async_fn_in_trait)] - async fn request<'a>( + async fn request( &self, - exchange: &mut Exchange<'a>, - ) -> Result<(), ClientError>; + request: ExtendedMessageBuilder<'_, '_>, + ) -> Result, ClientError>; +} + +pub trait BoxClient { + fn dyn_request<'a, 'b: 'a, 'c: 'a>( + &'a self, + request: ExtendedMessageBuilder<'b, 'c>, + ) -> Pin, ClientError>> + 'a>>; +} + +impl BoxClient for T { + fn dyn_request<'a, 'b: 'a, 'c: 'a>( + &'a self, + request: ExtendedMessageBuilder<'b, 'c>, + ) -> Pin, ClientError>> + 'a>> + { + Box::pin(self.request(request)) + } +} + +#[derive(Clone, Debug, Default)] +pub struct EdnsRecordBuilder<'a> { + pub header: EdnsHeader, + pub options: Vec>, +} + +#[derive(Clone, Debug, Default)] +pub struct EdnsHeader { + pub max_udp_payload_size: U16, + pub ext_rcode: u8, + pub version: u8, + pub flags: EdnsFlags, +} + +/// A message with the OPT data kept separately for easy access and modification. +pub struct ExtendedMessageBuilder<'b, 'c> { + pub builder: MessageBuilder<'b, 'c>, + pub edns_record: Option>, +} + +impl<'b, 'c> ExtendedMessageBuilder<'b, 'c> { + pub fn build(self) -> Result<&'b mut Message, TruncationError> { + let Self { + mut builder, + edns_record, + } = self; + if let Some(edns_record) = edns_record { + let h = edns_record.header; + let record = EdnsRecord { + max_udp_payload: h.max_udp_payload_size, + ext_rcode: h.ext_rcode, + version: h.version, + flags: h.flags, + options: SizePrefixed::new(&Opt::EMPTY), + }; + let record: Record<&RevName, RecordData<'_, &RevName>> = + record.into(); + let mut builder = builder.build_additional(&record)?; + let mut delegate = builder.delegate(); + delegate.append_built_bytes(&&*edns_record.options)?; + delegate.commit(); + builder.commit(); + } + Ok(builder.finish()) + } + + pub fn set_id(&mut self, id: u16) { + self.builder.header_mut().id.set(id); + } + + pub fn set_udp_max_payload_size(&mut self, size: u16) { + self.get_edns_mut().header.max_udp_payload_size.set(size); + } + + pub fn get_edns_mut(&mut self) -> &mut EdnsRecordBuilder<'b> { + self.edns_record.get_or_insert_default() + } } #[derive(Clone, Debug)] @@ -42,6 +127,8 @@ pub enum ClientError { Closed, TimedOut, + + NoTransportAvailable, } impl From for ClientError { diff --git a/src/new_client/multi_tcp.rs b/src/new_client/multi_tcp.rs new file mode 100644 index 00000000..b4a4774c --- /dev/null +++ b/src/new_client/multi_tcp.rs @@ -0,0 +1,427 @@ +//! Multiplexed TCP Client +//! +//! A [`TcpClient`] maintains a TCP connection. Each [`TcpClient::request`] +//! call sends a request over that connection and returns a future for the +//! response. +//! +//! If you require a long-lived connection with a server. You probably want +//! to use a multi TCP stream (to be implemented). +//! +//! Characteristics of this implementation: +//! +//! - Messages sent with this client should not contain TSIG records, +//! because the message id will be modified, invalidating the signature. +//! - Each [`TcpClient`] spawns a background task for reading the incoming +//! messages. +//! - The background task will abort when [`TcpClient`] is dropped. +//! - The ids assigned to each message will usually be low and may be reused. +//! - If the connection is found to be in a broken state. All requests will +//! receive errors. A new [`TcpClient`] should be created at this point. +//! - `edns-tcp-keepalive` is ignored, because we simply keep the connection +//! around for as long as we need it. +//! +//! # Relevant RFC excerpts +//! +//! RFC 1035, Section 4.2.2: +//! +//! > Messages sent over TCP connections use server port 53 (decimal). The +//! > message is prefixed with a two byte length field which gives the +//! > message length, excluding the two byte length field. This length field +//! > allows the low-level processing to assemble a complete message before +//! > beginning to parse it. +//! +//! RCF 7766, Section 6.2.1: +//! +//! > To amortise connection setup costs, both clients and servers SHOULD +//! > support connection reuse by sending multiple queries and responses over +//! > a single persistent TCP connection. +//! > +//! > When sending multiple queries over a TCP connection, clients MUST NOT +//! > reuse the DNS Message ID of an in-flight query on that connection in +//! > order to avoid Message ID collisions. +//! +//! RFC 7766, Section 6.2.1.1: +//! +//! > In order to achieve performance on par with UDP, DNS clients SHOULD +//! > pipeline their queries. When a DNS client sends multiple queries to +//! > a server, it SHOULD NOT wait for an outstanding reply before sending +//! > the next query. +//! +//! > It is likely that DNS servers need to process pipelined queries +//! > concurrently and also send out-of-order responses over TCP in order +//! > to provide the level of performance possible with UDP transport. +//! +//! RFC 7766, Secton 6.2.3: +//! +//! > DNS clients SHOULD close the TCP connection of an idle session, unless +//! > an idle timeout has been established using some other signalling +//! > mechanism, for example, edns-tcp-keepalive. +//! +//! RFC 7858, Section 3.4: +//! +//! > In order to amortize TCP and TLS connection setup costs, clients and +//! > servers SHOULD NOT immediately close a connection after each response. +//! > Instead, clients and servers SHOULD reuse existing connections for +//! > subsequent queries as long as they have sufficient resources. +//! +//! RFC 7766, Section 8: +//! +//! > DNS clients and servers SHOULD pass the two-octet length field, and +//! > the message described by that length field, to the TCP layer at the +//! > same time (e.g., in a single "write" system call) to make it more +//! > likely that all the data will be transmitted in a single TCP segment. +use core::convert::Infallible; +use core::mem; +use core::net::SocketAddr; +use std::boxed::Box; +use std::io; +use std::sync::Arc; +use std::time::Duration; +use std::vec::Vec; + +use futures_util::{stream, Stream, StreamExt}; +use slab::Slab; +use tokio::io::{split, AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf}; +use tokio::net::TcpStream; +use tokio::sync::{mpsc, oneshot}; +use tokio::task::JoinHandle; +use tokio::time::{timeout, timeout_at, Instant}; + +use crate::new_base::wire::{BuildBytes, ParseBytesByRef, SizePrefixed, U16}; +use crate::new_base::Message; +use crate::utils::CloneFrom; + +use super::{Client, ClientError, SocketError}; + +#[derive(Clone, Debug)] +pub struct TcpConfig { + /// Response timeout currently in effect. + pub response_timeout: Duration, + + /// Time until the connection will close if there are no requests waiting + /// for a response. + /// + /// Setting this to a low value might leads to the connection being closed + /// before the first request is sent. + /// + /// Setting this to `None` will close the connection when the + /// client and all requests are dropped. + pub idle_timeout: Option, +} + +impl Default for TcpConfig { + fn default() -> Self { + Self { + response_timeout: Duration::from_secs(19), + idle_timeout: None, + } + } +} + +struct AbortJoinHandle(JoinHandle); + +impl Drop for AbortJoinHandle { + fn drop(&mut self) { + self.0.abort(); + } +} + +struct Request { + callback_send: oneshot::Sender, ClientError>>, + message: Box, +} + +pub struct TcpClient { + /// Ensure that the read loop lives as long as the client + _background: AbortJoinHandle<()>, + config: TcpConfig, + send: mpsc::Sender, +} + +impl TcpClient { + pub fn new(addr: SocketAddr) -> Self { + Self::with_config(addr, Default::default()) + } + + pub fn with_config(addr: SocketAddr, config: TcpConfig) -> Self { + // The buffer size is chosen arbitrarily. + let (send, recv) = mpsc::channel(100); + + let background = Background::new(addr, config.clone(), recv); + let background = AbortJoinHandle(tokio::spawn(background.run())); + + Self { + _background: background, + config, + send, + } + } +} + +impl Client for TcpClient { + async fn request( + &self, + request: super::ExtendedMessageBuilder<'_, '_>, + ) -> Result, ClientError> { + let message = request.build().unwrap(); + let message = CloneFrom::clone_from(message); + + let (callback_send, callback_recv) = oneshot::channel(); + self.send + .send(Request { + callback_send, + message, + }) + .await + .unwrap(); + + callback_recv.await.unwrap() + } +} + +enum Waiting { + AwaitingReponse { + timeout_at: Instant, + callback_send: oneshot::Sender, ClientError>>, + }, + TimedOut, +} + +struct Background { + addr: SocketAddr, + config: TcpConfig, + requests: mpsc::Receiver, +} + +impl Background { + fn new( + addr: SocketAddr, + config: TcpConfig, + requests: mpsc::Receiver, + ) -> Self { + Self { + addr, + config, + requests, + current_requests: Slab::new(), + } + } + + async fn run(mut self) { + // This loop just waits for the next request to come in because we + // don't have a connection open. + loop { + let Some(req) = self.next_request().await else { + return; + }; + let Ok(mut connection) = self.connect().await else { + req.callback_send.send(Err(ClientError::Broken)); + continue; + }; + if let Err(_) = connection.send_request(req).await { + // Consider this connection broken and make a new one + continue; + } + + connection.run().await + } + } + + async fn next_request(&mut self) -> Option { + self.requests.recv().await + } + + async fn connect<'a>( + &'a mut self, + ) -> io::Result< + Connection<'a, impl Stream, ClientError>>>, + > { + let stream = TcpStream::connect(self.addr).await?; + let (read_half, write_half) = split(stream); + let read_stream = read_stream(read_half); + Ok(Connection { + stream: write_half, + background: self, + read_stream, + }) + } +} + +struct Connection< + 'a, + S: Stream, ClientError>> + Unpin, +> { + stream: WriteHalf, + background: &'a mut Background, + read_stream: S, + current_requests: Slab, +} + +impl, ClientError>> + Unpin> + Connection<'_, S> +{ + async fn send_request( + &mut self, + mut req: Request, + ) -> Result<(), ClientError> { + let timeout_at = + Instant::now() + self.background.config.response_timeout; + let id = self.current_requests.insert(Waiting::AwaitingReponse { + timeout_at, + callback_send: req.callback_send, + }); + + let Ok(id) = u16::try_from(id) else { + return Err(ClientError::TooManyRequests); + }; + req.message.header.id.set(id); + + // We allocate the space for the maximum DNS message size and 2 + // additional bytes for the shim. + let mut buffer = vec![0u8; 65535 + 2]; + + // XXX: remove unwraps + let request = SizePrefixed::::new(req.message); + let _ = request.build_bytes(&mut buffer).unwrap(); + + let size = u16::from_be_bytes(*buffer.first_chunk::<2>().unwrap()); + buffer.truncate(2 + size as usize); + + let res = self + .stream + .write_all(&buffer) + .await + .map_err(|e| SocketError::Send(e.kind())); + + Ok(res?) + } + + fn earliest_timeout(&self) -> Option<(usize, Instant)> { + self.current_requests + .iter() + .filter_map(|(idx, w)| match w { + Waiting::AwaitingReponse { timeout_at, .. } => { + Some((idx, *timeout_at)) + } + Waiting::TimedOut => None, + }) + .min_by_key(|(_, instant)| *instant) + } + + async fn run(&mut self) { + let error = loop { + if let Some((idx, earliest_timeout)) = self.earliest_timeout() { + let fut = + timeout_at(earliest_timeout, self.read_stream.next()); + let Ok(res) = fut.await else { + let r = &mut self.current_requests[idx]; + if let Waiting::AwaitingReponse { + timeout_at, + callback_send, + } = r + { + callback_send.send(Err(ClientError::TimedOut)); + *r = Waiting::TimedOut; + } + continue; + }; + match res { + // We got a response, send it to the waiting client. + Some(Ok(res)) => { + let id = res.header.id; + + // If we don't know the id we got, the other side is sending garbage, + // so close the connection. + let Some(r) = self + .current_requests + .try_remove(id.get() as usize) + else { + break ClientError::Broken; + }; + if let Waiting::AwaitingReponse { + callback_send, + .. + } = r + { + callback_send.send(Ok(res)); + } + } + Some(Err(e)) => { + break e; + } + None => { + break ClientError::Closed; + } + } + } else if let Some(idle_timeout) = + self.background.config.idle_timeout + { + match timeout(idle_timeout, self.background.next_request()) + .await + { + Ok(Some(req)) => { + if let Err(e) = self.send_request(req).await { + break e; + } + } + // We didn't get another request for some reason, this + // probably means every client is dropped, so we just + // close. + Ok(None) => { + break ClientError::Closed; + } + // We hit idle timeout, close the connection + Err(_) => break ClientError::Closed, + } + } else { + match self.background.next_request().await { + Some(req) => { + if let Err(e) = self.send_request(req).await { + break e; + } + } + None => { + break ClientError::Closed; + } + } + } + }; + } +} + +fn read_stream( + reader: ReadHalf, +) -> impl Stream, ClientError>> { + stream::unfold(Some(reader), |reader| async { + let Some(mut reader) = reader else { + return None; + }; + + let res: Result, ClientError> = loop { + // First read the shim + let mut shim_buf = [0u8; 2]; + if let Err(err) = reader.read_exact(&mut shim_buf).await { + break Err(ClientError::from(SocketError::Receive( + err.kind(), + ))); + } + let shim = u16::from_be_bytes(shim_buf) as usize; + + // Read a response + let mut buf = vec![0u8; shim]; + if let Err(err) = reader.read_exact(&mut buf).await { + break Err(SocketError::Receive(err.kind()).into()); + } + + let Ok(msg) = Message::parse_bytes_by_ref(&buf) else { + break Err(ClientError::GarbageResponse); + }; + + break Ok(CloneFrom::clone_from(msg)); + }; + + let reader = if res.is_ok() { Some(reader) } else { None }; + + Some((res, reader)) + }) +} diff --git a/src/new_client/redundant.rs b/src/new_client/redundant.rs new file mode 100644 index 00000000..5bb72a4e --- /dev/null +++ b/src/new_client/redundant.rs @@ -0,0 +1,362 @@ +//! Multiplexing requests over redundant transports. +//! +//! This module offers a client-side transport for adding redundancy to a DNS +//! query pipeline. A [`RedundantClient`] can be created and multiple +//! equivalent transports can be added to it. Requests routed through the +//! [`RedundantClient`] will first be sent to the fastest transport, then to +//! the second-fastest, etc. Statistics about the response time for each +//! transport are collected and used to estimate the average and an upper +//! bound; [`RedundantClient`] uses this to decide how long to wait before +//! trying each next transport. + +use core::time::Duration; +use std::{boxed::Box, sync::Arc, vec::Vec}; + +use futures_util::{stream::FuturesUnordered, StreamExt}; +use parking_lot::{Mutex, RwLock}; +use rand::Rng; +use tokio::time::{timeout_at, Instant}; + +use crate::new_base::name::UnparsedName; +use crate::new_base::parse::SplitMessageBytes; +use crate::new_base::{Message, Question, RType, Record, UnparsedRecordData}; + +use super::{BoxClient, Client, ClientError}; + +#[derive(Clone, Debug, Default)] +pub struct RedundantConfig { + /// Defer transport errors. + pub defer_transport_error: bool, + + /// Defer replies that report Refused. + pub defer_refused: bool, + + /// Defer replies that report ServFail. + pub defer_servfail: bool, +} + +/// A client containing multiple sub-clients it will query. +/// +/// The fastest connection will generally be used by this transport. +/// +/// The clients to put into this client should generally be long-lived. For +/// example, adding a single TCP client might not be good idea, because that +/// connection might get closed. A multi TCP client is therefore a better +/// fit. +#[derive(Default)] +pub struct RedundantClient { + config: RedundantConfig, + clients: RwLock>>, +} + +impl RedundantClient { + pub fn new() -> Self { + Self::default() + } + + pub fn with_config(config: RedundantConfig) -> Self { + Self { + config, + clients: Default::default(), + } + } + + pub fn add_client(&self, client: impl BoxClient + 'static) { + let subclient = Arc::new(SubClient { + inner: Box::new(client), + stats: Default::default(), + }); + self.clients.write().push(subclient) + } + + fn sort_clients(&self) -> Vec { + let mut clients: Vec = + self.clients.read().iter().map(RequestClient::new).collect(); + + // Occasionally probe a random transport. + let mut rng = rand::thread_rng(); + if clients.len() > 1 && rng.gen_bool(0.05) { + let swap_idx = rng.gen_range(0..clients.len()); + clients.swap(0, swap_idx); + clients[1..].sort_unstable_by_key(|c| c.timeout); + clients[0].timeout = clients[0].timeout.min(clients[1].timeout); + } else { + // Sort the client by lowest timeout; we will query in this order. + clients.sort_unstable_by_key(|c| c.timeout); + } + + clients + } + + /// Determine whether a successful response should be skipped. + /// + /// We skip a response if the `RCODE` is `SERVFAIL` or `REFUSED` + fn skip(&self, msg: &Message) -> bool { + // We match on SERVFAIL and REFUSED. If the normal rcode matches that + // we have to ensure that the extended rcode is 0. + match msg.header.flags.rcode() { + 2 /* SERVFAIL */ if self.config.defer_servfail => { + matches!(find_opt_rcode(msg), Some(0) | None) + } + 5 /* REFUSED */ if self.config.defer_refused => { + matches!(find_opt_rcode(msg), Some(0) | None) + } + _ => false + } + } +} + +impl Client for RedundantClient { + async fn request( + &self, + message: &Message, + ) -> Result, ClientError> { + // This will be our view of the clients for this request. + // We sort them based on their timeout and iterate over them in that + // order. + let clients = self.sort_clients(); + if clients.is_empty() { + return Err(ClientError::NoTransportAvailable); + } + let mut clients = clients.into_iter(); + + let mut futs = FuturesUnordered::new(); + + // The time at which the next request should be sent out. + let mut next_request_time = Instant::now(); + + // This will hold the result of requests that fail or that we skip, + // so we can return them later when the subsequent requests also + // fail. + let mut deferred_result = None; + + loop { + match timeout_at(next_request_time, futs.next()).await { + Ok(Some(res)) => { + // got some response, so we decide whether to return it, + // store it into the deferred_result or discard it. + let res: Result, _> = res; + match res { + Ok(msg) if self.skip(&msg) => { + if let Some(Err(_)) | None = deferred_result { + deferred_result = Some(Ok(msg)); + } + } + Err(err) if self.config.defer_transport_error => { + if deferred_result.is_none() { + deferred_result = Some(Err(err)); + } + } + // It's not one of the cases we defer or skip, so we + // return it! + result => { + return result; + } + } + } + // On a timeout or empty set of futures we start a new + // request. + // + // An empty set of futures happens in two cases: + // 1. We haven't send out any requests yet + // 2. All sent requests have been resolved, which means + // we can send more requests immediately. + Ok(None) | Err(_) => { + if let Some(RequestClient { client, timeout }) = + clients.next() + { + futs.push(client.request(message)); + next_request_time = Instant::now() + timeout; + } else { + return deferred_result + .unwrap_or(Err(ClientError::Bug)); + } + } + } + } + } +} + +struct RequestClient { + client: Arc, + timeout: Duration, +} + +impl RequestClient { + fn new(client: &Arc) -> Self { + let timeout = client.stats.lock().timeout; + RequestClient { + timeout, + client: client.clone(), + } + } +} + +struct SubClient { + inner: Box, + stats: Mutex, +} + +impl SubClient { + async fn request( + self: Arc, + request: &Message, + ) -> Result, ClientError> { + /// A drop guard for collecting statistics. + struct Guard<'a> { + /// Whether the request actually finished. + finished: bool, + + /// When the request started. + start_time: Instant, + + /// The transport statistics. + stats: &'a Mutex, + } + + impl Drop for Guard<'_> { + fn drop(&mut self) { + let elapsed = self.start_time.elapsed(); + let mut stats = self.stats.lock(); + + // Update on completion, or if the request took too long. + if self.finished || elapsed.as_secs_f64() > stats.mean { + stats.account(elapsed); + } + } + } + + // Collect statistics even if the future is canceled. + let mut guard = Guard { + finished: false, + start_time: Instant::now(), + stats: &self.stats, + }; + + // Perform the actual request. + let result = self.inner.dyn_request(request).await; + + // Inform the drop guard that the request completed. + guard.finished = true; + + result + } +} + +/// Statistics about a transport. +#[derive(Clone, Debug)] +struct SubClientStats { + /// The average response time in the window. + /// + /// If this is NaN, the window was empty. + mean: f64, + + /// The average of the square of the response time in the window. + /// + /// If this is NaN, the window was empty. + mean_sq: f64, + + /// A computed timeout for requests to the transport. + /// + /// This value is three standard deviations past the mean. Assuming the + /// transport request times follow a normal distribution, there is a 99.7% + /// chance a random transport request will fit within this timeout. + timeout: Duration, +} + +impl Default for SubClientStats { + fn default() -> Self { + Self { + mean: f64::NAN, + mean_sq: f64::NAN, + timeout: Duration::from_millis(300), + } + } +} + +impl SubClientStats { + /// Account for the given response time. + fn account(&mut self, rt: Duration) { + let rt = rt.as_secs_f64(); + + if self.mean.is_nan() { + // This is the first response time -- overwrite the averages. + self.mean = rt; + self.mean_sq = rt * rt; + } else { + // Adjust the averages by 1/8th. + // + // After 8 iterations of the same response time, the previous + // average has a weight of about 34%. After 8 more iterations, + // its weight is about 12%. + self.mean = (rt + 7. * self.mean) / 8.; + self.mean_sq = (rt * rt + 7. * self.mean_sq) / 8.; + } + + // Compute the variance and standard deviation. + let variance = self.mean_sq - self.mean * self.mean; + let std_dev = variance.max(0.).sqrt(); + + // Determine the appropriate timeout value. + self.timeout = Duration::from_secs_f64(self.mean + 3. * std_dev); + } +} + +/// Find the extended RCODE in the message. +/// +/// Note that the returned `u8` only contains the upper 8 bits of the +/// `RCODE`, i.e. the part stored in the `OPT` record. +/// +/// `None` is returned on parse errors. +/// +/// We have to write this here because new_base is lacking some proper +/// handling of the (extended) `RCODE`. +fn find_opt_rcode(msg: &Message) -> Option { + let counts = msg.header.counts; + + let mut offset = 0; + for _ in 0..counts.questions.get() { + let (_, rest) = Question::<&UnparsedName>::split_message_bytes( + &msg.contents, + offset, + ) + .ok()?; + offset = rest; + } + + for _ in 0..counts.answers.get() { + let (_, rest) = Record::<&UnparsedName, &UnparsedRecordData>::split_message_bytes( + &msg.contents, + offset, + ).ok()?; + offset = rest; + } + + for _ in 0..counts.authorities.get() { + let (_, rest) = Record::<&UnparsedName, &UnparsedRecordData>::split_message_bytes( + &msg.contents, + offset, + ).ok()?; + offset = rest; + } + + for _ in 0..counts.additional.get() { + let (r, rest) = + Record::<&UnparsedName, &UnparsedRecordData>::split_message_bytes( + &msg.contents, + offset, + ) + .ok()?; + + if let RType::OPT = r.rtype { + // The extension of the rcode is specified as the first 8 bits + // of the TTL field in RFC 6891. + let ttl_bytes: [u8; 4] = r.ttl.value.get().to_be_bytes(); + return Some(ttl_bytes[0]); + } + + offset = rest; + } + + None +} diff --git a/src/new_client/tcp.rs b/src/new_client/tcp.rs index dd316339..6fbfdeb6 100644 --- a/src/new_client/tcp.rs +++ b/src/new_client/tcp.rs @@ -73,6 +73,7 @@ use core::convert::Infallible; use core::mem; +use std::boxed::Box; use std::sync::Arc; use std::time::Duration; use std::vec::Vec; @@ -87,12 +88,11 @@ use tokio::sync::{Mutex, Notify, Semaphore}; use tokio::task::{yield_now, JoinHandle}; use tokio::time::timeout; -use crate::new_base::build::BuilderContext; -use crate::new_base::wire::{AsBytes, ParseBytesByRef}; +use crate::new_base::wire::{BuildBytes, ParseBytesByRef, SizePrefixed, U16}; use crate::new_base::Message; +use crate::utils::CloneFrom; -use super::exchange::{Exchange, ParsedMessage}; -use super::{Client, ClientError, SocketError}; +use super::{Client, ClientError, ExtendedMessageBuilder, SocketError}; /// Configuration for a stream transport connection. #[derive(Clone, Debug)] @@ -147,7 +147,11 @@ pub struct TcpClient { } impl TcpClient { - pub fn new(stream: TcpStream, config: TcpConfig) -> Self { + pub fn new(stream: TcpStream) -> Self { + Self::with_config(stream, Default::default()) + } + + pub fn with_config(stream: TcpStream, config: TcpConfig) -> Self { let (read, write) = tokio::io::split(stream); let state = Arc::new(TcpClientState { @@ -172,10 +176,10 @@ impl TcpClient { } impl Client for TcpClient { - async fn request<'a>( + async fn request( &self, - exchange: &mut Exchange<'a>, - ) -> Result<(), ClientError> { + mut request: ExtendedMessageBuilder<'_, '_>, + ) -> Result, ClientError> { let _permit = self .state .request_count @@ -193,16 +197,7 @@ impl Client for TcpClient { // The channels are stored in a slab. The indices generated by the // slab are used as message IDs. - // We allocate the space for the maximum DNS message size and 2 - // additional bytes for the shim. We start writing the message with - // an offset of 2 - let mut buffer = vec![0u8; 65536 + 2]; - let mut context = BuilderContext::default(); - let mut request_builder = exchange - .request - .build(&mut context, &mut buffer[2..]) - .map_err(|_| ClientError::TruncatedRequest)?; - + // First determine the id for the message let (tx, rx) = oneshot::channel(); let message_id = { let mut ids = self.state.ids.lock().await; @@ -219,18 +214,24 @@ impl Client for TcpClient { } }; - request_builder.header_mut().id.set(message_id); - let request_len = request_builder.message().as_bytes().len(); - *buffer.first_chunk_mut().unwrap() = - (request_len as u16).to_be_bytes(); + request.set_id(message_id); - let buffer = &buffer[..request_len + 2]; + // We allocate the space for the maximum DNS message size and 2 + // additional bytes for the shim. + let mut buffer = vec![0u8; 65536 + 2]; + + // XXX: remove unwraps + let request = SizePrefixed::::new(request.build().unwrap()); + let _ = request.build_bytes(&mut buffer).unwrap(); + + let size = u16::from_be_bytes(*buffer.first_chunk::<2>().unwrap()); + buffer.truncate(2 + size as usize); // Temporary scope to drop the lock on write early { let mut write = self.state.write.lock().await; write - .write_all(buffer) + .write_all(&buffer) .await .map_err(|e| SocketError::Send(e.kind()))?; } @@ -240,7 +241,7 @@ impl Client for TcpClient { match res { // We have a message with our id, we parse it and return - Ok(Ok(Ok(msg))) => self.return_answer(msg, exchange), + Ok(Ok(Ok(msg))) => self.return_answer(msg), // We have received an error Ok(Ok(Err(err))) => Err(err), // A receive error happened on the channel. This just shouldn't @@ -261,19 +262,13 @@ impl Client for TcpClient { impl TcpClient { fn return_answer( &self, - msg: Vec, - exchange: &mut Exchange, - ) -> Result<(), ClientError> { - let Ok(msg) = Message::parse_bytes_by_ref(&msg) else { - return Err(ClientError::GarbageResponse); - }; - let Ok(parsed) = ParsedMessage::parse(msg, &mut exchange.alloc) - else { + buf: Vec, + ) -> Result, ClientError> { + let Ok(msg) = Message::parse_bytes_by_ref(&buf) else { return Err(ClientError::GarbageResponse); }; - exchange.response = parsed; - return Ok(()); + return Ok(CloneFrom::clone_from(msg)); } } diff --git a/src/new_client/udp.rs b/src/new_client/udp.rs index d0cab78a..ab613b12 100644 --- a/src/new_client/udp.rs +++ b/src/new_client/udp.rs @@ -49,18 +49,18 @@ use core::net::SocketAddr; use core::time::Duration; +use std::boxed::Box; use std::io; use tokio::net::UdpSocket; use tokio::sync::Semaphore; use tracing::trace; -use crate::new_base::build::BuilderContext; use crate::new_base::wire::{AsBytes, ParseBytesByRef}; use crate::new_base::Message; +use crate::utils::CloneFrom; -use super::exchange::{Exchange, ParsedMessage}; -use super::{Client, ClientError, SocketError}; +use super::{Client, ClientError, ExtendedMessageBuilder, SocketError}; #[derive(Clone, Debug)] pub struct UdpConfig { @@ -101,10 +101,10 @@ impl UdpClient { } impl Client for UdpClient { - async fn request<'a>( + async fn request( &self, - exchange: &mut Exchange<'a>, - ) -> Result<(), ClientError> { + mut request: ExtendedMessageBuilder<'_, '_>, + ) -> Result, ClientError> { let _permit = self .semaphore .acquire() @@ -112,27 +112,22 @@ impl Client for UdpClient { .expect("the semaphore is never closed and not exposed"); if let Some(size) = self.config.udp_payload_size { - exchange.request.set_max_udp_payload_size(size); + request.set_udp_max_payload_size(size); } - let mut buffer = vec![0u8; 65536]; - let mut context = BuilderContext::default(); - let mut request_builder = exchange - .request - .build(&mut context, &mut buffer) - .map_err(|_| ClientError::TruncatedRequest)?; + // XXX: remove unwrap + let request = request.build().unwrap(); let mut response_buffer = vec![0u8; self.config.recv_size]; for _ in 0..(1 + self.config.max_retries) { - request_builder.header_mut().id.set(rand::random()); - let request_message = request_builder.message(); + request.header.id.set(rand::random()); // We create a new UDP socket for each retry, to follow // RFC 5452's recommendations of using unpredictable source port // numbers. let response_result = send_udp_request( &mut *response_buffer, - request_message, + request, self.addr, self.config.read_timeout, ) @@ -145,16 +140,7 @@ impl Client for UdpClient { let response_message = response_result?; - let Ok(parsed) = - ParsedMessage::parse(response_message, &mut exchange.alloc) - else { - // The message turned out to be garbage, continue the loop - // to ask the server again. - continue; - }; - - exchange.response = parsed; - return Ok(()); + return Ok(CloneFrom::clone_from(response_message)); } drop(_permit);