use crate::constants::*;
use crate::dns_sector::*;
use crate::errors::*;
use crate::rr_iterator::*;
use byteorder::{BigEndian, ByteOrder};
use std::cmp;
#[derive(Copy, Clone, Debug)]
pub struct UncompressedNameResult {
pub name_len: usize,
pub final_offset: usize,
}
#[derive(Copy, Clone, Debug)]
pub struct CompressedNameResult {
pub name_len: usize,
pub final_offset: usize,
}
pub struct Compress;
impl Compress {
pub fn check_compressed_name(packet: &[u8], mut offset: usize) -> Result<usize, Error> {
let packet_len = packet.len();
let mut name_len = 0;
let (mut barrier_offset, mut lowest_offset, mut final_offset) = (packet_len, offset, None);
let mut refs_allowed = DNS_MAX_HOSTNAME_INDIRECTIONS;
if offset >= packet_len {
bail!(DSError::InternalError("Offset outside packet boundaries"));
}
if 1 > packet_len - offset {
bail!(DSError::InvalidName("Empty name"));
}
loop {
if offset >= barrier_offset {
if offset >= packet_len {
bail!(DSError::InvalidName("Truncated name"));
}
bail!(DSError::InvalidName("Cycle"));
}
let label_len = match packet[offset] {
len if len & 0xc0 == 0xc0 => {
if refs_allowed <= 0 {
bail!(DSError::InvalidName("Too many indirections"));
}
refs_allowed -= 1;
if 2 > packet_len - offset {
bail!(DSError::InvalidName("Invalid internal offset"));
}
let ref_offset =
((((len & 0x3f) as u16) << 8) | (packet[offset + 1]) as u16) as usize;
if ref_offset == offset || ref_offset >= lowest_offset {
bail!(DSError::InvalidName("Forward/self reference"));
}
if packet[ref_offset] & 0xc0 != 0xc0 && packet[ref_offset] < 1 {
bail!(DSError::InvalidName("Reference to an empty label"));
}
final_offset = final_offset.or(Some(offset + 2));
offset = ref_offset;
barrier_offset = lowest_offset;
lowest_offset = ref_offset;
continue;
}
len if len > 0x3f => bail!(DSError::InvalidName("Label length too long")),
len => len as usize,
};
if label_len >= packet_len - offset {
bail!(DSError::InvalidName("Out-of-bounds name"));
}
name_len += label_len + 1;
if name_len > DNS_MAX_HOSTNAME_LEN {
bail!(DSError::InvalidName("Name too long"));
}
if packet[offset + 1..offset + label_len + 1]
.iter()
.any(|&c| c.is_ascii_control() || c == b'.' || c == b'\\' || c == 0)
{
bail!(DSError::InvalidName("Unexpected character in name"));
}
offset += label_len + 1;
if label_len == 0 {
break;
}
}
let final_offset = final_offset.unwrap_or(offset);
Ok(final_offset)
}
pub fn copy_uncompressed_name(
name: &mut Vec<u8>,
packet: &[u8],
mut offset: usize,
) -> UncompressedNameResult {
let mut name_len = 0;
let mut final_offset = None;
loop {
let label_len = match packet[offset] {
len if len & 0xc0 == 0xc0 => {
final_offset = final_offset.or(Some(offset + 2));
let new_offset = (BigEndian::read_u16(&packet[offset..]) & 0x3fff) as usize;
assert!(new_offset < offset);
offset = new_offset;
continue;
}
len => len,
} as usize;
let prefixed_label_len = 1 + label_len;
name.extend(&packet[offset..offset + prefixed_label_len]);
name_len += prefixed_label_len;
offset += prefixed_label_len;
if label_len == 0 {
break;
}
}
let final_offset = final_offset.unwrap_or(offset);
UncompressedNameResult {
name_len,
final_offset,
}
}
fn uncompress_rdata(
uncompressed: &mut Vec<u8>,
raw: RRRaw<'_>,
rr_type: Option<u16>,
rr_rdlen: Option<usize>,
) {
let packet = &raw.packet;
let offset_rdata = raw.name_end;
let rdata = &packet[offset_rdata..];
match rr_type {
None => {
debug_assert!(rr_rdlen.is_none());
uncompressed.extend_from_slice(&rdata[..DNS_RR_QUESTION_HEADER_SIZE]);
}
Some(x) if x == Type::NS.into() || x == Type::CNAME.into() || x == Type::PTR.into() => {
let offset = uncompressed.len();
uncompressed.extend_from_slice(&rdata[..DNS_RR_HEADER_SIZE]);
let new_rdlen = Compress::copy_uncompressed_name(
uncompressed,
packet,
offset_rdata + DNS_RR_HEADER_SIZE,
)
.name_len;
BigEndian::write_u16(
&mut uncompressed[offset + DNS_RR_RDLEN_OFFSET..],
new_rdlen as u16,
);
}
Some(x) if x == Type::MX.into() => {
let offset = uncompressed.len();
uncompressed.extend_from_slice(&rdata[..DNS_RR_HEADER_SIZE + 2]);
let new_rdlen = 2 + Compress::copy_uncompressed_name(
uncompressed,
packet,
offset_rdata + DNS_RR_HEADER_SIZE + 2,
)
.name_len;
BigEndian::write_u16(
&mut uncompressed[offset + DNS_RR_RDLEN_OFFSET..],
new_rdlen as u16,
);
}
Some(x) if x == Type::SOA.into() => {
let offset = uncompressed.len();
uncompressed.extend_from_slice(&rdata[..DNS_RR_HEADER_SIZE]);
let u1 = Compress::copy_uncompressed_name(
uncompressed,
packet,
offset_rdata + DNS_RR_HEADER_SIZE,
);
let u2 =
Compress::copy_uncompressed_name(uncompressed, packet, u1.final_offset);
uncompressed.extend_from_slice(&packet[u2.final_offset..u2.final_offset + 20]);
let new_rdlen = u1.name_len + u2.name_len + 20;
BigEndian::write_u16(
&mut uncompressed[offset + DNS_RR_RDLEN_OFFSET..],
new_rdlen as u16,
);
}
_ => {
uncompressed.extend_from_slice(&rdata[..DNS_RR_HEADER_SIZE + rr_rdlen.unwrap()]);
}
}
}
pub fn compress_rdata(
dict: &mut SuffixDict,
compressed: &mut Vec<u8>,
raw: RRRaw<'_>,
rr_type: Option<u16>,
rr_rdlen: Option<usize>,
) {
let packet = &raw.packet;
let offset_rdata = raw.name_end;
let rdata = &packet[offset_rdata..];
match rr_type {
None => {
debug_assert!(rr_rdlen.is_none());
compressed.extend_from_slice(&rdata[..DNS_RR_QUESTION_HEADER_SIZE]);
}
Some(x) if x == Type::NS.into() || x == Type::CNAME.into() || x == Type::PTR.into() => {
let offset = compressed.len();
compressed.extend_from_slice(&rdata[..DNS_RR_HEADER_SIZE]);
let new_rdlen = Compress::copy_compressed_name(
dict,
compressed,
packet,
offset_rdata + DNS_RR_HEADER_SIZE,
)
.name_len;
BigEndian::write_u16(
&mut compressed[offset + DNS_RR_RDLEN_OFFSET..],
new_rdlen as u16,
);
}
Some(x) if x == Type::MX.into() => {
let offset = compressed.len();
compressed.extend_from_slice(&rdata[..DNS_RR_HEADER_SIZE + 2]);
let new_rdlen = 2 + Compress::copy_compressed_name(
dict,
compressed,
packet,
offset_rdata + DNS_RR_HEADER_SIZE + 2,
)
.name_len;
BigEndian::write_u16(
&mut compressed[offset + DNS_RR_RDLEN_OFFSET..],
new_rdlen as u16,
);
}
Some(x) if x == Type::SOA.into() => {
let offset = compressed.len();
compressed.extend_from_slice(&rdata[..DNS_RR_HEADER_SIZE]);
let u1 = Compress::copy_compressed_name(
dict,
compressed,
packet,
offset_rdata + DNS_RR_HEADER_SIZE,
);
let u2 = Compress::copy_compressed_name(
dict,
compressed,
packet,
u1.final_offset,
);
compressed.extend_from_slice(&packet[u2.final_offset..u2.final_offset + 20]);
let new_rdlen = u1.name_len + u2.name_len + 20;
BigEndian::write_u16(
&mut compressed[offset + DNS_RR_RDLEN_OFFSET..],
new_rdlen as u16,
);
}
_ => {
compressed.extend_from_slice(&rdata[..DNS_RR_HEADER_SIZE + rr_rdlen.unwrap()]);
}
}
}
pub fn uncompress_with_previous_offset(
packet: &[u8],
ref_offset: usize,
) -> Result<(Vec<u8>, usize), Error> {
let packet = packet.to_owned();
if packet.len() < DNS_HEADER_SIZE {
bail!(DSError::PacketTooSmall);
}
let mut new_offset = None;
let mut uncompressed = Vec::new();
uncompressed.extend_from_slice(&packet[..DNS_HEADER_SIZE]);
let mut parsed_packet = DNSSector::new(packet)?.parse()?;
{
let mut it = parsed_packet.into_iter_question();
while let Some(item) = it {
if Some(ref_offset) == item.offset() {
new_offset = Some(uncompressed.len());
}
item.copy_raw_name(&mut uncompressed);
Self::uncompress_rdata(&mut uncompressed, item.raw(), None, None);
it = item.next();
}
}
{
let mut it = parsed_packet.into_iter_answer();
while let Some(item) = it {
if Some(ref_offset) == item.offset() {
new_offset = Some(uncompressed.len());
}
item.copy_raw_name(&mut uncompressed);
Self::uncompress_rdata(
&mut uncompressed,
item.raw(),
Some(item.rr_type()),
Some(item.rr_rdlen()),
);
it = item.next();
}
}
{
let mut it = parsed_packet.into_iter_nameservers();
while let Some(item) = it {
if Some(ref_offset) == item.offset() {
new_offset = Some(uncompressed.len());
}
item.copy_raw_name(&mut uncompressed);
Self::uncompress_rdata(
&mut uncompressed,
item.raw(),
Some(item.rr_type()),
Some(item.rr_rdlen()),
);
it = item.next();
}
}
{
let mut it = parsed_packet.into_iter_additional_including_opt();
while let Some(item) = it {
if Some(ref_offset) == item.offset() {
new_offset = Some(uncompressed.len());
}
item.copy_raw_name(&mut uncompressed);
Self::uncompress_rdata(
&mut uncompressed,
item.raw(),
Some(item.rr_type()),
Some(item.rr_rdlen()),
);
it = item.next_including_opt();
}
}
if ref_offset == parsed_packet.packet().len() {
new_offset = Some(uncompressed.len());
}
Ok((
uncompressed,
new_offset.expect("Previous offset not found at a record boundary"),
))
}
pub fn uncompress(packet: &[u8]) -> Result<Vec<u8>, Error> {
Self::uncompress_with_previous_offset(packet, DNS_HEADER_SIZE).map(|x| x.0)
}
pub fn compress(packet: &[u8]) -> Result<Vec<u8>, Error> {
let packet = packet.to_owned();
if packet.len() < DNS_HEADER_SIZE {
bail!(DSError::PacketTooSmall);
}
let mut compressed = Vec::new();
compressed.extend_from_slice(&packet[..DNS_HEADER_SIZE]);
let mut parsed_packet = DNSSector::new(packet)?.parse()?;
let mut dict = SuffixDict::new();
{
let mut it = parsed_packet.into_iter_question();
while let Some(item) = it {
{
let mut raw = item.raw();
raw.offset = Self::copy_compressed_name(
&mut dict,
&mut compressed,
raw.packet,
raw.offset,
)
.final_offset;
Self::compress_rdata(&mut dict, &mut compressed, raw, None, None);
}
it = item.next();
}
}
{
let mut it = parsed_packet.into_iter_answer();
while let Some(item) = it {
{
let mut raw = item.raw();
raw.offset = Self::copy_compressed_name(
&mut dict,
&mut compressed,
raw.packet,
raw.offset,
)
.final_offset;
Self::compress_rdata(
&mut dict,
&mut compressed,
raw,
Some(item.rr_type()),
Some(item.rr_rdlen()),
);
}
it = item.next();
}
}
{
let mut it = parsed_packet.into_iter_nameservers();
while let Some(item) = it {
{
let mut raw = item.raw();
raw.offset = Self::copy_compressed_name(
&mut dict,
&mut compressed,
raw.packet,
raw.offset,
)
.final_offset;
Self::compress_rdata(
&mut dict,
&mut compressed,
raw,
Some(item.rr_type()),
Some(item.rr_rdlen()),
);
}
it = item.next();
}
}
{
let mut it = parsed_packet.into_iter_additional_including_opt();
while let Some(item) = it {
{
let mut raw = item.raw();
raw.offset = Self::copy_compressed_name(
&mut dict,
&mut compressed,
raw.packet,
raw.offset,
)
.final_offset;
Self::compress_rdata(
&mut dict,
&mut compressed,
raw,
Some(item.rr_type()),
Some(item.rr_rdlen()),
);
}
it = item.next();
}
}
Ok(compressed)
}
pub fn raw_name_len(name: &[u8]) -> usize {
let mut i = 0;
while name[i] != 0 {
let len = name[i] as usize;
if len & 0xc0 == 0xc0 {
i += 1;
break;
}
i += len + 1
}
i + 1
}
pub fn raw_name_len_after_decompression(packet: &[u8], mut offset: usize) -> usize {
let mut name_len = 0;
loop {
let label_len = match packet[offset] {
len if len & 0xc0 == 0xc0 => {
let new_offset = (BigEndian::read_u16(&packet[offset..]) & 0x3fff) as usize;
assert!(new_offset < offset);
offset = new_offset;
continue;
}
len => len,
} as usize;
let prefixed_label_len = 1 + label_len;
name_len += prefixed_label_len;
offset += prefixed_label_len;
if label_len == 0 {
break;
}
}
name_len
}
pub fn raw_name_to_str(packet: &[u8], mut offset: usize) -> Vec<u8> {
let mut indirections = 0;
let mut res: Vec<u8> = Vec::with_capacity(64);
loop {
let label_len = match packet[offset] {
0 => break,
len if len & 0xc0 == 0xc0 => {
let new_offset = (BigEndian::read_u16(&packet[offset..]) & 0x3fff) as usize;
if new_offset == offset || indirections > DNS_MAX_HOSTNAME_INDIRECTIONS {
return res;
}
indirections += 1;
offset = new_offset;
continue;
}
len => len,
} as usize;
offset += 1;
let label = &packet[offset..offset + label_len];
offset += label_len;
if !res.is_empty() {
res.push(b'.');
}
for &c in label {
match c {
b'.' => res.extend(b"\\046"),
_ => res.push(c),
}
}
}
res
}
pub fn copy_compressed_name_with_base_offset(
dict: &mut SuffixDict,
compressed: &mut Vec<u8>,
packet: &[u8],
mut offset: usize,
base_offset: usize,
) -> CompressedNameResult {
let uncompressed_name_len = Compress::raw_name_len_after_decompression(packet, offset);
let initial_compressed_len = compressed.len();
let final_offset = offset + uncompressed_name_len;
loop {
let label_len = packet[offset] as usize;
if label_len & 0xc0 == 0xc0 {
panic!("copy_compressed_name() called on an already compressed name");
}
if let Some(ref_offset) =
dict.insert(&packet[offset..final_offset], base_offset + offset)
{
assert!(offset < 65536 >> 2); compressed.push((ref_offset >> 8) as u8 | 0xc0);
compressed.push((ref_offset & 0xff) as u8);
break;
}
let offset_next = offset + 1 + label_len;
compressed.extend_from_slice(&packet[offset..offset_next]);
offset = offset_next;
if label_len == 0 {
break;
}
}
CompressedNameResult {
name_len: compressed.len() - initial_compressed_len,
final_offset,
}
}
pub fn copy_compressed_name(
dict: &mut SuffixDict,
compressed: &mut Vec<u8>,
packet: &[u8],
offset: usize,
) -> CompressedNameResult {
Self::copy_compressed_name_with_base_offset(dict, compressed, packet, offset, 0)
}
}
const MAX_SUFFIX_LEN: usize = 127;
const MAX_SUFFIXES: usize = 32;
struct Suffix {
offset: usize,
len: usize,
suffix: [u8; MAX_SUFFIX_LEN],
}
impl Default for Suffix {
fn default() -> Self {
Self {
offset: 0,
len: 0,
suffix: [0u8; MAX_SUFFIX_LEN],
}
}
}
#[derive(Default)]
pub struct SuffixDict {
count: usize,
index: usize,
suffixes: [Suffix; MAX_SUFFIXES],
}
impl SuffixDict {
pub fn new() -> Self {
SuffixDict::default()
}
fn insert(&mut self, suffix: &[u8], offset: usize) -> Option<usize> {
if offset >= 65536 >> 2 {
return None;
}
let suffix_len = suffix.len();
if suffix_len <= 2 || suffix_len > MAX_SUFFIX_LEN {
return None;
}
for i in 0..self.count {
let candidate = &self.suffixes[i];
if candidate.len <= suffix_len
&& Self::raw_names_eq_ignore_case(suffix, &candidate.suffix[..candidate.len])
{
return Some(candidate.offset);
}
}
let entry = &mut self.suffixes[self.index];
let len = Self::raw_name_copy(&mut entry.suffix, suffix);
debug_assert_eq!(len, suffix_len);
entry.len = suffix_len;
entry.offset = offset;
self.index += 1;
self.count = cmp::max(self.index, self.count);
if self.index == MAX_SUFFIXES {
debug_assert!(MAX_SUFFIXES > 1);
self.index = 1; }
None
}
fn raw_name_copy(to: &mut [u8], name: &[u8]) -> usize {
let len = Compress::raw_name_len(name);
to[..len].copy_from_slice(&name[..len]);
len
}
fn raw_names_eq_ignore_case(name1: &[u8], name2: &[u8]) -> bool {
let mut label_len = 0;
for (&c1, &c2) in name1.iter().zip(name2.iter()) {
if !c1.eq_ignore_ascii_case(&c2) {
return false;
}
if label_len == 0 {
if c1 == 0 {
return true;
}
label_len = c1;
} else {
label_len -= 1
}
}
false
}
}