dns-message-parser 0.2.0

Libary to encode and decode DNS packets
Documentation
use bytes::BytesMut;

use crate::{Class, DomainName, Type};

use std::collections::HashMap;
use std::mem::size_of;

use super::{encode_string, encode_u16, encode_u32, EncodeError, EncodeResult};

pub(super) struct EncodeData<'a> {
    pub(super) bytes: &'a mut BytesMut,
    pub(super) bytes_rdata: BytesMut,
    pub(super) class: &'a Class,
    pub(super) ttl: u32,
}

impl<'a> EncodeData<'a> {
    pub(super) fn new(bytes: &'a mut BytesMut, class: &'a Class, ttl: u32) -> EncodeData<'a> {
        EncodeData {
            bytes,
            bytes_rdata: BytesMut::new(),
            class,
            ttl,
        }
    }

    pub(super) fn encode_generic_rr_header(&mut self, type_: Type) -> EncodeResult {
        type_.encode(self.bytes)?;
        self.class.encode(self.bytes)?;
        encode_u32(self.bytes, self.ttl);
        Ok(())
    }

    pub(super) fn get_offset(&self) -> usize {
        self.bytes.len() + size_of::<u16>()
    }

    pub(super) fn encode_domain(
        &mut self,
        type_: Type,
        domain_name: &DomainName,
        compression: &mut HashMap<String, usize>,
    ) -> EncodeResult {
        self.encode_generic_rr_header(type_)?;
        let offset = self.get_offset();
        domain_name.encode(&mut self.bytes_rdata, &offset, compression)
    }

    pub(super) fn encode_u16_domain(
        &mut self,
        type_: Type,
        u: u16,
        domain_name: &DomainName,
        compression: &mut HashMap<String, usize>,
    ) -> EncodeResult {
        self.encode_generic_rr_header(type_)?;
        let offset = self.get_offset();
        encode_u16(&mut self.bytes_rdata, u);
        domain_name.encode(&mut self.bytes_rdata, &offset, compression)
    }

    pub(super) fn encode_domain_domain(
        &mut self,
        type_: Type,
        domain_name_1: &DomainName,
        domain_name_2: &DomainName,
        compression: &mut HashMap<String, usize>,
    ) -> EncodeResult {
        self.encode_generic_rr_header(type_)?;
        let offset = self.get_offset();
        domain_name_1.encode(&mut self.bytes_rdata, &offset, compression)?;
        domain_name_2.encode(&mut self.bytes_rdata, &offset, compression)
    }

    pub(super) fn encode_vec(&mut self, type_: Type, vec: &[u8]) -> EncodeResult {
        self.encode_generic_rr_header(type_)?;
        self.bytes_rdata.extend(vec);
        Ok(())
    }

    pub(super) fn encode_string(&mut self, type_: Type, string: &str) -> EncodeResult {
        self.encode_generic_rr_header(type_)?;
        encode_string(&mut self.bytes_rdata, string)
    }

    pub(super) fn add_rdata(&mut self) -> EncodeResult {
        let length = self.bytes_rdata.len();
        if (std::u16::MAX as usize) < length {
            return Err(EncodeError::TooMuchData);
        }
        encode_u16(self.bytes, length as u16);
        self.bytes.extend(&self.bytes_rdata);
        Ok(())
    }
}