use super::header::{Header, HeaderCounts, HeaderSection};
#[cfg(feature = "random")]
use super::iana::Rtype;
use super::iana::{OptRcode, OptionCode, Rcode};
use super::message::Message;
use super::name::{Label, ToDname};
#[cfg(feature = "std")]
use super::octets::Octets64;
use super::octets::{Compose, OctetsBuilder, OctetsRef, ShortBuf};
use super::opt::{OptData, OptHeader};
use super::question::AsQuestion;
use super::record::AsRecord;
#[cfg(feature = "bytes")]
use bytes::BytesMut;
#[cfg(feature = "std")]
use core::convert::TryInto;
use core::mem;
use core::ops::{Deref, DerefMut};
#[cfg(feature = "std")]
use std::collections::HashMap;
#[cfg(feature = "std")]
use std::vec::Vec;
#[derive(Clone, Debug)]
pub struct MessageBuilder<Target> {
target: Target,
}
impl<Target: OctetsBuilder> MessageBuilder<Target> {
pub fn from_target(mut target: Target) -> Result<Self, ShortBuf> {
target.truncate(0);
target.append_slice(HeaderSection::new().as_slice())?;
Ok(MessageBuilder { target })
}
}
#[cfg(feature = "std")]
impl MessageBuilder<Vec<u8>> {
pub fn new_vec() -> Self {
Self::from_target(Vec::new()).unwrap()
}
}
#[cfg(feature = "std")]
impl MessageBuilder<StreamTarget<Vec<u8>>> {
pub fn new_stream_vec() -> Self {
Self::from_target(StreamTarget::new(Vec::new()).unwrap()).unwrap()
}
}
#[cfg(feature = "bytes")]
impl MessageBuilder<BytesMut> {
pub fn new_bytes() -> Self {
Self::from_target(BytesMut::new()).unwrap()
}
}
#[cfg(feature = "bytes")]
impl MessageBuilder<StreamTarget<BytesMut>> {
pub fn new_stream_bytes() -> Self {
Self::from_target(StreamTarget::new(BytesMut::new()).unwrap())
.unwrap()
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> MessageBuilder<Target> {
pub fn start_answer<Octets>(
mut self,
msg: &Message<Octets>,
rcode: Rcode,
) -> Result<AnswerBuilder<Target>, ShortBuf>
where
Octets: AsRef<[u8]>,
for<'a> &'a Octets: OctetsRef,
{
{
let header = self.header_mut();
header.set_id(msg.header().id());
header.set_qr(true);
header.set_opcode(msg.header().opcode());
header.set_rd(msg.header().rd());
header.set_rcode(rcode);
}
let mut builder = self.question();
for item in msg.question().flatten() {
builder.push(item)?;
}
Ok(builder.answer())
}
#[cfg(feature = "random")]
pub fn request_axfr<N: ToDname>(
mut self,
apex: N,
) -> Result<AnswerBuilder<Target>, ShortBuf> {
self.header_mut().set_random_id();
let mut builder = self.question();
builder.push((apex, Rtype::Axfr))?;
Ok(builder.answer())
}
}
impl<Target: OctetsBuilder + AsRef<[u8]>> MessageBuilder<Target> {
pub fn header(&self) -> Header {
*Header::for_message_slice(self.target.as_ref())
}
pub fn counts(&self) -> HeaderCounts {
*HeaderCounts::for_message_slice(self.target.as_ref())
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> MessageBuilder<Target> {
pub fn header_mut(&mut self) -> &mut Header {
Header::for_message_slice_mut(self.target.as_mut())
}
fn counts_mut(&mut self) -> &mut HeaderCounts {
HeaderCounts::for_message_slice_mut(self.target.as_mut())
}
}
impl<Target: OctetsBuilder> MessageBuilder<Target> {
pub fn builder(self) -> MessageBuilder<Target> {
self
}
pub fn question(self) -> QuestionBuilder<Target> {
QuestionBuilder::new(self)
}
pub fn answer(self) -> AnswerBuilder<Target> {
self.question().answer()
}
pub fn authority(self) -> AuthorityBuilder<Target> {
self.question().answer().authority()
}
pub fn additional(self) -> AdditionalBuilder<Target> {
self.question().answer().authority().additional()
}
pub fn finish(self) -> Target {
self.target
}
pub fn into_message(self) -> Message<Target::Octets> {
unsafe { Message::from_octets_unchecked(self.target.freeze()) }
}
}
impl<Target> MessageBuilder<Target> {
pub fn as_target(&self) -> &Target {
&self.target
}
fn as_target_mut(&mut self) -> &mut Target {
&mut self.target
}
pub fn as_slice(&self) -> &[u8]
where
Target: AsRef<[u8]>,
{
self.as_target().as_ref()
}
pub fn as_message(&self) -> Message<&[u8]>
where
Target: AsRef<[u8]>,
{
unsafe { Message::from_octets_unchecked(self.target.as_ref()) }
}
}
impl<Target> From<QuestionBuilder<Target>> for MessageBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: QuestionBuilder<Target>) -> Self {
src.builder()
}
}
impl<Target> From<AnswerBuilder<Target>> for MessageBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AnswerBuilder<Target>) -> Self {
src.builder()
}
}
impl<Target> From<AuthorityBuilder<Target>> for MessageBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AuthorityBuilder<Target>) -> Self {
src.builder()
}
}
impl<Target> From<AdditionalBuilder<Target>> for MessageBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AdditionalBuilder<Target>) -> Self {
src.builder()
}
}
impl<Target> AsRef<Target> for MessageBuilder<Target> {
fn as_ref(&self) -> &Target {
self.as_target()
}
}
impl<Target: AsRef<[u8]>> AsRef<[u8]> for MessageBuilder<Target> {
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
#[derive(Clone, Debug)]
pub struct QuestionBuilder<Target> {
builder: MessageBuilder<Target>,
}
impl<Target: OctetsBuilder> QuestionBuilder<Target> {
fn new(builder: MessageBuilder<Target>) -> Self {
Self { builder }
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> QuestionBuilder<Target> {
pub fn push(
&mut self,
question: impl AsQuestion,
) -> Result<(), ShortBuf> {
let pos = self.as_target().len();
question.compose_question(self.as_target_mut())?;
self.counts_mut().inc_qdcount().map_err(|err| {
self.as_target_mut().truncate(pos);
err
})
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> QuestionBuilder<Target> {
pub fn rewind(&mut self) {
self.as_target_mut()
.truncate(mem::size_of::<HeaderSection>());
self.counts_mut().set_qdcount(0);
}
pub fn builder(mut self) -> MessageBuilder<Target> {
self.rewind();
self.builder
}
}
impl<Target: OctetsBuilder> QuestionBuilder<Target> {
pub fn question(self) -> QuestionBuilder<Target> {
self
}
pub fn answer(self) -> AnswerBuilder<Target> {
AnswerBuilder::new(self.builder)
}
pub fn authority(self) -> AuthorityBuilder<Target> {
self.answer().authority()
}
pub fn additional(self) -> AdditionalBuilder<Target> {
self.answer().authority().additional()
}
pub fn finish(self) -> Target {
self.builder.finish()
}
pub fn into_message(self) -> Message<Target::Octets> {
self.builder.into_message()
}
}
impl<Target> QuestionBuilder<Target> {
pub fn as_builder(&self) -> &MessageBuilder<Target> {
&self.builder
}
pub fn as_builder_mut(&mut self) -> &mut MessageBuilder<Target> {
&mut self.builder
}
}
impl<Target> From<MessageBuilder<Target>> for QuestionBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: MessageBuilder<Target>) -> Self {
src.question()
}
}
impl<Target> From<AnswerBuilder<Target>> for QuestionBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AnswerBuilder<Target>) -> Self {
src.question()
}
}
impl<Target> From<AuthorityBuilder<Target>> for QuestionBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AuthorityBuilder<Target>) -> Self {
src.question()
}
}
impl<Target> From<AdditionalBuilder<Target>> for QuestionBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AdditionalBuilder<Target>) -> Self {
src.question()
}
}
impl<Target> Deref for QuestionBuilder<Target> {
type Target = MessageBuilder<Target>;
fn deref(&self) -> &Self::Target {
&self.builder
}
}
impl<Target> DerefMut for QuestionBuilder<Target> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.builder
}
}
impl<Target> AsRef<MessageBuilder<Target>> for QuestionBuilder<Target> {
fn as_ref(&self) -> &MessageBuilder<Target> {
self.as_builder()
}
}
impl<Target> AsMut<MessageBuilder<Target>> for QuestionBuilder<Target> {
fn as_mut(&mut self) -> &mut MessageBuilder<Target> {
self.as_builder_mut()
}
}
impl<Target> AsRef<Target> for QuestionBuilder<Target> {
fn as_ref(&self) -> &Target {
self.as_target()
}
}
impl<Target: AsRef<[u8]>> AsRef<[u8]> for QuestionBuilder<Target> {
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
#[derive(Clone, Debug)]
pub struct AnswerBuilder<Target> {
builder: MessageBuilder<Target>,
start: usize,
}
impl<Target: OctetsBuilder> AnswerBuilder<Target> {
fn new(builder: MessageBuilder<Target>) -> Self {
AnswerBuilder {
start: builder.target.len(),
builder,
}
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> AnswerBuilder<Target> {
pub fn push(&mut self, record: impl AsRecord) -> Result<(), ShortBuf> {
let pos = self.as_target().len();
record.compose_record(self.as_target_mut())?;
self.counts_mut().inc_ancount().map_err(|err| {
self.as_target_mut().truncate(pos);
err
})
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> AnswerBuilder<Target> {
pub fn rewind(&mut self) {
self.builder.target.truncate(self.start);
self.counts_mut().set_ancount(0);
}
pub fn builder(self) -> MessageBuilder<Target> {
self.question().builder()
}
pub fn question(mut self) -> QuestionBuilder<Target> {
self.rewind();
QuestionBuilder::new(self.builder)
}
}
impl<Target: OctetsBuilder> AnswerBuilder<Target> {
pub fn answer(self) -> AnswerBuilder<Target> {
self
}
pub fn authority(self) -> AuthorityBuilder<Target> {
AuthorityBuilder::new(self)
}
pub fn additional(self) -> AdditionalBuilder<Target> {
self.authority().additional()
}
pub fn finish(self) -> Target {
self.builder.finish()
}
pub fn into_message(self) -> Message<Target::Octets> {
self.builder.into_message()
}
}
impl<Target> AnswerBuilder<Target> {
pub fn as_builder(&self) -> &MessageBuilder<Target> {
&self.builder
}
pub fn as_builder_mut(&mut self) -> &mut MessageBuilder<Target> {
&mut self.builder
}
}
impl<Target> From<MessageBuilder<Target>> for AnswerBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: MessageBuilder<Target>) -> Self {
src.answer()
}
}
impl<Target> From<QuestionBuilder<Target>> for AnswerBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: QuestionBuilder<Target>) -> Self {
src.answer()
}
}
impl<Target> From<AuthorityBuilder<Target>> for AnswerBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AuthorityBuilder<Target>) -> Self {
src.answer()
}
}
impl<Target> From<AdditionalBuilder<Target>> for AnswerBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AdditionalBuilder<Target>) -> Self {
src.answer()
}
}
impl<Target> Deref for AnswerBuilder<Target> {
type Target = MessageBuilder<Target>;
fn deref(&self) -> &Self::Target {
&self.builder
}
}
impl<Target> DerefMut for AnswerBuilder<Target> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.builder
}
}
impl<Target> AsRef<MessageBuilder<Target>> for AnswerBuilder<Target> {
fn as_ref(&self) -> &MessageBuilder<Target> {
self.as_builder()
}
}
impl<Target> AsMut<MessageBuilder<Target>> for AnswerBuilder<Target> {
fn as_mut(&mut self) -> &mut MessageBuilder<Target> {
self.as_builder_mut()
}
}
impl<Target> AsRef<Target> for AnswerBuilder<Target> {
fn as_ref(&self) -> &Target {
self.as_target()
}
}
impl<Target: AsRef<[u8]>> AsRef<[u8]> for AnswerBuilder<Target> {
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
#[derive(Clone, Debug)]
pub struct AuthorityBuilder<Target> {
answer: AnswerBuilder<Target>,
start: usize,
}
impl<Target: OctetsBuilder> AuthorityBuilder<Target> {
fn new(answer: AnswerBuilder<Target>) -> Self {
AuthorityBuilder {
start: answer.as_target().len(),
answer,
}
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> AuthorityBuilder<Target> {
pub fn push(&mut self, record: impl AsRecord) -> Result<(), ShortBuf> {
let pos = self.as_target().len();
record.compose_record(self.as_target_mut())?;
self.counts_mut().inc_nscount().map_err(|err| {
self.as_target_mut().truncate(pos);
err
})
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> AuthorityBuilder<Target> {
pub fn rewind(&mut self) {
self.answer.as_target_mut().truncate(self.start);
self.counts_mut().set_nscount(0);
}
pub fn builder(self) -> MessageBuilder<Target> {
self.question().builder()
}
pub fn question(self) -> QuestionBuilder<Target> {
self.answer().question()
}
pub fn answer(mut self) -> AnswerBuilder<Target> {
self.rewind();
self.answer
}
}
impl<Target: OctetsBuilder> AuthorityBuilder<Target> {
pub fn authority(self) -> AuthorityBuilder<Target> {
self
}
pub fn additional(self) -> AdditionalBuilder<Target> {
AdditionalBuilder::new(self)
}
pub fn finish(self) -> Target {
self.answer.finish()
}
pub fn into_message(self) -> Message<Target::Octets> {
self.answer.into_message()
}
}
impl<Target> AuthorityBuilder<Target> {
pub fn as_builder(&self) -> &MessageBuilder<Target> {
self.answer.as_builder()
}
pub fn as_builder_mut(&mut self) -> &mut MessageBuilder<Target> {
self.answer.as_builder_mut()
}
}
impl<Target> From<MessageBuilder<Target>> for AuthorityBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: MessageBuilder<Target>) -> Self {
src.authority()
}
}
impl<Target> From<QuestionBuilder<Target>> for AuthorityBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: QuestionBuilder<Target>) -> Self {
src.authority()
}
}
impl<Target> From<AnswerBuilder<Target>> for AuthorityBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: AnswerBuilder<Target>) -> Self {
src.authority()
}
}
impl<Target> From<AdditionalBuilder<Target>> for AuthorityBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn from(src: AdditionalBuilder<Target>) -> Self {
src.authority()
}
}
impl<Target> Deref for AuthorityBuilder<Target> {
type Target = MessageBuilder<Target>;
fn deref(&self) -> &Self::Target {
self.answer.deref()
}
}
impl<Target> DerefMut for AuthorityBuilder<Target> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.answer.deref_mut()
}
}
impl<Target> AsRef<MessageBuilder<Target>> for AuthorityBuilder<Target> {
fn as_ref(&self) -> &MessageBuilder<Target> {
self.as_builder()
}
}
impl<Target> AsMut<MessageBuilder<Target>> for AuthorityBuilder<Target> {
fn as_mut(&mut self) -> &mut MessageBuilder<Target> {
self.as_builder_mut()
}
}
impl<Target> AsRef<Target> for AuthorityBuilder<Target> {
fn as_ref(&self) -> &Target {
self.as_target()
}
}
impl<Target: AsRef<[u8]>> AsRef<[u8]> for AuthorityBuilder<Target> {
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
#[derive(Clone, Debug)]
pub struct AdditionalBuilder<Target> {
authority: AuthorityBuilder<Target>,
start: usize,
}
impl<Target: OctetsBuilder> AdditionalBuilder<Target> {
fn new(authority: AuthorityBuilder<Target>) -> Self {
AdditionalBuilder {
start: authority.as_target().len(),
authority,
}
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> AdditionalBuilder<Target> {
pub fn push(&mut self, record: impl AsRecord) -> Result<(), ShortBuf> {
let pos = self.as_target().len();
record.compose_record(self.as_target_mut())?;
self.counts_mut().inc_arcount().map_err(|err| {
self.as_target_mut().truncate(pos);
err
})
}
}
impl<Target> AdditionalBuilder<Target>
where
Target: OctetsBuilder + AsRef<[u8]> + AsMut<[u8]>,
{
pub fn opt<F, R>(&mut self, build: F) -> Result<R, ShortBuf>
where
F: FnOnce(&mut OptBuilder<Target>) -> Result<R, ShortBuf>,
{
build(&mut OptBuilder::new(self)?)
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> AdditionalBuilder<Target> {
pub fn rewind(&mut self) {
self.authority.as_target_mut().truncate(self.start);
self.counts_mut().set_arcount(0);
}
pub fn builder(self) -> MessageBuilder<Target> {
self.question().builder()
}
pub fn question(self) -> QuestionBuilder<Target> {
self.answer().question()
}
pub fn answer(self) -> AnswerBuilder<Target> {
self.authority().answer()
}
pub fn authority(mut self) -> AuthorityBuilder<Target> {
self.rewind();
self.authority
}
}
impl<Target: OctetsBuilder> AdditionalBuilder<Target> {
pub fn additional(self) -> AdditionalBuilder<Target> {
self
}
pub fn finish(self) -> Target {
self.authority.finish()
}
pub fn into_message(self) -> Message<Target::Octets> {
self.authority.into_message()
}
}
impl<Target> AdditionalBuilder<Target> {
pub fn as_builder(&self) -> &MessageBuilder<Target> {
self.authority.as_builder()
}
pub fn as_builder_mut(&mut self) -> &mut MessageBuilder<Target> {
self.authority.as_builder_mut()
}
}
impl<Target> From<MessageBuilder<Target>> for AdditionalBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: MessageBuilder<Target>) -> Self {
src.additional()
}
}
impl<Target> From<QuestionBuilder<Target>> for AdditionalBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: QuestionBuilder<Target>) -> Self {
src.additional()
}
}
impl<Target> From<AnswerBuilder<Target>> for AdditionalBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: AnswerBuilder<Target>) -> Self {
src.additional()
}
}
impl<Target> From<AuthorityBuilder<Target>> for AdditionalBuilder<Target>
where
Target: OctetsBuilder,
{
fn from(src: AuthorityBuilder<Target>) -> Self {
src.additional()
}
}
impl<Target> Deref for AdditionalBuilder<Target> {
type Target = MessageBuilder<Target>;
fn deref(&self) -> &Self::Target {
self.as_builder()
}
}
impl<Target> DerefMut for AdditionalBuilder<Target> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.as_builder_mut()
}
}
impl<Target> AsRef<MessageBuilder<Target>> for AdditionalBuilder<Target> {
fn as_ref(&self) -> &MessageBuilder<Target> {
self.as_builder()
}
}
impl<Target> AsMut<MessageBuilder<Target>> for AdditionalBuilder<Target> {
fn as_mut(&mut self) -> &mut MessageBuilder<Target> {
self.as_builder_mut()
}
}
impl<Target> AsRef<Target> for AdditionalBuilder<Target> {
fn as_ref(&self) -> &Target {
self.as_target()
}
}
impl<Target: AsRef<[u8]>> AsRef<[u8]> for AdditionalBuilder<Target> {
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
pub trait RecordSectionBuilder {
fn push(&mut self, record: impl AsRecord) -> Result<(), ShortBuf>;
}
impl<Target> RecordSectionBuilder for AnswerBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn push(&mut self, record: impl AsRecord) -> Result<(), ShortBuf> {
Self::push(self, record)
}
}
impl<Target: OctetsBuilder + AsMut<[u8]>> RecordSectionBuilder
for AuthorityBuilder<Target>
{
fn push(&mut self, record: impl AsRecord) -> Result<(), ShortBuf> {
Self::push(self, record)
}
}
impl<Target> RecordSectionBuilder for AdditionalBuilder<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
fn push(&mut self, record: impl AsRecord) -> Result<(), ShortBuf> {
Self::push(self, record)
}
}
#[derive(Debug)]
pub struct OptBuilder<'a, Target> {
additional: &'a mut AdditionalBuilder<Target>,
start: usize,
#[allow(dead_code)] arcount: u16,
}
impl<'a, Target> OptBuilder<'a, Target>
where
Target: OctetsBuilder + AsRef<[u8]> + AsMut<[u8]>,
{
fn new(
additional: &'a mut AdditionalBuilder<Target>,
) -> Result<Self, ShortBuf> {
let start = additional.as_target().len();
let arcount = additional.counts().arcount();
let err = additional
.as_target_mut()
.append_all(|target| {
OptHeader::default().compose(target)?;
0u16.compose(target)
})
.is_err();
if err {
return Err(ShortBuf);
}
if additional.counts_mut().inc_arcount().is_err() {
additional.as_target_mut().truncate(start);
return Err(ShortBuf);
}
Ok(OptBuilder {
additional,
start,
arcount,
})
}
pub fn push<Opt: OptData>(&mut self, opt: &Opt) -> Result<(), ShortBuf> {
self.push_raw_option(opt.code(), |target| opt.compose(target))
}
pub fn push_raw_option<F>(
&mut self,
code: OptionCode,
op: F,
) -> Result<(), ShortBuf>
where
F: FnOnce(&mut Target) -> Result<(), ShortBuf>,
{
let pos = self.as_target().as_ref().len();
self.as_target_mut().append_all(|target| {
code.compose(target)?;
target.u16_len_prefixed(|target| op(target))
})?;
let len = self.as_target().as_ref().len()
- self.start
- (mem::size_of::<OptHeader>() + 2);
if len > usize::from(u16::max_value()) {
self.as_target_mut().truncate(pos);
return Err(ShortBuf);
}
let start = self.start + mem::size_of::<OptHeader>();
self.as_target_mut().as_mut()[start..start + 2]
.copy_from_slice(&(len as u16).to_be_bytes());
Ok(())
}
pub fn udp_payload_size(&self) -> u16 {
self.opt_header().udp_payload_size()
}
pub fn set_udp_payload_size(&mut self, value: u16) {
self.opt_header_mut().set_udp_payload_size(value)
}
pub fn rcode(&self) -> OptRcode {
self.opt_header().rcode(self.additional.header())
}
pub fn set_rcode(&mut self, rcode: OptRcode) {
self.additional.header_mut().set_rcode(rcode.rcode());
self.opt_header_mut().set_rcode(rcode)
}
pub fn version(&self) -> u8 {
self.opt_header().version()
}
pub fn set_version(&mut self, version: u8) {
self.opt_header_mut().set_version(version)
}
pub fn dnssec_ok(&self) -> bool {
self.opt_header().dnssec_ok()
}
pub fn set_dnssec_ok(&mut self, value: bool) {
self.opt_header_mut().set_dnssec_ok(value)
}
fn opt_header(&self) -> &OptHeader {
OptHeader::for_record_slice(&self.as_target().as_ref()[self.start..])
}
fn opt_header_mut(&mut self) -> &mut OptHeader {
let start = self.start;
OptHeader::for_record_slice_mut(
&mut self.as_target_mut().as_mut()[start..],
)
}
pub fn as_target(&self) -> &Target {
self.additional.as_target()
}
fn as_target_mut(&mut self) -> &mut Target {
self.additional.as_target_mut()
}
}
#[derive(Clone, Debug)]
pub struct StreamTarget<Target> {
target: Target,
}
impl<Target: OctetsBuilder + AsMut<[u8]>> StreamTarget<Target> {
pub fn new(mut target: Target) -> Result<Self, ShortBuf> {
target.truncate(0);
0u16.compose(&mut target)?;
Ok(StreamTarget { target })
}
}
#[cfg(feature = "std")]
impl StreamTarget<Vec<u8>> {
pub fn new_vec() -> Self {
Self::new(Vec::new()).unwrap()
}
}
impl<Target: OctetsBuilder> StreamTarget<Target> {
pub fn as_target(&self) -> &Target {
&self.target
}
pub fn into_target(self) -> Target {
self.target
}
fn update_shim(&mut self)
where
Target: AsMut<[u8]>,
{
let len = (self.target.len() - 2) as u16;
self.target.as_mut()[..2].copy_from_slice(&len.to_be_bytes())
}
}
impl<Target: OctetsBuilder + AsRef<[u8]>> StreamTarget<Target> {
pub fn as_stream_slice(&self) -> &[u8] {
self.target.as_ref()
}
pub fn as_dgram_slice(&self) -> &[u8] {
&self.target.as_ref()[2..]
}
}
impl<Target: AsRef<[u8]>> AsRef<[u8]> for StreamTarget<Target> {
fn as_ref(&self) -> &[u8] {
&self.target.as_ref()[2..]
}
}
impl<Target: AsMut<[u8]>> AsMut<[u8]> for StreamTarget<Target> {
fn as_mut(&mut self) -> &mut [u8] {
&mut self.target.as_mut()[2..]
}
}
impl<Target> OctetsBuilder for StreamTarget<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
type Octets = Target::Octets;
fn append_slice(&mut self, slice: &[u8]) -> Result<(), ShortBuf> {
match self.target.append_slice(slice) {
Ok(()) => {
self.update_shim();
Ok(())
}
Err(ShortBuf) => Err(ShortBuf),
}
}
fn truncate(&mut self, len: usize) {
self.target.truncate(len + 2);
self.update_shim();
}
fn freeze(self) -> Self::Octets {
self.target.freeze()
}
fn len(&self) -> usize {
self.target.len() - 2
}
fn is_empty(&self) -> bool {
self.target.len() > 2
}
}
#[derive(Clone, Debug)]
pub struct StaticCompressor<Target> {
target: Target,
entries: [u16; 24],
len: usize,
}
impl<Target> StaticCompressor<Target> {
pub fn new(target: Target) -> Self {
StaticCompressor {
target,
entries: Default::default(),
len: 0,
}
}
pub fn as_target(&self) -> &Target {
&self.target
}
pub fn into_target(self) -> Target {
self.target
}
pub fn as_slice(&self) -> &[u8]
where
Target: AsRef<[u8]>,
{
self.target.as_ref()
}
pub fn as_slice_mut(&mut self) -> &mut [u8]
where
Target: AsMut<[u8]>,
{
self.target.as_mut()
}
fn get<'a, N: Iterator<Item = &'a Label> + Clone>(
&self,
name: N,
) -> Option<u16>
where
Target: AsRef<[u8]>,
{
self.entries[..self.len].iter().find_map(|&pos| {
if name
.clone()
.eq(Label::iter_slice(self.target.as_ref(), pos as usize))
{
Some(pos)
} else {
None
}
})
}
fn insert(&mut self, pos: usize) -> bool {
if pos < 0xc000 && self.len < self.entries.len() {
self.entries[self.len] = pos as u16;
self.len += 1;
true
} else {
false
}
}
}
impl<Target: AsRef<[u8]>> AsRef<[u8]> for StaticCompressor<Target> {
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
impl<Target: AsMut<[u8]>> AsMut<[u8]> for StaticCompressor<Target> {
fn as_mut(&mut self) -> &mut [u8] {
self.as_slice_mut()
}
}
impl<Target> OctetsBuilder for StaticCompressor<Target>
where
Target: OctetsBuilder + AsRef<[u8]> + AsMut<[u8]>,
{
type Octets = Target::Octets;
fn append_slice(&mut self, slice: &[u8]) -> Result<(), ShortBuf> {
self.target.append_slice(slice)
}
fn truncate(&mut self, len: usize) {
self.target.truncate(len);
if len < 0xC000 {
let len = len as u16;
for i in 0..self.len {
if self.entries[i] >= len {
self.len = i;
break;
}
}
}
}
fn append_compressed_dname<N: ToDname>(
&mut self,
name: &N,
) -> Result<(), ShortBuf> {
let mut name = name.iter_labels().peekable();
loop {
if let Some(label) = name.peek() {
if label.is_root() {
label.compose(self)?;
return Ok(());
}
}
if let Some(pos) = self.get(name.clone()) {
return (pos | 0xC000).compose(self);
}
if !self.insert(self.target.len()) {
for label in &mut name {
label.compose(self)?;
}
return Ok(());
}
let label = name.next().unwrap();
label.compose(self)?;
}
}
fn freeze(self) -> Self::Octets {
self.target.freeze()
}
fn len(&self) -> usize {
self.target.len()
}
fn is_empty(&self) -> bool {
self.target.is_empty()
}
}
#[cfg(feature = "std")]
#[derive(Clone, Debug)]
pub struct TreeCompressor<Target> {
target: Target,
start: Node,
}
#[cfg(feature = "std")]
#[derive(Clone, Debug, Default)]
struct Node {
parents: HashMap<Octets64, Self>,
value: Option<u16>,
}
#[cfg(feature = "std")]
impl Node {
fn drop_above(&mut self, len: u16) {
self.value = match self.value {
Some(value) if value < len => Some(value),
_ => None,
};
self.parents
.values_mut()
.for_each(|node| node.drop_above(len))
}
}
#[cfg(feature = "std")]
impl<Target> TreeCompressor<Target> {
pub fn new(target: Target) -> Self {
TreeCompressor {
target,
start: Default::default(),
}
}
pub fn as_target(&self) -> &Target {
&self.target
}
pub fn into_target(self) -> Target {
self.target
}
pub fn as_slice(&self) -> &[u8]
where
Target: AsRef<[u8]>,
{
self.target.as_ref()
}
pub fn as_slice_mut(&mut self) -> &mut [u8]
where
Target: AsMut<[u8]>,
{
self.target.as_mut()
}
fn get<'a, N: Iterator<Item = &'a Label> + Clone>(
&self,
name: N,
) -> Option<u16> {
let mut node = &self.start;
for label in name {
if label.is_root() {
return node.value;
}
node = node.parents.get(label.as_ref())?;
}
None
}
fn insert<'a, N: Iterator<Item = &'a Label> + Clone>(
&mut self,
name: N,
pos: usize,
) -> bool {
if pos >= 0xC000 {
return false;
}
let pos = pos as u16;
let mut node = &mut self.start;
for label in name {
if label.is_root() {
node.value = Some(pos);
break;
}
node = node
.parents
.entry(label.as_ref().try_into().unwrap())
.or_default();
}
true
}
}
#[cfg(feature = "std")]
impl<Target: AsRef<[u8]>> AsRef<[u8]> for TreeCompressor<Target> {
fn as_ref(&self) -> &[u8] {
self.as_slice()
}
}
#[cfg(feature = "std")]
impl<Target: AsMut<[u8]>> AsMut<[u8]> for TreeCompressor<Target> {
fn as_mut(&mut self) -> &mut [u8] {
self.as_slice_mut()
}
}
#[cfg(feature = "std")]
impl<Target> OctetsBuilder for TreeCompressor<Target>
where
Target: OctetsBuilder + AsMut<[u8]>,
{
type Octets = Target::Octets;
fn append_slice(&mut self, slice: &[u8]) -> Result<(), ShortBuf> {
self.target.append_slice(slice)
}
fn truncate(&mut self, len: usize) {
self.target.truncate(len);
if len < 0xC000 {
self.start.drop_above(len as u16)
}
}
fn append_compressed_dname<N: ToDname>(
&mut self,
name: &N,
) -> Result<(), ShortBuf> {
let mut name = name.iter_labels().peekable();
loop {
if let Some(label) = name.peek() {
if label.is_root() {
label.compose(self)?;
return Ok(());
}
}
if let Some(pos) = self.get(name.clone()) {
return (pos | 0xC000).compose(self);
}
if !self.insert(name.clone(), self.target.len()) {
for label in &mut name {
label.compose(self)?;
}
return Ok(());
}
let label = name.next().unwrap();
label.compose(self)?;
}
}
fn freeze(self) -> Self::Octets {
self.target.freeze()
}
fn len(&self) -> usize {
self.target.len()
}
fn is_empty(&self) -> bool {
self.target.is_empty()
}
}
#[cfg(test)]
#[cfg(feature = "std")]
mod test {
use super::*;
use crate::base::Serial;
use crate::base::{iana::Rtype, opt, Dname};
use crate::rdata::{Ns, Soa, A};
use core::str::FromStr;
use std::vec::Vec;
#[test]
fn message_builder() {
let name = Dname::<Vec<u8>>::from_str("example.com").unwrap();
let mut msg = MessageBuilder::from_target(StaticCompressor::new(
StreamTarget::new_vec(),
))
.unwrap();
msg.header_mut().set_rd(true);
let mut msg = msg.question();
msg.push((&name, Rtype::A)).unwrap();
let mut msg = msg.answer();
msg.push((&name, 86400, A::from_octets(192, 0, 2, 1)))
.unwrap();
msg.push((&name, 86400, A::from_octets(192, 0, 2, 2)))
.unwrap();
let mut msg = msg.authority();
msg.push((&name, 0, Ns::from(name.clone()))).unwrap();
let mut msg = msg.additional();
msg.push((&name, 86400, A::from_octets(192, 0, 2, 1)))
.unwrap();
let target = msg.finish().into_target();
eprintln!("target {}", target.len());
let msg = Message::from_octets(target.as_dgram_slice()).unwrap();
let q = msg.first_question().unwrap();
assert_eq!(q.qname(), &name);
assert_eq!(q.qtype(), Rtype::A);
let section = msg.answer().unwrap();
let mut records = section.limit_to::<A>();
assert_eq!(
records.next().unwrap().unwrap().data(),
&A::from_octets(192, 0, 2, 1)
);
assert_eq!(
records.next().unwrap().unwrap().data(),
&A::from_octets(192, 0, 2, 2)
);
let section = msg.authority().unwrap();
let mut records = section.limit_to::<Ns<_>>();
let rr = records.next().unwrap().unwrap();
assert_eq!(rr.owner(), &name);
assert_eq!(rr.data().nsdname(), &name);
let section = msg.additional().unwrap();
let mut records = section.limit_to::<A>();
let rr = records.next().unwrap().unwrap();
assert_eq!(rr.owner(), &name);
assert_eq!(rr.data(), &A::from_octets(192, 0, 2, 1));
}
#[test]
fn opt_builder() {
let mut msg = MessageBuilder::new_vec().additional();
let nsid = opt::rfc5001::Nsid::from_octets(&b"example"[..]);
msg.opt(|o| {
o.set_udp_payload_size(4096);
o.push(&nsid)?;
Ok(())
})
.unwrap();
let msg = Message::from_octets(msg.finish()).unwrap();
let opt = msg.opt().unwrap();
assert_eq!(opt.udp_payload_size(), 4096);
let mut opts = opt.as_opt().iter::<opt::rfc5001::Nsid<_>>();
assert_eq!(opts.next(), Some(Ok(nsid)));
}
fn create_compressed<T>(target: T) -> T
where
T: OctetsBuilder + AsRef<[u8]> + AsMut<[u8]>,
{
let mut msg = MessageBuilder::from_target(target).unwrap().question();
msg.header_mut().set_rcode(Rcode::NXDomain);
msg.header_mut().set_rd(true);
msg.header_mut().set_ra(true);
msg.header_mut().set_qr(true);
msg.push((&"example".parse::<Dname<Vec<u8>>>().unwrap(), Rtype::Ns))
.unwrap();
let mut msg = msg.authority();
let mname: Dname<Vec<u8>> = "a.root-servers.net".parse().unwrap();
let rname = "nstld.verisign-grs.com".parse().unwrap();
msg.push((
Dname::root_slice(),
86390,
Soa::new(
mname,
rname,
Serial(2020081701),
1800,
900,
604800,
86400,
),
))
.unwrap();
msg.finish()
}
#[test]
fn compressor() {
let expect = &[
0x00, 0x00, 0x81, 0x83, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00,
0x00, 0x07, 0x65, 0x78, 0x61, 0x6d, 0x70, 0x6c, 0x65, 0x00, 0x00,
0x02, 0x00, 0x01, 0x00, 0x00, 0x06, 0x00, 0x01, 0x00, 0x01, 0x51,
0x76, 0x00, 0x40, 0x01, 0x61, 0x0c, 0x72, 0x6f, 0x6f, 0x74, 0x2d,
0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x03, 0x6e, 0x65, 0x74,
0x00, 0x05, 0x6e, 0x73, 0x74, 0x6c, 0x64, 0x0c, 0x76, 0x65, 0x72,
0x69, 0x73, 0x69, 0x67, 0x6e, 0x2d, 0x67, 0x72, 0x73, 0x03, 0x63,
0x6f, 0x6d, 0x00, 0x78, 0x68, 0x00, 0x25, 0x00, 0x00, 0x07, 0x08,
0x00, 0x00, 0x03, 0x84, 0x00, 0x09, 0x3a, 0x80, 0x00, 0x01, 0x51,
0x80,
];
let msg = create_compressed(StaticCompressor::new(Vec::new()));
assert_eq!(&expect[..], msg.as_ref());
let msg = create_compressed(TreeCompressor::new(Vec::new()));
assert_eq!(&expect[..], msg.as_ref());
}
}