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

use bytes::{Bytes, BufMut};
use data_encoding::{self, HEXLOWER_PERMISSIVE};
use crate::errors::*;
use failure::{Fail, ResultExt};
use crate::ser::packet::{DnsPacketData, DnsPacketWriteContext, remaining_bytes, short_blob, write_short_blob, get_blob};
use crate::ser::text::*;
use std::fmt;
use std::io::Cursor;

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

lazy_static::lazy_static!{
	pub(crate) static ref HEXLOWER_PERMISSIVE_ALLOW_WS: data_encoding::Encoding = {
		let mut spec = data_encoding::Specification::new();
		spec.symbols.push_str("0123456789abcdef");
		spec.translate.from.push_str("ABCDEF");
		spec.translate.to.push_str("abcdef");
		spec.ignore.push_str(WHITESPACE);
		spec.encoding().unwrap()
	};

	static ref BASE64_ALLOW_WS: data_encoding::Encoding = {
		let mut spec = data_encoding::Specification::new();
		spec.symbols.push_str("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/");
		spec.padding = Some('=');
		spec.ignore.push_str(WHITESPACE);
		spec.encoding().unwrap()
	};
}

// Similar to `ShortText`, but uses (unquoted, no spaces allowed) hex
// for text representation; "-" when empty
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct HexShortBlob(Bytes);

impl DnsPacketData for HexShortBlob {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
		Ok(HexShortBlob(short_blob(data)?))
	}

	fn serialize(&self, _context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
		write_short_blob(&self.0, packet)
	}
}

impl DnsTextData for HexShortBlob {
	fn dns_parse(_context: &DnsTextContext, data: &mut &str) -> crate::errors::Result<Self> {
		let s = next_field(data)?;
		if s == "-" {
			Ok(HexShortBlob(Bytes::new()))
		} else {
			let raw = HEXLOWER_PERMISSIVE.decode(s.as_bytes())
				.with_context(|e| e.context(format!("invalid hex: {:?}", s)))?;
			failure::ensure!(raw.len() < 256, "short hex field must be at most 255 bytes long");
			Ok(HexShortBlob(raw.into()))
		}
	}

	fn dns_format(&self, f: &mut DnsTextFormatter) -> fmt::Result {
		if self.0.is_empty() {
			write!(f, "-")
		} else {
			write!(f, "{}", HEXLOWER_PERMISSIVE.encode(&self.0))
		}
	}
}

impl std::ops::Deref for HexShortBlob {
	type Target = [u8];

	fn deref(&self) -> &Self::Target {
		&self.0
	}
}

// 16-bit length, uses decimal length + base64 encoding (if length > 0)
// for text representation.
// 
// In base64 encoding no whitespace allowed and padding required.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Base64LongBlob(Bytes);

impl DnsPacketData for Base64LongBlob {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
		let len = u16::deserialize(data)? as usize;
		check_enough_data!(data, len, "data for long blob");

		Ok(Base64LongBlob(get_blob(data, len)?))
	}

	fn serialize(&self, context: &mut DnsPacketWriteContext, packet: &mut Vec<u8>) -> Result<()> {
		let len = self.0.len();
		failure::ensure!(len < 0x1_0000, "blob too long");
		(len as u16).serialize(context, packet)?;
		packet.reserve(len);
		packet.put_slice(&self.0);
		Ok(())
	}
}

impl DnsTextData for Base64LongBlob {
	fn dns_parse(_context: &DnsTextContext, data: &mut &str) -> crate::errors::Result<Self> {
		let length_field = next_field(data)?;
		let length = length_field.parse::<u16>()
			.with_context(|_| format!("invalid length for blob: {:?}", length_field))?;

		if length > 0 {
			let blob_field = next_field(data)?;
			let result = BASE64_ALLOW_WS.decode(blob_field.as_bytes())
				.with_context(|e| e.context(format!("invalid base64: {:?}", blob_field)))?;
			Ok(Base64LongBlob(result.into()))
		} else {
			Ok(Base64LongBlob(Bytes::new()))
		}
	}

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

impl std::ops::Deref for Base64LongBlob {
	type Target = [u8];

	fn deref(&self) -> &Self::Target {
		&self.0
	}
}

/// No length byte (or restriction), just all data to end of record. uses
/// base64 encoding for text representation, whitespace allowed, padding
/// required.
///
/// No following field allowed, i.e. last field in the record.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Base64RemainingBlob(Bytes);

impl DnsPacketData for Base64RemainingBlob {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
		Ok(Base64RemainingBlob(remaining_bytes(data)))
	}

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

impl DnsTextData for Base64RemainingBlob {
	fn dns_parse(_context: &DnsTextContext, data: &mut &str) -> crate::errors::Result<Self> {
		skip_whitespace(data);
		let result = BASE64_ALLOW_WS.decode(data.as_bytes())
			.with_context(|e| e.context(format!("invalid base64: {:?}", data)))?;
		*data = "";
		Ok(Base64RemainingBlob(result.into()))
	}

	fn dns_format(&self, f: &mut DnsTextFormatter) -> fmt::Result {
		if !self.0.is_empty() {
			write!(f, "{}", BASE64_ALLOW_WS.encode(&self.0))
		} else {
			Ok(())
		}
	}
}

impl std::ops::Deref for Base64RemainingBlob {
	type Target = [u8];

	fn deref(&self) -> &Self::Target {
		&self.0
	}
}

/// No length byte (or restriction), just all data to end of record. uses
/// hex encoding for text representation, whitespace allowed.
///
/// No following field allowed, i.e. last field in the record.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct HexRemainingBlob(Bytes);

impl DnsPacketData for HexRemainingBlob {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
		Ok(HexRemainingBlob(remaining_bytes(data)))
	}

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

impl DnsTextData for HexRemainingBlob {
	fn dns_parse(_context: &DnsTextContext, data: &mut &str) -> crate::errors::Result<Self> {
		skip_whitespace(data);
		let result = HEXLOWER_PERMISSIVE_ALLOW_WS.decode(data.as_bytes())
			.with_context(|e| e.context(format!("invalid hex: {:?}", data)))?;
		*data = "";
		Ok(HexRemainingBlob(result.into()))
	}

	fn dns_format(&self, f: &mut DnsTextFormatter) -> fmt::Result {
		write!(f, "{}", HEXLOWER_PERMISSIVE_ALLOW_WS.encode(&self.0))
	}
}

impl std::ops::Deref for HexRemainingBlob {
	type Target = [u8];

	fn deref(&self) -> &Self::Target {
		&self.0
	}
}

/// No length byte (or restriction), just all data to end of record. uses
/// hex encoding for text representation, whitespace allowed.
///
/// No following field allowed, i.e. last field in the record.
///
/// Must contain at least one byte
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct HexRemainingBlobNotEmpty(Bytes);

impl DnsPacketData for HexRemainingBlobNotEmpty {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
		let data = remaining_bytes(data);
		failure::ensure!(!data.is_empty(), "must not be empty");
		Ok(HexRemainingBlobNotEmpty(data))
	}

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

impl DnsTextData for HexRemainingBlobNotEmpty {
	fn dns_parse(_context: &DnsTextContext, data: &mut &str) -> crate::errors::Result<Self> {
		skip_whitespace(data);
		let result = HEXLOWER_PERMISSIVE_ALLOW_WS.decode(data.as_bytes())
			.with_context(|e| e.context(format!("invalid hex: {:?}", data)))?;
		*data = "";
		failure::ensure!(!result.is_empty(), "must not be empty");
		Ok(HexRemainingBlobNotEmpty(result.into()))
	}

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

impl std::ops::Deref for HexRemainingBlobNotEmpty {
	type Target = [u8];

	fn deref(&self) -> &Self::Target {
		&self.0
	}
}