WEBDIENSTE/NLnetLabs-domain
1
0
Fork 0
mirror of https://github.com/NLnetLabs/domain.git synced 2026-08-17 21:15:41 +02:00
Code Issues Packages Projects Releases Wiki Activity

client stuuffff

Browse Source
This commit is contained in:
Terts Diepraam
2025-04-09 16:51:47 +02:00
parent 7aaadb4679
commit d18253a9d6
8 changed files with 1316 additions and 719 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/new_client.rs
+125 -54
View File
@@ -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::<RevNameBuf>::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),
}
}
src/new_client/exchange.rs
-599
View File
@@ -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<Metadata>,
}
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<Metadata>,
}
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<Question<&'a RevName>>,
/// Answer records in the message.
pub answers: Vec<Record<&'a RevName, RecordData<'a, &'a RevName>>>,
/// Authority records in the message.
pub authorities: Vec<Record<&'a RevName, RecordData<'a, &'a RevName>>>,
/// 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<Record<&'a RevName, RecordData<'a, &'a RevName>>>,
/// 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<EdnsOption<'a>>,
}
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<Self, ParseError> {
type ParsedQuestion = Question<RevNameBuf>;
type ParsedRecord<'a> =
Record<RevNameBuf, RecordData<'a, RevNameBuf>>;
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<MessageBuilder<'b, 'c>, 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<dyn Any + Send + 'static>` 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<dyn Any + Send + 'static>,
}
impl Metadata {
/// Wrap an object in [`Metadata`].
pub fn new<T: Any + Send + 'static>(object: T) -> Self {
let type_id = TypeId::of::<T>();
let object = Box::new(object) as Box<dyn Any + Send + 'static>;
Self { type_id, object }
}
/// Check whether this is metadata of a certain type.
pub fn is<T: Any + Send + 'static>(&self) -> bool {
self.type_id == TypeId::of::<T>()
}
/// Try downcasting to a reference of a particular type.
pub fn try_as<T: Any + Send + 'static>(&self) -> Option<&T> {
if !self.is::<T>() {
return None;
}
let pointer: *const (dyn Any + Send + 'static) = &*self.object;
// SAFETY: 'pointer' was created by 'Box<T>::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::<T>() })
}
/// Try downcasting to a mutable reference of a particular type.
pub fn try_as_mut<T: Any + Send + 'static>(&mut self) -> Option<&mut T> {
if !self.is::<T>() {
return None;
}
let pointer: *mut (dyn Any + Send + 'static) = &mut *self.object;
// SAFETY: 'pointer' was created by 'Box<T>::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::<T>() })
}
/// Try moving this object out of the [`Metadata`].
pub fn try_into<T: Any + Send + 'static>(self) -> Result<T, Self> {
if !self.is::<T>() {
return Err(self);
}
let pointer: *mut _ = Box::into_raw(self.object);
// SAFETY: 'pointer' was created by 'Box<T>::into_raw()', and thus is
// safe to move into the same 'Box<T>'.
Ok(*unsafe { Box::from_raw(pointer.cast::<T>()) })
}
}
//----------- 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<T>(&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<T: Copy>(&mut self, src: &[T]) -> &'a mut [T] {
self.inner.alloc_slice_copy(src)
}
pub fn alloc_unsized<T: ?Sized + UnsizedClone>(
&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.
src/new_client/load_balancer.rs
+268
View File
@@ -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<u64>,
/// Interval over which the burst is counted.
burst_interval: Duration,
}
pub struct SubClient {
config: SubClientConfig,
client: Box<dyn BoxClient>,
stats: Mutex<SubClientStats>,
}
pub struct LoadBalancerClient {
config: LoadBalancerConfig,
clients: RwLock<Vec<Arc<SubClient>>>,
}
impl<C: BoxClient + 'static> From<C> 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<RequestClient> {
let mut clients: Vec<RequestClient> = 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<u8>) -> 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<Vec<u8>, 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<SubClient>,
timeout: Duration,
}
impl RequestClient {
fn new(client: &Arc<SubClient>) -> 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<u8> {
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
}
src/new_client/mod.rs
+92 -5
View File
@@ -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<Box<Message>, ClientError>;
}
pub trait BoxClient {
fn dyn_request<'a, 'b: 'a, 'c: 'a>(
&'a self,
request: ExtendedMessageBuilder<'b, 'c>,
) -> Pin<Box<dyn Future<Output = Result<Box<Message>, ClientError>> + 'a>>;
}
impl<T: Client> BoxClient for T {
fn dyn_request<'a, 'b: 'a, 'c: 'a>(
&'a self,
request: ExtendedMessageBuilder<'b, 'c>,
) -> Pin<Box<dyn Future<Output = Result<Box<Message>, ClientError>> + 'a>>
{
Box::pin(self.request(request))
}
}
#[derive(Clone, Debug, Default)]
pub struct EdnsRecordBuilder<'a> {
pub header: EdnsHeader,
pub options: Vec<EdnsOption<'a>>,
}
#[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<EdnsRecordBuilder<'b>>,
}
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<SocketError> for ClientError {
src/new_client/multi_tcp.rs
+427
View File
@@ -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<Duration>,
}
impl Default for TcpConfig {
fn default() -> Self {
Self {
response_timeout: Duration::from_secs(19),
idle_timeout: None,
}
}
}
struct AbortJoinHandle<T>(JoinHandle<T>);
impl<T> Drop for AbortJoinHandle<T> {
fn drop(&mut self) {
self.0.abort();
}
}
struct Request {
callback_send: oneshot::Sender<Result<Box<Message>, ClientError>>,
message: Box<Message>,
}
pub struct TcpClient {
/// Ensure that the read loop lives as long as the client
_background: AbortJoinHandle<()>,
config: TcpConfig,
send: mpsc::Sender<Request>,
}
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<Box<Message>, 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<Result<Box<Message>, ClientError>>,
},
TimedOut,
}
struct Background {
addr: SocketAddr,
config: TcpConfig,
requests: mpsc::Receiver<Request>,
}
impl Background {
fn new(
addr: SocketAddr,
config: TcpConfig,
requests: mpsc::Receiver<Request>,
) -> 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<Request> {
self.requests.recv().await
}
async fn connect<'a>(
&'a mut self,
) -> io::Result<
Connection<'a, impl Stream<Item = Result<Box<Message>, 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<Item = Result<Box<Message>, ClientError>> + Unpin,
> {
stream: WriteHalf<TcpStream>,
background: &'a mut Background,
read_stream: S,
current_requests: Slab<Waiting>,
}
impl<S: Stream<Item = Result<Box<Message>, 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::<U16, _>::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<TcpStream>,
) -> impl Stream<Item = Result<Box<Message>, ClientError>> {
stream::unfold(Some(reader), |reader| async {
let Some(mut reader) = reader else {
return None;
};
let res: Result<Box<Message>, 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))
})
}
src/new_client/redundant.rs
+362
View File
@@ -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<Vec<Arc<SubClient>>>,
}
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<RequestClient> {
let mut clients: Vec<RequestClient> =
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<Box<Message>, 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<Box<Message>, _> = 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<SubClient>,
timeout: Duration,
}
impl RequestClient {
fn new(client: &Arc<SubClient>) -> Self {
let timeout = client.stats.lock().timeout;
RequestClient {
timeout,
client: client.clone(),
}
}
}
struct SubClient {
inner: Box<dyn BoxClient>,
stats: Mutex<SubClientStats>,
}
impl SubClient {
async fn request(
self: Arc<Self>,
request: &Message,
) -> Result<Box<Message>, 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<SubClientStats>,
}
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<u8> {
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
}
src/new_client/tcp.rs
+30 -35
View File
@@ -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<Box<Message>, 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::<U16, _>::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<u8>,
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<u8>,
) -> Result<Box<Message>, 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));
}
}
src/new_client/udp.rs
+12 -26
View File
@@ -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<Box<Message>, 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);