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

use bytes::Bytes;
use errors::*;
use ser::DnsPacketData;
use ser::text::{DnsTextData, DnsTextFormatter, next_field};
use std::fmt;
use std::io::Cursor;
use std::borrow::Cow;

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct Class(pub u16);

pub const IN   : Class = Class(KnownClass::IN    as u16);
pub const CH   : Class = Class(KnownClass::CH    as u16);
pub const HS   : Class = Class(KnownClass::HS    as u16);
pub const NONE : Class = Class(KnownQClass::NONE as u16);
pub const ANY  : Class = Class(KnownQClass::ANY  as u16);

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
#[repr(u16)]
#[allow(non_camel_case_types)]
pub enum KnownClass {
	// try to list "original" rfc
	IN = 0x0001, // RFC 1035
	// CS = 0x0002, // "CSNET"
	CH = 0x0003, // "Chaos"
	HS = 0x0004, // "Hesiod"
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
#[repr(u16)]
#[allow(non_camel_case_types)]
pub enum KnownQClass {
	// try to list "original" rfc
	NONE = 0x00fe, // RFC 2136
	ANY  = 0x00ff, // RFC 1035: "*"
}

impl Class {
	pub fn name(self) -> Cow<'static, str> {
		Cow::Borrowed(match self {
			IN => "IN",
			CH => "CH",
			HS => "HS",
			NONE => "NONE",
			ANY => "ANY",
			_ => return Cow::Owned(self.generic_name()),
		})
	}

	/// uses generic name for QCLASS values (`is_qclass`)
	pub fn class_name(self) -> Cow<'static, str> {
		if self.is_qclass() {
			Cow::Owned(self.generic_name())
		} else {
			self.name()
		}
	}

	/// QCLASS names can overlap with (Q)TYPE names
	///
	/// classes unknown to this implementation never count as QCLASS,
	/// but they also are only represented using the generic "CLASS..."
	/// names, which don't overlap with (Q)TYPE names.
	pub fn is_qclass(self) -> bool {
		match self {
			NONE => true,
			ANY => true,
			_ => false,
		}
	}

	pub fn from_name(name: &str) -> Option<Self> {
		use std::ascii::AsciiExt;
		if let Some(n) = Self::class_from_name(name) { return Some(n); }
		// explicit QCLASS names
		if name.eq_ignore_ascii_case("NONE") { return Some(NONE); }
		if name.eq_ignore_ascii_case("ANY") { return Some(ANY); }
		None
	}

	/// similar to `from_name`, but doesn't accept QCLASS names (it
	/// always accepts "CLASS..." names though, even if they are known to
	/// be of type QCLASS)
	pub fn class_from_name(name: &str) -> Option<Self> {
		use std::ascii::AsciiExt;
		if name.eq_ignore_ascii_case("IN") { return Some(IN); }
		if name.eq_ignore_ascii_case("CH") { return Some(CH); }
		if name.eq_ignore_ascii_case("HS") { return Some(HS); }
		if name.as_bytes()[0..5].eq_ignore_ascii_case(b"CLASS") {
			if let Ok(c) = name[5..].parse::<u16>() {
				return Some(Class(c))
			}
		}
		None
	}

	pub fn write_class_name(self, f: &mut fmt::Formatter) -> fmt::Result {
		if self.is_qclass() {
			self.write_generic_name(f)
		} else {
			write!(f, "{}", self)
		}
	}

	pub fn write_generic_name(self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "CLASS{}", self.0)
	}

	pub fn generic_name(self) -> String {
		format!("CLASS{}", self.0)
	}
}

impl From<KnownClass> for Class {
	fn from(value: KnownClass) -> Self {
		Class(value as u16)
	}
}

impl From<KnownQClass> for Class {
	fn from(value: KnownQClass) -> Self {
		Class(value as u16)
	}
}

impl fmt::Display for Class {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		let n = match *self {
			IN => "IN",
			CH => "CH",
			HS => "HS",
			NONE => "NONE",
			ANY => "ANY",
			_ => {
				return self.write_generic_name(f)
			},
		};
		write!(f, "{}", n)
	}
}

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

impl DnsTextData for Class {
	fn dns_parse(data: &mut &str) -> Result<Self> {
		let field = next_field(data)?;
		Class::from_name(field).ok_or_else(|| format_err!("unknown class {:?}", field))
	}

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