diff --git a/Cargo.lock b/Cargo.lock index 49778ae01..c53a0803b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -235,6 +235,7 @@ dependencies = [ "bytes", "chrono", "domain-macros", + "either", "futures-util", "hashbrown 0.14.5", "heapless", diff --git a/Cargo.toml b/Cargo.toml index 481da842c..071e17a21 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,8 @@ license = "BSD-3-Clause" [dependencies] domain-macros = { path = "./macros", version = "0.10.3" } +either = { version = "1.10.0", default-features = false } + arbitrary = { version = "1.4.1", optional = true, features = ["derive"] } octseq = { version = "0.5.2", default-features = false } time = { version = "0.3.1", default-features = false } @@ -58,7 +60,7 @@ bytes = ["dep:bytes", "octseq/bytes"] heapless = ["dep:heapless", "octseq/heapless"] serde = ["dep:serde", "octseq/serde"] smallvec = ["dep:smallvec", "octseq/smallvec"] -std = ["dep:hashbrown", "bytes?/std", "octseq/std", "time/std"] +std = ["dep:hashbrown", "either/use_std", "bytes?/std", "octseq/std", "time/std"] tracing = ["dep:log", "dep:tracing"] # Cryptographic backends diff --git a/src/lib.rs b/src/lib.rs index e38ebcaa7..0709aa1a3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -212,4 +212,5 @@ pub mod zonetree; pub mod new_base; pub mod new_edns; +pub mod new_net; pub mod new_rdata; diff --git a/src/new_base/question.rs b/src/new_base/question.rs index 34e6dc282..fab29aaa8 100644 --- a/src/new_base/question.rs +++ b/src/new_base/question.rs @@ -59,6 +59,7 @@ where impl<'a, N> ParseMessageBytes<'a> for Question where + // TODO: Reduce to 'ParseMessageBytes'. N: SplitMessageBytes<'a>, { fn parse_message_bytes( diff --git a/src/new_net/mod.rs b/src/new_net/mod.rs new file mode 100644 index 000000000..a28ed0b4b --- /dev/null +++ b/src/new_net/mod.rs @@ -0,0 +1,6 @@ +//! Sending and receiving DNS messages. + +#![cfg(feature = "net")] +#![cfg_attr(docsrs, doc(cfg(feature = "net")))] + +pub mod server; diff --git a/src/new_net/server/impls.rs b/src/new_net/server/impls.rs new file mode 100644 index 000000000..88293d24a --- /dev/null +++ b/src/new_net/server/impls.rs @@ -0,0 +1,1249 @@ +//! Blanket implementations for service traits. + +use core::ops::ControlFlow; + +#[cfg(feature = "std")] +use std::{boxed::Box, rc::Rc, sync::Arc, vec::Vec}; + +use either::Either::{self, Left, Right}; + +use crate::new_base::{ + build::{MessageBuilder, QuestionBuilder, RecordBuilder}, + Header, +}; + +use super::{ + LocalService, LocalServiceLayer, ProduceMessage, RequestMessage, Service, + ServiceLayer, TransformMessage, +}; + +//----------- impl Service --------------------------------------------------- + +impl Service for &T { + async fn respond(&self, request: &RequestMessage<'_>) -> Self::Producer { + T::respond(self, request).await + } +} + +impl Service for &mut T { + async fn respond(&self, request: &RequestMessage<'_>) -> Self::Producer { + T::respond(self, request).await + } +} + +#[cfg(feature = "std")] +impl Service for Box { + async fn respond(&self, request: &RequestMessage<'_>) -> Self::Producer { + T::respond(self, request).await + } +} + +#[cfg(feature = "std")] +impl Service for Arc { + async fn respond(&self, request: &RequestMessage<'_>) -> Self::Producer { + T::respond(self, request).await + } +} + +impl Service for (A, S) +where + A: ServiceLayer, + S: Service, +{ + async fn respond(&self, request: &RequestMessage<'_>) -> Self::Producer { + match self.0.respond(request).await { + ControlFlow::Continue(t) => { + Right((t, self.1.respond(request).await)) + } + ControlFlow::Break(p) => Left(p), + } + } +} + +macro_rules! impl_service_tuple { + ($($layers:ident)* .. $service:ident) => { + impl<$($layers,)* $service: ?Sized> + Service for ($($layers,)* $service,) + where + $($layers: ServiceLayer,)* + $service: Service, + { + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer { + #[allow(non_snake_case)] + let ($($layers,)* $service,) = self; + (($($layers),*,), $service).respond(request).await + } + } + }; +} + +impl_service_tuple!(A B..S); +impl_service_tuple!(A B C..S); +impl_service_tuple!(A B C D..S); +impl_service_tuple!(A B C D E..S); +impl_service_tuple!(A B C D E F..S); +impl_service_tuple!(A B C D E F G..S); +impl_service_tuple!(A B C D E F G H..S); +impl_service_tuple!(A B C D E F G H I..S); +impl_service_tuple!(A B C D E F G H I J..S); +impl_service_tuple!(A B C D E F G H I J K..S); + +//----------- impl LocalService ---------------------------------------------- + +impl LocalService for &T { + type Producer = T::Producer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer { + T::respond_local(self, request).await + } +} + +impl LocalService for &mut T { + type Producer = T::Producer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer { + T::respond_local(self, request).await + } +} + +#[cfg(feature = "std")] +impl LocalService for Box { + type Producer = T::Producer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer { + T::respond_local(self, request).await + } +} + +#[cfg(feature = "std")] +impl LocalService for Rc { + type Producer = T::Producer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer { + T::respond_local(self, request).await + } +} + +#[cfg(feature = "std")] +impl LocalService for Arc { + type Producer = T::Producer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer { + T::respond_local(self, request).await + } +} + +impl LocalService for (A, S) +where + A: LocalServiceLayer, + S: LocalService, +{ + type Producer = Either; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer { + match self.0.respond_local(request).await { + ControlFlow::Continue(t) => { + Right((t, self.1.respond_local(request).await)) + } + ControlFlow::Break(p) => Left(p), + } + } +} + +macro_rules! impl_local_service_tuple { + ($($layers:ident)* .. $service:ident) => { + impl<$($layers,)* $service: ?Sized> + LocalService for ($($layers,)* $service,) + where + $($layers: LocalServiceLayer,)* + $service: LocalService, + { + type Producer = + <(($($layers),*,), $service) as LocalService>::Producer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer { + #[allow(non_snake_case)] + let ($($layers,)* $service,) = self; + (($($layers),*,), $service).respond_local(request).await + } + } + }; +} + +impl_local_service_tuple!(A B..S); +impl_local_service_tuple!(A B C..S); +impl_local_service_tuple!(A B C D..S); +impl_local_service_tuple!(A B C D E..S); +impl_local_service_tuple!(A B C D E F..S); +impl_local_service_tuple!(A B C D E F G..S); +impl_local_service_tuple!(A B C D E F G H..S); +impl_local_service_tuple!(A B C D E F G H I..S); +impl_local_service_tuple!(A B C D E F G H I J..S); +impl_local_service_tuple!(A B C D E F G H I J K..S); + +//----------- impl ServiceLayer ---------------------------------------------- + +impl ServiceLayer for &T { + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond(self, request).await + } +} + +impl ServiceLayer for &mut T { + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond(self, request).await + } +} + +#[cfg(feature = "std")] +impl ServiceLayer for Box { + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond(self, request).await + } +} + +#[cfg(feature = "std")] +impl ServiceLayer for Arc { + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond(self, request).await + } +} + +impl ServiceLayer for (A, B) +where + A: ServiceLayer, + B: ServiceLayer, +{ + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + let at = match ::respond(&self.0, request).await { + ControlFlow::Continue(at) => at, + ControlFlow::Break(ap) => return ControlFlow::Break(Left(ap)), + }; + + match ::respond(&self.1, request).await { + ControlFlow::Continue(bt) => ControlFlow::Continue((at, bt)), + ControlFlow::Break(bp) => ControlFlow::Break(Right((at, bp))), + } + } +} + +macro_rules! impl_service_layer_tuple { + ($first:ident .. $last:ident: $($middle:ident)+) => { + impl<$first, $($middle,)+ $last: ?Sized> + ServiceLayer for ($first, $($middle,)+ $last) + where + $first: ServiceLayer, + $($middle: ServiceLayer,)+ + $last: ServiceLayer, + { + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow + { + #[allow(non_snake_case)] + let ($first, $($middle,)+ ref $last) = self; + ($first, ($($middle,)+ $last)).respond(request).await + } + } + } +} + +impl_service_layer_tuple!(A..C: B); +impl_service_layer_tuple!(A..D: B C); +impl_service_layer_tuple!(A..E: B C D); +impl_service_layer_tuple!(A..F: B C D E); +impl_service_layer_tuple!(A..G: B C D E F); +impl_service_layer_tuple!(A..H: B C D E F G); +impl_service_layer_tuple!(A..I: B C D E F G H); +impl_service_layer_tuple!(A..J: B C D E F G H I); +impl_service_layer_tuple!(A..K: B C D E F G H I J); +impl_service_layer_tuple!(A..L: B C D E F G H I J K); + +#[cfg(feature = "std")] +impl ServiceLayer for [T] { + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + let mut transformers = Vec::new(); + for layer in self { + match layer.respond(request).await { + ControlFlow::Continue(t) => transformers.push(t), + ControlFlow::Break(p) => { + return ControlFlow::Break((transformers.into(), p)); + } + } + } + ControlFlow::Continue(transformers.into()) + } +} + +#[cfg(feature = "std")] +impl ServiceLayer for Vec { + async fn respond( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + self.as_slice().respond(request).await + } +} + +//----------- impl LocalServiceLayer ----------------------------------------- + +impl LocalServiceLayer for &T { + type Producer = T::Producer; + + type Transformer = T::Transformer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond_local(self, request).await + } +} + +impl LocalServiceLayer for &mut T { + type Producer = T::Producer; + + type Transformer = T::Transformer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond_local(self, request).await + } +} + +#[cfg(feature = "std")] +impl LocalServiceLayer for Box { + type Producer = T::Producer; + + type Transformer = T::Transformer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond_local(self, request).await + } +} + +#[cfg(feature = "std")] +impl LocalServiceLayer for Rc { + type Producer = T::Producer; + + type Transformer = T::Transformer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond_local(self, request).await + } +} + +#[cfg(feature = "std")] +impl LocalServiceLayer for Arc { + type Producer = T::Producer; + + type Transformer = T::Transformer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + T::respond_local(self, request).await + } +} + +impl LocalServiceLayer for (A, B) +where + A: LocalServiceLayer, + B: LocalServiceLayer, +{ + type Producer = Either; + + type Transformer = (A::Transformer, B::Transformer); + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + let at = match self.0.respond_local(request).await { + ControlFlow::Continue(at) => at, + ControlFlow::Break(ap) => return ControlFlow::Break(Left(ap)), + }; + + match self.1.respond_local(request).await { + ControlFlow::Continue(bt) => ControlFlow::Continue((at, bt)), + ControlFlow::Break(bp) => ControlFlow::Break(Right((at, bp))), + } + } +} + +macro_rules! impl_local_service_layer_tuple { + ($first:ident .. $last:ident: $($middle:ident)+) => { + impl<$first, $($middle,)+ $last: ?Sized> + LocalServiceLayer for ($first, $($middle,)+ $last) + where + $first: LocalServiceLayer, + $($middle: LocalServiceLayer,)+ + $last: LocalServiceLayer, + { + type Producer = + <($first, ($($middle,)+ $last)) as LocalServiceLayer> + ::Producer; + + type Transformer = + <($first, ($($middle,)+ $last)) as LocalServiceLayer> + ::Transformer; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow + { + #[allow(non_snake_case)] + let ($first, $($middle,)+ ref $last) = self; + ($first, ($($middle,)+ $last)).respond_local(request).await + } + } + } +} + +impl_local_service_layer_tuple!(A..C: B); +impl_local_service_layer_tuple!(A..D: B C); +impl_local_service_layer_tuple!(A..E: B C D); +impl_local_service_layer_tuple!(A..F: B C D E); +impl_local_service_layer_tuple!(A..G: B C D E F); +impl_local_service_layer_tuple!(A..H: B C D E F G); +impl_local_service_layer_tuple!(A..I: B C D E F G H); +impl_local_service_layer_tuple!(A..J: B C D E F G H I); +impl_local_service_layer_tuple!(A..K: B C D E F G H I J); +impl_local_service_layer_tuple!(A..L: B C D E F G H I J K); + +#[cfg(feature = "std")] +impl LocalServiceLayer for [T] { + type Producer = (Box<[T::Transformer]>, T::Producer); + type Transformer = Box<[T::Transformer]>; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + let mut transformers = Vec::new(); + for layer in self { + match layer.respond_local(request).await { + ControlFlow::Continue(t) => transformers.push(t), + ControlFlow::Break(p) => { + return ControlFlow::Break((transformers.into(), p)); + } + } + } + ControlFlow::Continue(transformers.into()) + } +} + +#[cfg(feature = "std")] +impl LocalServiceLayer for Vec { + type Producer = (Box<[T::Transformer]>, T::Producer); + type Transformer = Box<[T::Transformer]>; + + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow { + self.as_slice().respond_local(request).await + } +} + +//----------- impl ProduceMessage -------------------------------------------- + +impl ProduceMessage for &mut T { + fn header(&mut self, header: &mut Header) { + T::header(self, header); + } + + fn next_question<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + T::next_question(self, builder) + } + + fn next_answer<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + T::next_answer(self, builder) + } + + fn next_authority<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + T::next_authority(self, builder) + } + + fn next_additional<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + T::next_additional(self, builder) + } +} + +#[cfg(feature = "std")] +impl ProduceMessage for Box { + fn header(&mut self, header: &mut Header) { + T::header(self, header); + } + + fn next_question<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + T::next_question(self, builder) + } + + fn next_answer<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + T::next_answer(self, builder) + } + + fn next_authority<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + T::next_authority(self, builder) + } + + fn next_additional<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + T::next_additional(self, builder) + } +} + +impl ProduceMessage for Either +where + A: ProduceMessage, + B: ProduceMessage, +{ + fn header(&mut self, header: &mut Header) { + either::for_both!(self, x => x.header(header)); + } + + fn next_question<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + either::for_both!(self, x => x.next_question(builder)) + } + + fn next_answer<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + either::for_both!(self, x => x.next_answer(builder)) + } + + fn next_authority<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + either::for_both!(self, x => x.next_authority(builder)) + } + + fn next_additional<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + either::for_both!(self, x => x.next_additional(builder)) + } +} + +impl ProduceMessage for (A, B) +where + A: TransformMessage, + B: ProduceMessage, +{ + fn header(&mut self, header: &mut Header) { + B::header(&mut self.1, header); + A::modify_header(&mut self.0, header); + } + + fn next_question<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + loop { + let mut delegate = builder.reborrow(); + let mut question = match self.1.next_question(&mut delegate) { + Some(question) => question, + None => break, + }; + + if self.0.modify_question(&mut question).is_break() { + continue; + } + + return builder.resume_question(); + } + + self.0.next_question(builder) + } + + fn next_answer<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + loop { + let mut delegate = builder.reborrow(); + let mut answer = match self.1.next_answer(&mut delegate) { + Some(answer) => answer, + None => break, + }; + + if self.0.modify_answer(&mut answer).is_break() { + continue; + } + + core::mem::drop(answer); + return builder.resume_answer(); + } + + self.0.next_answer(builder) + } + + fn next_authority<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + loop { + let mut delegate = builder.reborrow(); + let mut authority = match self.1.next_authority(&mut delegate) { + Some(authority) => authority, + None => break, + }; + + if self.0.modify_authority(&mut authority).is_break() { + continue; + } + + core::mem::drop(authority); + return builder.resume_authority(); + } + + self.0.next_authority(builder) + } + + fn next_additional<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + loop { + let mut delegate = builder.reborrow(); + let mut additional = match self.1.next_additional(&mut delegate) { + Some(additional) => additional, + None => break, + }; + + if self.0.modify_additional(&mut additional).is_break() { + continue; + } + + core::mem::drop(additional); + return builder.resume_additional(); + } + + self.0.next_additional(builder) + } +} + +macro_rules! impl_produce_message_tuple { + ($first:ident .. $last:ident: $($middle:ident)*) => { + impl<$first, $($middle,)* $last: ?Sized> + ProduceMessage for ($first, $($middle,)* $last) + where + $first: TransformMessage, + $($middle: TransformMessage,)* + $last: ProduceMessage, + { + fn header(&mut self, header: &mut Header) { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .header(header) + } + + fn next_question<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .next_question(builder) + } + + fn next_answer<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .next_answer(builder) + } + + fn next_authority<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .next_authority(builder) + } + + fn next_additional<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .next_additional(builder) + } + } + }; +} + +impl_produce_message_tuple!(A..C: B); +impl_produce_message_tuple!(A..D: B C); +impl_produce_message_tuple!(A..E: B C D); +impl_produce_message_tuple!(A..F: B C D E); +impl_produce_message_tuple!(A..G: B C D E F); +impl_produce_message_tuple!(A..H: B C D E F G); +impl_produce_message_tuple!(A..I: B C D E F G H); +impl_produce_message_tuple!(A..J: B C D E F G H I); +impl_produce_message_tuple!(A..K: B C D E F G H I J); +impl_produce_message_tuple!(A..L: B C D E F G H I J K); + +impl ProduceMessage for [T] { + fn header(&mut self, header: &mut Header) { + if let [ref mut layers @ .., ref mut last] = self { + last.header(header); + for layer in layers.iter_mut().rev() { + layer.modify_header(header); + } + } + } + + fn next_question<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + let mut layers = self; + while let [ref mut nested @ .., ref mut last] = layers { + let mut delegate = builder.reborrow(); + let mut question = match last.next_question(&mut delegate) { + Some(question) => question, + None => break, + }; + + match nested + .iter_mut() + .rev() + .try_for_each(|layer| layer.modify_question(&mut question)) + { + ControlFlow::Continue(()) => { + return builder.resume_question(); + } + + ControlFlow::Break(()) => {} + } + + layers = nested; + } + + None + } + + fn next_answer<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + let mut layers = self; + while let [ref mut nested @ .., ref mut last] = layers { + let mut delegate = builder.reborrow(); + let mut answer = match last.next_answer(&mut delegate) { + Some(answer) => answer, + None => break, + }; + + match nested + .iter_mut() + .rev() + .try_for_each(|layer| layer.modify_answer(&mut answer)) + { + ControlFlow::Continue(()) => { + core::mem::drop(answer); + return builder.resume_answer(); + } + + ControlFlow::Break(()) => {} + } + + layers = nested; + } + + None + } + + fn next_authority<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + let mut layers = self; + while let [ref mut nested @ .., ref mut last] = layers { + let mut delegate = builder.reborrow(); + let mut authority = match last.next_authority(&mut delegate) { + Some(authority) => authority, + None => break, + }; + + match nested + .iter_mut() + .rev() + .try_for_each(|layer| layer.modify_authority(&mut authority)) + { + ControlFlow::Continue(()) => { + core::mem::drop(authority); + return builder.resume_authority(); + } + + ControlFlow::Break(()) => {} + } + + layers = nested; + } + + None + } + + fn next_additional<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + let mut layers = self; + while let [ref mut nested @ .., ref mut last] = layers { + let mut delegate = builder.reborrow(); + let mut additional = match last.next_additional(&mut delegate) { + Some(additional) => additional, + None => break, + }; + + match nested.iter_mut().rev().try_for_each(|layer| { + layer.modify_additional(&mut additional) + }) { + ControlFlow::Continue(()) => { + core::mem::drop(additional); + return builder.resume_additional(); + } + + ControlFlow::Break(()) => {} + } + + layers = nested; + } + + None + } +} + +#[cfg(feature = "std")] +impl ProduceMessage for Vec { + fn header(&mut self, header: &mut Header) { + self.as_mut_slice().header(header) + } + + fn next_question<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + self.as_mut_slice().next_question(builder) + } + + fn next_answer<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + self.as_mut_slice().next_answer(builder) + } + + fn next_authority<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + self.as_mut_slice().next_authority(builder) + } + + fn next_additional<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + self.as_mut_slice().next_additional(builder) + } +} + +//----------- TransformMessage ----------------------------------------------- + +impl TransformMessage for &mut T { + fn modify_header(&mut self, header: &mut Header) { + T::modify_header(self, header); + } + + fn modify_question( + &mut self, + builder: &mut QuestionBuilder<'_>, + ) -> ControlFlow<()> { + T::modify_question(self, builder) + } + + fn modify_answer( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + T::modify_answer(self, builder) + } + + fn modify_authority( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + T::modify_authority(self, builder) + } + + fn modify_additional( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + T::modify_additional(self, builder) + } +} + +#[cfg(feature = "std")] +impl TransformMessage for Box { + fn modify_header(&mut self, header: &mut Header) { + T::modify_header(self, header); + } + + fn modify_question( + &mut self, + builder: &mut QuestionBuilder<'_>, + ) -> ControlFlow<()> { + T::modify_question(self, builder) + } + + fn modify_answer( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + T::modify_answer(self, builder) + } + + fn modify_authority( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + T::modify_authority(self, builder) + } + + fn modify_additional( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + T::modify_additional(self, builder) + } +} + +impl TransformMessage for Either +where + A: TransformMessage, + B: TransformMessage, +{ + fn modify_header(&mut self, header: &mut Header) { + either::for_both!(self, x => x.modify_header(header)); + } + + fn modify_question( + &mut self, + builder: &mut QuestionBuilder<'_>, + ) -> ControlFlow<()> { + either::for_both!(self, x => x.modify_question(builder)) + } + + fn modify_answer( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + either::for_both!(self, x => x.modify_answer(builder)) + } + + fn modify_authority( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + either::for_both!(self, x => x.modify_authority(builder)) + } + + fn modify_additional( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + either::for_both!(self, x => x.modify_additional(builder)) + } +} + +impl TransformMessage for (A, B) +where + A: TransformMessage, + B: TransformMessage, +{ + fn modify_header(&mut self, header: &mut Header) { + self.1.modify_header(header); + self.0.modify_header(header); + } + + fn modify_question( + &mut self, + builder: &mut QuestionBuilder<'_>, + ) -> ControlFlow<()> { + self.1.modify_question(builder)?; + self.0.modify_question(builder)?; + ControlFlow::Continue(()) + } + + fn modify_answer( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.1.modify_answer(builder)?; + self.0.modify_answer(builder)?; + ControlFlow::Continue(()) + } + + fn modify_authority( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.1.modify_authority(builder)?; + self.0.modify_authority(builder)?; + ControlFlow::Continue(()) + } + + fn modify_additional( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.1.modify_additional(builder)?; + self.0.modify_additional(builder)?; + ControlFlow::Continue(()) + } +} + +macro_rules! impl_transform_message_tuple { + ($first:ident .. $last:ident: $($middle:ident)*) => { + impl<$first, $($middle,)* $last: ?Sized> + TransformMessage for ($first, $($middle,)* $last) + where + $first: TransformMessage, + $($middle: TransformMessage,)* + $last: TransformMessage, + { + fn modify_header(&mut self, header: &mut Header) { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .modify_header(header) + } + + fn modify_question( + &mut self, + builder: &mut QuestionBuilder<'_>, + ) -> ControlFlow<()> { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .modify_question(builder) + } + + fn modify_answer( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .modify_answer(builder) + } + + fn modify_authority( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .modify_authority(builder) + } + + fn modify_additional( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + #[allow(non_snake_case)] + let ($first, $($middle,)* ref mut $last) = self; + ($first, ($($middle,)* $last)) + .modify_additional(builder) + } + } + }; +} + +impl_transform_message_tuple!(A..C: B); +impl_transform_message_tuple!(A..D: B C); +impl_transform_message_tuple!(A..E: B C D); +impl_transform_message_tuple!(A..F: B C D E); +impl_transform_message_tuple!(A..G: B C D E F); +impl_transform_message_tuple!(A..H: B C D E F G); +impl_transform_message_tuple!(A..I: B C D E F G H); +impl_transform_message_tuple!(A..J: B C D E F G H I); +impl_transform_message_tuple!(A..K: B C D E F G H I J); +impl_transform_message_tuple!(A..L: B C D E F G H I J K); + +impl TransformMessage for [T] { + fn modify_header(&mut self, header: &mut Header) { + self.iter_mut() + .rev() + .for_each(|layer| layer.modify_header(header)); + } + + fn modify_question( + &mut self, + builder: &mut QuestionBuilder<'_>, + ) -> ControlFlow<()> { + self.iter_mut() + .rev() + .try_for_each(|layer| layer.modify_question(builder)) + } + + fn modify_answer( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.iter_mut() + .rev() + .try_for_each(|layer| layer.modify_answer(builder)) + } + + fn modify_authority( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.iter_mut() + .rev() + .try_for_each(|layer| layer.modify_authority(builder)) + } + + fn modify_additional( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.iter_mut() + .rev() + .try_for_each(|layer| layer.modify_additional(builder)) + } +} + +#[cfg(feature = "std")] +impl TransformMessage for Vec { + fn modify_header(&mut self, header: &mut Header) { + self.as_mut_slice().modify_header(header) + } + + fn modify_question( + &mut self, + builder: &mut QuestionBuilder<'_>, + ) -> ControlFlow<()> { + self.as_mut_slice().modify_question(builder) + } + + fn modify_answer( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.as_mut_slice().modify_answer(builder) + } + + fn modify_authority( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.as_mut_slice().modify_authority(builder) + } + + fn modify_additional( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + self.as_mut_slice().modify_additional(builder) + } +} diff --git a/src/new_net/server/mod.rs b/src/new_net/server/mod.rs new file mode 100644 index 000000000..e02688d62 --- /dev/null +++ b/src/new_net/server/mod.rs @@ -0,0 +1,409 @@ +//! Responding to DNS requests. +//! +//! # Architecture +//! +//! A _transport_ implements a network interface allowing it to receive DNS +//! requests and return DNS responses. Transports can be implemented on UDP, +//! TCP, TLS, etc., and users can implement their own transports. +//! +//! A _service_ implements the business logic of handling a DNS request and +//! building a DNS response. A service can be composed of multiple _layers_, +//! each of which can inspect the request and prepare part of the response. +//! Many common layers are already implemented, but users can define more. + +#![cfg(feature = "unstable-server-transport")] +#![cfg_attr(docsrs, doc(cfg(feature = "unstable-server-transport")))] + +use core::{future::Future, ops::ControlFlow}; + +use crate::new_base::{ + build::{MessageBuilder, QuestionBuilder, RecordBuilder}, + Header, +}; + +mod impls; + +pub mod request; +pub use request::RequestMessage; + +//----------- Service -------------------------------------------------------- + +/// A (multi-threaded) DNS service, that computes responses for requests. +/// +/// Given a DNS request message, a service computes an appropriate response. +/// Services are usually wrapped in a network transport that receives requests +/// and returns the service's responses. +/// +/// Use [`LocalService`] for a single-threaded equivalent. +/// +/// # Layering +/// +/// Additional functionality can be added to a service by prefixing it with +/// service layers, usually in a tuple. A number of blanket implementations +/// are provided to simplify this. +pub trait Service: LocalService + Sync { + /// Respond to a DNS request. + /// + /// The provided consumer must have been provided the entire DNS request + /// message. This method will use the extracted information to formulate + /// a response message, in the form of a producer type. + /// + /// The returned future implements [`Send`]. Use [`LocalService`] and + /// [`LocalService::respond_local()`] if [`Send`] is not necessary. + fn respond( + &self, + request: &RequestMessage<'_>, + ) -> impl Future + Send; +} + +//----------- LocalService --------------------------------------------------- + +/// A (single-threaded) DNS service, that computes responses for requests. +/// +/// Given a DNS request message, a service computes an appropriate response. +/// Services are usually wrapped in a network transport that receives requests +/// and returns the service's responses. +/// +/// Use [`Service`] for a multi-threaded equivalent. +/// +/// # Layering +/// +/// Additional functionality can be added to a service by prefixing it with +/// service layers, usually in a tuple. A number of blanket implementations +/// are provided to simplify this. +pub trait LocalService { + /// A producer of DNS responses. + /// + /// This type returns components to insert in a DNS response message. It + /// is constructed by [`Self::respond_local()`]. + /// + /// # Lifetimes + /// + /// The producer can borrow from the request message (`'req`). Note that + /// it cannot borrow from the response message. + type Producer: ProduceMessage; + + /// Respond to a DNS request. + /// + /// The provided consumer must have been provided the entire DNS request + /// message. This method will use the extracted information to formulate + /// a response message, in the form of a producer type. + /// + /// The returned future does not implement [`Send`]. Use [`Service`] and + /// [`Service::respond()`] for a [`Send`]-implementing version. + #[allow(async_fn_in_trait)] + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> Self::Producer; +} + +//----------- ServiceLayer --------------------------------------------------- + +/// A (multi-threaded) layer wrapping a DNS [`Service`]. +/// +/// A layer can be wrapped around a service, inspecting the requests sent to +/// it and transforming the responses returned by it. +/// +/// Use [`LocalServiceLayer`] for a single-threaded equivalent. +/// +/// # Combinations +/// +/// Layers can be combined (usually in a tuple) into larger layers. A number +/// of blanket implementations are provided to simplify this. +pub trait ServiceLayer: + LocalServiceLayer + Sync +{ + /// Respond to a DNS request. + /// + /// The provided consumer must have been provided the entire DNS request + /// message. If the request should be forwarded through to the wrapped + /// service, [`ControlFlow::Continue`] is returned, with a transformer to + /// modify the wrapped service's response. However, if the request should + /// be responded to directly by this layer, without any interaction from + /// the wrapped service, [`ControlFlow::Break`] is returned. + /// + /// The returned future implements [`Send`]. Use [`LocalServiceLayer`] + /// and [`LocalServiceLayer::respond_local()`] if [`Send`] is not + /// necessary. + fn respond( + &self, + request: &RequestMessage<'_>, + ) -> impl Future> + Send; +} + +//----------- LocalServiceLayer ---------------------------------------------- + +/// A (single-threaded) layer wrapping a DNS [`Service`]. +/// +/// A layer can be wrapped around a service, inspecting the requests sent to +/// it and transforming the responses returned by it. +/// +/// Use [`ServiceLayer`] for a multi-threaded equivalent. +/// +/// # Combinations +/// +/// Layers can be combined (usually in a tuple) into larger layers. A number +/// of blanket implementations are provided to simplify this. +pub trait LocalServiceLayer { + /// A producer of DNS responses. + /// + /// This type returns components to insert in a DNS response message. It + /// is constructed by [`Self::respond_local()`], if a response is returned + /// early (without the wrapped service interacting with it). + /// + /// # Lifetimes + /// + /// The producer can borrow from the request message (`'req`). Note that + /// it cannot borrow from the response message. + type Producer: ProduceMessage; + + /// A transformer of DNS responses. + /// + /// This type modifies the response from the wrapped service, by adding, + /// removing, or modifying the components of the response message. It is + /// constructed by [`Self::respond_local()`], if an early return does not + /// occur. + /// + /// # Lifetimes + /// + /// The transformer can borrow from the request message (`'req`). Note + /// that it cannot borrow from the response message. + type Transformer: TransformMessage; + + /// Respond to a DNS request. + /// + /// The provided consumer must have been provided the entire DNS request + /// message. If the request should be forwarded through to the wrapped + /// service, [`ControlFlow::Continue`] is returned, with a transformer to + /// modify the wrapped service's response. However, if the request should + /// be responded to directly by this layer, without any interaction from + /// the wrapped service, [`ControlFlow::Break`] is returned. + /// + /// The returned future does not implement [`Send`]. Use [`ServiceLayer`] + /// and [`ServiceLayer::respond()`] for a [`Send`]-implementing version. + #[allow(async_fn_in_trait)] + async fn respond_local( + &self, + request: &RequestMessage<'_>, + ) -> ControlFlow; +} + +//----------- ProduceMessage ------------------------------------------------- + +/// A type that produces a DNS message. +/// +/// This interface is similar to [`Iterator`], except that it can iterate over +/// the different components of a message (questions, answers, authorities, +/// and additional records). +/// +/// # Architecture +/// +/// This interface is convenient when multiple transformers need to modify the +/// message as it is being built. Rather than forcing each transformer to +/// parse and rewrite the message, this interface allows the message to built +/// up over a single iteration, with every transformer directly examining each +/// component added to the message. +/// +/// # Examples +pub trait ProduceMessage { + /// The header of the message. + /// + /// The provided header will be uninitialized, and this method is expected + /// to reset it entirely. The default implementation does nothing. + fn header(&mut self, header: &mut Header) { + let _ = header; + } + + /// The next DNS question in the message. + /// + /// This method is expected to add at most one question using the given + /// message builder. If a question is added, its builder is returned so + /// that it can be modified or filtered before being finalized. + /// + /// This must act like a fused iterator; if no question is added, then + /// future calls to the same method will also add no questions. + /// + /// The default implementation of this method will add no questions. + /// + /// # Errors + /// + /// If new records cannot be inserted in the message (because it is full), + /// the method is responsible for returning gracefully. The message may + /// be marked as truncated, for example. + fn next_question<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + let _ = builder; + None + } + + /// The next answer record in the message. + /// + /// This method is expected to add at most one answer record using the + /// given message builder. If a record is added, its builder is returned + /// so that it can be modified or filtered before being finalized. + /// + /// This must act like a fused iterator; if no record is added, then + /// future calls to the same method will also add no records. + /// + /// The default implementation of this method will add no records. + /// + /// # Errors + /// + /// If new records cannot be inserted in the message (because it is full), + /// the method is responsible for returning gracefully. The message may + /// be marked as truncated, for example. + fn next_answer<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + let _ = builder; + None + } + + /// The next authority record in the message. + /// + /// This method is expected to add at most one authority record using the + /// given message builder. If a record is added, its builder is returned + /// so that it can be modified or filtered before being finalized. + /// + /// This must act like a fused iterator; if no record is added, then + /// future calls to the same method will also add no records. + /// + /// The default implementation of this method will add no records. + /// + /// # Errors + /// + /// If new records cannot be inserted in the message (because it is full), + /// the method is responsible for returning gracefully. The message may + /// be marked as truncated, for example. + fn next_authority<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + let _ = builder; + None + } + + /// The next additional record in the message. + /// + /// This method is expected to add at most one additional record using the + /// given message builder. If a record is added, its builder is returned + /// so that it can be modified or filtered before being finalized. + /// + /// This must act like a fused iterator; if no record is added, then + /// future calls to the same method will also add no records. + /// + /// The default implementation of this method will add no records. + /// + /// # Errors + /// + /// If new records cannot be inserted in the message (because it is full), + /// the method is responsible for returning gracefully. The message may + /// be marked as truncated, for example. + fn next_additional<'b>( + &mut self, + builder: &'b mut MessageBuilder<'_>, + ) -> Option> { + let _ = builder; + None + } +} + +//----------- TransformMessage ----------------------------------------------- + +/// A type that modifies a DNS message as it is being built. +/// +/// This interface is designed around [`ProduceMessage`]: as the components of +/// the message are produced, they are passed through methods of this trait to +/// be modified or filtered out. Furthermore, implementing types can add more +/// components to the message as they also implement [`ProduceMessage`]. +/// +/// # Examples +pub trait TransformMessage: ProduceMessage { + /// Modify the header of the message. + /// + /// The provided header has been initialized; this method can choose to + /// modify it. The default implementation does nothing. + fn modify_header(&mut self, header: &mut Header) { + let _ = header; + } + + /// Modify a question added to the message. + /// + /// This method is called when a question is being added to the message. + /// The question can be modified. + /// + /// If [`ControlFlow::Continue`] is returned, the question is preserved, + /// and can be passed through future transformations. Otherwise, the + /// question is removed. + /// + /// The default implementation of this method passes the question through + /// transparently. + fn modify_question( + &mut self, + builder: &mut QuestionBuilder<'_>, + ) -> ControlFlow<()> { + let _ = builder; + ControlFlow::Continue(()) + } + + /// Modify an answer record added to the message. + /// + /// This method is called when an answer record is being added to the + /// message. The record (and its data) can be modified. + /// + /// If [`ControlFlow::Continue`] is returned, the record is preserved, + /// and can be passed through future transformations. Otherwise, the + /// record is removed. + /// + /// The default implementation of this method passes the record through + /// transparently. + fn modify_answer( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + let _ = builder; + ControlFlow::Continue(()) + } + + /// Modify an authority record added to the message. + /// + /// This method is called when an authority record is being added to the + /// message. The record (and its data) can be modified. + /// + /// If [`ControlFlow::Continue`] is returned, the record is preserved, + /// and can be passed through future transformations. Otherwise, the + /// record is removed. + /// + /// The default implementation of this method passes the record through + /// transparently. + fn modify_authority( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + let _ = builder; + ControlFlow::Continue(()) + } + + /// Modify an additional record added to the message. + /// + /// This method is called when an additional record is being added to the + /// message. The record (and its data) can be modified. + /// + /// If [`ControlFlow::Continue`] is returned, the record is preserved, + /// and can be passed through future transformations. Otherwise, the + /// record is removed. + /// + /// The default implementation of this method passes the record through + /// transparently. + fn modify_additional( + &mut self, + builder: &mut RecordBuilder<'_>, + ) -> ControlFlow<()> { + let _ = builder; + ControlFlow::Continue(()) + } +} diff --git a/src/new_net/server/request.rs b/src/new_net/server/request.rs new file mode 100644 index 000000000..8a0eadc48 --- /dev/null +++ b/src/new_net/server/request.rs @@ -0,0 +1,591 @@ +//! DNS request messages. + +use core::{iter::FusedIterator, marker::PhantomData, ops::Range}; + +use crate::{ + new_base::{ + name::{Name, UnparsedName}, + parse::{ParseMessageBytes, SplitMessageBytes}, + wire::{AsBytes, ParseBytes, ParseError, SizePrefixed, U16}, + Message, ParseRecordData, QClass, QType, Question, RClass, RType, + Record, SectionCounts, UnparsedRecordData, TTL, + }, + new_edns::{EdnsOption, EdnsRecord}, + new_rdata::EdnsOptionsIter, +}; + +/// A DNS request message. +#[derive(Clone)] +pub struct RequestMessage<'b> { + /// The underlying [`Message`]. + pub message: &'b Message, + + /// Cached offsets for the question section. + questions: (Range, [Range; 1]), + + /// Cached offsets for the answer section. + answers: (Range, [Range; 0]), + + /// Cached offsets for the authority section. + authorities: (Range, [Range; 0]), + + /// Cached offsets for the additional section. + additional: (Range, [Range; 2]), + + /// Cached offsets for the EDNS record. + edns: (Range, u16, [Range; 4]), +} + +//--- Construction + +impl<'b> RequestMessage<'b> { + /// Wrap a raw [`Message`] into a [`RequestMessage`]. + /// + /// This will iterate through the message, pre-filling some caches for + /// efficient access in the future. + pub fn new(message: &'b Message) -> Result { + /// Parse the question section into cached offsets. + fn parse_questions( + contents: &[u8], + range: &mut Range, + number: u16, + indices: &mut [Range], + ) -> Result<(), ParseError> { + let mut indices = indices.iter_mut(); + let mut offset = range.start as usize; + + for _ in 0..number { + let (question, rest) = + Question::<&UnparsedName>::split_message_bytes( + contents, offset, + )?; + + if let Some(indices) = indices.next() { + let fields = offset + question.qname.len(); + *indices = offset as u16..fields as u16; + } + + offset = rest; + } + + range.end = offset as u16; + Ok(()) + } + + /// Parse a record section into cached offsets. + fn parse_records( + contents: &[u8], + section: u8, + range: &mut Range, + number: u16, + indices: &mut [Range], + edns_range: &mut Option>, + ) -> Result<(), ParseError> { + let mut indices = indices.iter_mut(); + let mut offset = range.start as usize; + + for _ in 0..number { + let (record, rest) = Record::< + &UnparsedName, + &UnparsedRecordData, + >::split_message_bytes( + contents, offset + )?; + + if let Some(indices) = indices.next() { + let fields = offset + record.rname.len(); + *indices = offset as u16..fields as u16; + } + + if section == 3 && record.rtype == RType::OPT { + if edns_range.is_some() { + // A DNS message can only contain one EDNS record. + return Err(ParseError); + } + + *edns_range = Some(offset as u16..rest as u16); + } + + offset = rest; + } + + range.end = offset as u16; + Ok(()) + } + + /// Parse the EDNS record into cached offsets. + fn parse_edns( + contents: &[u8], + range: Range, + number: &mut u16, + indices: &mut [Range], + ) -> Result<(), ParseError> { + let mut indices = indices.iter_mut(); + let mut offset = range.start as usize + 11; + + while offset < range.end as usize { + let (_type, rest) = + <&U16>::split_message_bytes(contents, offset)?; + let (_data, rest) = + >::split_message_bytes( + contents, rest, + )?; + + *number += 1; + + if let Some(indices) = indices.next() { + *indices = offset as u16..rest as u16; + } + + offset = rest; + } + + Ok(()) + } + + // DNS messages are 64KiB at the largest. + let _ = u16::try_from(message.as_bytes().len()) + .map_err(|_| ParseError)?; + + let mut this = Self { + message, + questions: Default::default(), + answers: Default::default(), + authorities: Default::default(), + additional: Default::default(), + edns: Default::default(), + }; + + let mut edns_range = None; + + parse_questions( + &message.contents, + &mut this.questions.0, + message.header.counts.questions.get(), + &mut this.questions.1, + )?; + + this.answers.0 = this.questions.0.end..0; + parse_records( + &message.contents, + 1, + &mut this.answers.0, + message.header.counts.answers.get(), + &mut this.answers.1, + &mut edns_range, + )?; + + this.authorities.0 = this.answers.0.end..0; + parse_records( + &message.contents, + 2, + &mut this.authorities.0, + message.header.counts.authorities.get(), + &mut this.authorities.1, + &mut edns_range, + )?; + + this.additional.0 = this.authorities.0.end..0; + parse_records( + &message.contents, + 2, + &mut this.additional.0, + message.header.counts.additional.get(), + &mut this.additional.1, + &mut edns_range, + )?; + + if let Some(edns_range) = edns_range { + this.edns.0 = edns_range.clone(); + parse_edns( + &message.contents, + edns_range, + &mut this.edns.1, + &mut this.edns.2, + )?; + } + + Ok(this) + } +} + +//--- Internals + +impl<'b> RequestMessage<'b> { + /// The section counts. + fn counts(&self) -> &'b SectionCounts { + &self.message.header.counts + } +} + +//--- Inspection + +impl<'b> RequestMessage<'b> { + /// The sole question in the message. + /// + /// # Name Compression + /// + /// Due to the restrictions around compressed domain names (in order to + /// prevent attackers from crafting compression pointer loops), it is + /// guaranteed that the first QNAME in the message is uncompressed. + /// + /// # Errors + /// + /// Fails if there are zero or more than one question in the message. + pub fn sole_question(&self) -> Result, ParseError> { + if self.message.header.counts.questions.get() != 1 { + return Err(ParseError); + } + + // SAFETY: 'RequestMessage' is pre-validated. + let range = self.questions.1[0].clone(); + let range = range.start as usize..range.end as usize; + let qname = &self.message.contents[range.clone()]; + let qname = unsafe { Name::from_bytes_unchecked(qname) }; + let fields = &self.message.contents[range.end..]; + let qtype = QType::parse_bytes(&fields[0..2]).unwrap(); + let qclass = QClass::parse_bytes(&fields[2..4]).unwrap(); + + Ok(Question { + qname, + qtype, + qclass, + }) + } + + /// The EDNS record in the message, if any. + pub fn edns_record(&self) -> Option> { + if self.edns.0.is_empty() { + return None; + } + + let range = self.edns.0.clone(); + let contents = &self.message.contents[..range.end as usize]; + EdnsRecord::parse_message_bytes(contents, range.start as usize) + .map(Some) + .expect("'RequestMessage' only holds well-formed EDNS records") + } + + /// The questions in the message. + /// + /// An iterator of [`Question`]s is returned. + /// + /// # Name Compression + /// + /// The returned questions use [`UnparsedName`] for the QNAMEs. These can + /// be resolved against the original message to determine the whole domain + /// name, if necessary. Note that decompression can fail. + pub fn questions(&self) -> RequestQuestions<'_, 'b> { + let contents = self.questions.0.clone(); + RequestQuestions { + message: self, + cache: self.questions.1.iter(), + contents: contents.start as usize..contents.end as usize, + indices: 0..self.counts().questions.get(), + } + } + + /// The answer records in the message. + /// + /// An iterator of [`Record`]s is returned. + /// + /// # Name Compression + /// + /// The returned records use [`UnparsedName`] for the RNAMEs. These can + /// be resolved against the original message to determine the whole domain + /// name, if necessary. Note that decompression can fail. + /// + /// # Record Data + /// + /// The caller can select an appropriate record data type to use. In most + /// cases, [`RecordData`](crate::new_rdata::RecordData) is appropriate; if + /// many records will be skipped, however, [`UnparsedRecordData`] might be + /// preferable. + pub fn answers(&self) -> RequestRecords<'_, 'b, D> + where + D: ParseRecordData<'b>, + { + let contents = self.answers.0.clone(); + RequestRecords { + message: self, + cache: self.answers.1.iter(), + contents: contents.start as usize..contents.end as usize, + indices: 0..self.counts().answers.get(), + _rdata: PhantomData, + } + } + + /// The authority records in the message. + /// + /// An iterator of [`Record`]s is returned. + /// + /// # Name Compression + /// + /// The returned records use [`UnparsedName`] for the RNAMEs. These can + /// be resolved against the original message to determine the whole domain + /// name, if necessary. Note that decompression can fail. + /// + /// # Record Data + /// + /// The caller can select an appropriate record data type to use. In most + /// cases, [`RecordData`](crate::new_rdata::RecordData) is appropriate; if + /// many records will be skipped, however, [`UnparsedRecordData`] might be + /// preferable. + pub fn authorities(&self) -> RequestRecords<'_, 'b, D> + where + D: ParseRecordData<'b>, + { + let contents = self.authorities.0.clone(); + RequestRecords { + message: self, + cache: self.authorities.1.iter(), + contents: contents.start as usize..contents.end as usize, + indices: 0..self.counts().authorities.get(), + _rdata: PhantomData, + } + } + + /// The additional records in the message. + /// + /// An iterator of [`Record`]s is returned. + /// + /// # Name Compression + /// + /// The returned records use [`UnparsedName`] for the RNAMEs. These can + /// be resolved against the original message to determine the whole domain + /// name, if necessary. Note that decompression can fail. + /// + /// # Record Data + /// + /// The caller can select an appropriate record data type to use. In most + /// cases, [`RecordData`](crate::new_rdata::RecordData) is appropriate; if + /// many records will be skipped, however, [`UnparsedRecordData`] might be + /// preferable. + pub fn additional(&self) -> RequestRecords<'_, 'b, D> + where + D: ParseRecordData<'b>, + { + let contents = self.additional.0.clone(); + RequestRecords { + message: self, + cache: self.additional.1.iter(), + contents: contents.start as usize..contents.end as usize, + indices: 0..self.counts().additional.get(), + _rdata: PhantomData, + } + } + + /// The EDNS options in the message. + /// + /// An iterator of (results of) [`EdnsOption`]s is returned. + /// + /// If the message does not contain an EDNS record, this is empty. + /// + /// # Errors + /// + /// While the overall structure of the EDNS record is verified when the + /// [`RequestMessage`] is constructed, option-specific validation (e.g. + /// checking the size of an EDNS cookie option) is not performed early. + /// Those errors will be caught and returned by the iterator. + pub fn edns_options(&self) -> RequestEdnsOptions<'b> { + let start = self.edns.0.start as usize + 11; + let end = self.edns.0.end as usize; + let options = &self.message.contents[start..end]; + RequestEdnsOptions { + inner: EdnsOptionsIter::new(options), + indices: 0..self.edns.1, + } + } +} + +//----------- RequestQuestions ----------------------------------------------- + +/// The questions in a [`RequestMessage`]. +#[derive(Clone)] +pub struct RequestQuestions<'r, 'b> { + /// The underlying request message. + message: &'r RequestMessage<'b>, + + /// The cached question ranges. + cache: core::slice::Iter<'r, Range>, + + /// The range of message contents to parse. + contents: Range, + + /// The range of record indices left. + indices: Range, +} + +impl<'b> Iterator for RequestQuestions<'_, 'b> { + type Item = Question<&'b UnparsedName>; + + fn next(&mut self) -> Option { + // Try loading a cached question. + if let Some(range) = self.cache.next().cloned() { + if range.is_empty() { + // There are no more questions, stop. + self.cache = Default::default(); + self.contents.start = self.contents.end; + return None; + } + + // SAFETY: 'RequestMessage' is pre-validated. + let range = range.start as usize..range.end as usize; + let qname = &self.message.message.contents[range.clone()]; + let qname = unsafe { UnparsedName::from_bytes_unchecked(qname) }; + let fields = &self.message.message.contents[range.end..]; + let qtype = QType::parse_bytes(&fields[0..2]).unwrap(); + let qclass = QClass::parse_bytes(&fields[2..4]).unwrap(); + + self.indices.start += 1; + return Some(Question { + qname, + qtype, + qclass, + }); + } + + let _ = self.indices.next()?; + let contents = &self.message.message.contents[..self.contents.end]; + let (question, rest) = + Question::split_message_bytes(contents, self.contents.start) + .expect("'RequestMessage' only contains valid questions"); + + self.contents.start = rest; + Some(question) + } +} + +impl ExactSizeIterator for RequestQuestions<'_, '_> { + fn len(&self) -> usize { + self.indices.len() + } +} + +impl FusedIterator for RequestQuestions<'_, '_> {} + +//----------- RequestRecords ------------------------------------------------- + +/// The records in a section of a [`RequestMessage`]. +#[derive(Clone)] +pub struct RequestRecords<'r, 'b, D> { + /// The underlying request message. + message: &'r RequestMessage<'b>, + + /// The cached record ranges. + cache: core::slice::Iter<'r, Range>, + + /// The range of message contents to parse. + contents: Range, + + /// The range of record indices left. + indices: Range, + + /// A representation of the record data held. + _rdata: PhantomData<&'r [D]>, +} + +impl<'b, D> Iterator for RequestRecords<'_, 'b, D> +where + D: ParseRecordData<'b>, +{ + type Item = Result, ParseError>; + + fn next(&mut self) -> Option { + // Try loading a cached record. + if let Some(range) = self.cache.next().cloned() { + if range.is_empty() { + // There are no more records, stop. + self.cache = Default::default(); + self.contents.start = self.contents.end; + return None; + } + + // SAFETY: 'RequestMessage' is pre-validated. + let range = range.start as usize..range.end as usize; + let rname = &self.message.message.contents[range.clone()]; + let rname = unsafe { UnparsedName::from_bytes_unchecked(rname) }; + let fields = &self.message.message.contents[range.end..]; + let rtype = RType::parse_bytes(&fields[0..2]).unwrap(); + let rclass = RClass::parse_bytes(&fields[2..4]).unwrap(); + let ttl = TTL::parse_bytes(&fields[4..8]).unwrap(); + let size = U16::parse_bytes(&fields[8..10]).unwrap(); + let rdata_end = range.end + 10 + size.get() as usize; + let rdata = &self.message.message.contents[..rdata_end]; + let rdata = + match D::parse_record_data(rdata, range.end + 10, rtype) { + Ok(rdata) => rdata, + Err(err) => return Some(Err(err)), + }; + + self.indices.start += 1; + return Some(Ok(Record { + rname, + rtype, + rclass, + ttl, + rdata, + })); + } + + let _ = self.indices.next()?; + let contents = &self.message.message.contents[..self.contents.end]; + let (record, rest) = match Record::split_message_bytes( + contents, + self.contents.start, + ) { + Ok((record, rest)) => (record, rest), + Err(err) => return Some(Err(err)), + }; + + self.contents.start = rest; + Some(Ok(record)) + } +} + +impl<'b, D> ExactSizeIterator for RequestRecords<'_, 'b, D> +where + D: ParseRecordData<'b>, +{ + fn len(&self) -> usize { + self.indices.len() + } +} + +impl<'b, D> FusedIterator for RequestRecords<'_, 'b, D> where + D: ParseRecordData<'b> +{ +} + +//----------- RequestEdnsOptions --------------------------------------------- + +/// The EDNS options in a [`RequestMessage`]. +/// +/// This is a wrapper around [`EdnsOptionsIter`] that also implements +/// [`ExactSizeIterator`], as the total number of EDNS options is cached in +/// the [`RequestMessage`]. +#[derive(Clone)] +pub struct RequestEdnsOptions<'b> { + /// The underlying iterator. + inner: EdnsOptionsIter<'b>, + + /// The range of option indices left. + indices: Range, +} + +impl<'b> Iterator for RequestEdnsOptions<'b> { + type Item = Result, ParseError>; + + fn next(&mut self) -> Option { + let _ = self.indices.next()?; + self.inner.next() + } +} + +impl ExactSizeIterator for RequestEdnsOptions<'_> { + fn len(&self) -> usize { + self.indices.len() + } +} + +impl FusedIterator for RequestEdnsOptions<'_> {}