rust-dnsbox/lib/dnsbox-base/src/common_types/name/name_packet_parser.rs

use super::*;
use bytes::Buf;

/// `data`: bytes of packet from beginning until at least the end of the name
/// `start_pos`: position of first byte of the name
/// `uncmpr_offsets`: offsets of uncompressed labels so far
/// `label_len`: first compressed  label length (`0xc0 | offset-high, offset-low`)
/// `total_len`: length of (uncompressed) label encoding so far
fn deserialize_name_compressed_cont(
	data: Bytes,
	start_pos: usize,
	uncmpr_offsets: SmallVec<[u8; 16]>,
	mut total_len: usize,
	mut label_len: u8,
) -> Result<DnsName> {
	let mut label_offsets = uncmpr_offsets
		.into_iter()
		.map(LabelOffset::LabelStart)
		.collect::<SmallVec<_>>();

	let mut pos = start_pos + total_len;
	'next_compressed: loop {
		{
			failure::ensure!(pos + 1 < data.len(), "not enough data for compressed label");
			let new_pos = ((label_len as usize & 0x3f) << 8) | (data[pos + 1] as usize);
			failure::ensure!(
				new_pos < pos,
				"Compressed label offset too big: {} >= {}",
				new_pos,
				pos
			);
			pos = new_pos;
		}

		loop {
			failure::ensure!(pos < data.len(), "not enough data for label");
			label_len = data[pos];

			if 0 == label_len {
				return Ok(DnsName {
					data: data,
					label_offsets: LabelOffsets::Compressed(start_pos, label_offsets),
					total_len: total_len as u8 + 1,
				});
			}

			if label_len & 0xc0 == 0xc0 {
				continue 'next_compressed;
			}
			failure::ensure!(label_len < 64, "Invalid label length {}", label_len);

			total_len += 1 + label_len as usize;
			// max len 255, but there also needs to be an empty label at the end
			if total_len > 254 {
				failure::bail!("DNS name too long")
			}

			label_offsets.push(LabelOffset::PacketStart(pos as u16));
			pos += 1 + label_len as usize;
		}
	}
}

pub fn deserialize_name(data: &mut Cursor<Bytes>, accept_compressed: bool) -> Result<DnsName> {
	check_enough_data!(data, 1, "DnsName");
	let start_pos = data.position() as usize;
	let mut total_len: usize = 0;
	let mut label_offsets = SmallVec::new();
	loop {
		check_enough_data!(data, 1, "DnsName label len");
		let label_len = data.get_u8() as usize;
		if 0 == label_len {
			let end_pos = data.position() as usize;
			return Ok(DnsName {
				data: data.get_ref().slice(start_pos, end_pos),
				label_offsets: LabelOffsets::Uncompressed(label_offsets),
				total_len: total_len as u8 + 1,
			});
		}
		if label_len & 0xc0 == 0xc0 {
			// compressed label
			if !accept_compressed {
				failure::bail!("Invalid label compression {}", label_len)
			}
			check_enough_data!(data, 1, "DnsName compressed label target");
			// eat second part of compressed label
			data.get_u8();

			let end_pos = data.position() as usize;
			let data = data.get_ref().slice(0, end_pos);

			return deserialize_name_compressed_cont(
				data,
				start_pos,
				label_offsets,
				total_len,
				label_len as u8,
			);
		}
		label_offsets.push(total_len as u8);
		if label_len > 63 {
			failure::bail!("Invalid label length {}", label_len)
		}
		total_len += 1 + label_len;
		// max len 255, but there also needs to be an empty label at the end
		if total_len > 254 {
			failure::bail! {"DNS name too long"}
		}
		check_enough_data!(data, label_len, "DnsName label");
		data.advance(label_len);
	}
}