rust-dnsbox/lib/dnsbox-base/src/common_types/nsec.rs

use bytes::{Bytes, Buf, BufMut};
use common_types::Type;
use data_encoding;
use errors::*;
use failure::{Fail, ResultExt};
use ser::packet::{DnsPacketData, DnsPacketWriteContext, remaining_bytes, short_blob, write_short_blob};
use ser::text::{DnsTextData, DnsTextFormatter, DnsTextContext, skip_whitespace, next_field};
use std::collections::BTreeSet;
use std::fmt;
use std::io::Cursor;

static WHITESPACE: &str = "\t\n\x0c\r "; // \f == \x0c formfeed

lazy_static!{
	static ref BASE32HEX_NOPAD_ALLOW_WS: data_encoding::Encoding = {
		let mut spec = data_encoding::Specification::new();
		spec.symbols.push_str("0123456789ABCDEFGHIJKLMNOPQRSTUV");
		spec.translate.from.push_str("abcdefghijklmnopqrstuv");
		spec.translate.to.push_str("ABCDEFGHIJKLMNOPQRSTUV");
		spec.ignore.push_str(WHITESPACE);
		spec.encoding().unwrap()
	};
}

/// Last field in a record!
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct NsecTypeBitmap {
	raw: Bytes,
	set: BTreeSet<Type>,
}

impl NsecTypeBitmap {
	pub fn from_set(set: BTreeSet<Type>) -> Self {
		let mut raw = Vec::new();
		let mut current_window_base = None;
		let mut window_pos = 0;
		let mut window_len = 0;
		for t in &set {
			let base = (t.0 >> 8) as u8;
			if current_window_base != Some(base) {
				// new window
				current_window_base = Some(base);
				window_pos = raw.len();
				window_len = 0;
				raw.push(base);
				raw.push(0);
			}

			let bit_ndx = t.0 as u8;
			let byte_ndx = bit_ndx / 8;
			if byte_ndx >= window_len {
				// make window bigger, fill with zeroes
				let new_window_len = byte_ndx + 1;
				let new_raw_len = raw.len() + (new_window_len - window_len) as usize;
				raw.resize(new_raw_len, 0);
				raw[window_pos + 1] = new_window_len;
				window_len = new_window_len;
			}

			let mask = 0x80 >> (bit_ndx % 8);
			raw[window_pos + 2 + byte_ndx as usize] |= mask;
		}

		NsecTypeBitmap {
			raw: raw.into(),
			set: set,
		}
	}
}

impl DnsPacketData for NsecTypeBitmap {
	fn deserialize(data: &mut Cursor<Bytes>) -> ::errors::Result<Self> {
		// remember raw encoding
		let raw = {
			let mut data: Cursor<Bytes> = data.clone();
			remaining_bytes(&mut data)
		};
		let mut set = BTreeSet::new();
		let mut prev_window = None;
		while data.has_remaining() {
			let window_base = (data.get_u8() as u16) << 8;
			ensure!(Some(window_base) > prev_window, "wrong nsec bitmap window order, {:?} <= {:?}", Some(window_base), prev_window);
			prev_window = Some(window_base);
			check_enough_data!(data, 1, "nsec bitmap window length");
			let window_len = data.get_u8() as u16;
			ensure!(window_len <= 32, "nsec bitmap window too long");
			check_enough_data!(data, window_len as usize, "nsec bitmap window length");
			for i in 0..window_len {
				let mut v = data.get_u8();
				for j in 0..8 {
					if 0 != v & 0x80 {
						set.insert(Type(window_base + i*8 + j));
					}
					v <<= 1;
				}
			}
		}
		Ok(NsecTypeBitmap{
			raw: raw,
			set: set,
		})
	}

	fn serialize(&self, _context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
		packet.reserve(self.raw.len());
		packet.put_slice(&self.raw);
		Ok(())
	}
}

impl DnsTextData for NsecTypeBitmap {
	fn dns_parse(context: &DnsTextContext, data: &mut &str) -> ::errors::Result<Self> {
		let mut set = BTreeSet::new();
		skip_whitespace(data);
		while !data.is_empty() {
			let t = Type::dns_parse(context, data)?;
			set.insert(t);
		}
		Ok(NsecTypeBitmap::from_set(set))
	}

	fn dns_format(&self, f: &mut DnsTextFormatter) -> fmt::Result {
		for t in &self.set {
			write!(f, "{}", t)?;
		}
		Ok(())
	}
}

/// `base32hex` encoding without padding for text representation (RFC
/// 4648). whitespaces not allowed (not the last field).
///
/// Between 1 and 255 bytes long.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct NextHashedOwnerName(Bytes);

impl DnsPacketData for NextHashedOwnerName {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
		let text = short_blob(data)?;
		ensure!(text.len() > 0, "NextHashedOwnerName must not be empty");
		Ok(NextHashedOwnerName(text))
	}

	fn serialize(&self, _context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
		ensure!(self.0.len() > 0, "NextHashedOwnerName must not be empty");
		write_short_blob(&self.0, packet)
	}
}

impl DnsTextData for NextHashedOwnerName {
	fn dns_parse(_context: &DnsTextContext, data: &mut &str) -> ::errors::Result<Self> {
		let field = next_field(data)?;
		let raw = BASE32HEX_NOPAD_ALLOW_WS.decode(field.as_bytes())
			.with_context(|e| e.context(format!("invalid base32hex (no padding): {:?}", field)))?;
		ensure!(raw.len() > 0, "NextHashedOwnerName must not be empty");
		ensure!(raw.len() < 256, "NextHashedOwnerName field must be at most 255 bytes long");
		Ok(NextHashedOwnerName(raw.into()))
	}

	fn dns_format(&self, f: &mut DnsTextFormatter) -> fmt::Result {
		if self.0.is_empty() { return Err(fmt::Error); }
		write!(f, "{}", BASE32HEX_NOPAD_ALLOW_WS.encode(&self.0))
	}
}