#![deny(missing_docs)]
//! Various structs to represents DNS names and labels

use bytes::Bytes;
use errors::*;
use ser::DnsPacketData;
use smallvec::SmallVec;
use std::fmt;
use std::io::Cursor;
use std::ops::{Deref, DerefMut};

pub use self::display::*;
pub use self::label::*;

mod display;
mod label;
mod name_mutations;
mod name_packet_parser;
mod name_text_parser;

#[derive(Clone,Copy,PartialEq,Eq,PartialOrd,Ord,Hash,Debug)]
enum LabelOffset {
	LabelStart(u8),
	PacketStart(u16),
}

// the heap meta data is usually at least 2*usize big; assuming 64-bit
// platforms it should be ok to use 16 bytes in the smallvec.
#[derive(Clone,PartialEq,Eq,PartialOrd,Ord,Hash,Debug)]
enum LabelOffsets {
	Uncompressed(SmallVec<[u8;16]>),
	Compressed(usize, SmallVec<[LabelOffset;8]>),
}

impl LabelOffsets {
	fn len(&self) -> u8 {
		let l = match *self {
			LabelOffsets::Uncompressed(ref offs) => offs.len(),
			LabelOffsets::Compressed(_, ref offs) => offs.len(),
		};
		debug_assert!(l < 128);
		l as u8
	}

	fn label_pos(&self, ndx: u8) -> usize {
		debug_assert!(ndx < 127);
		match *self {
			LabelOffsets::Uncompressed(ref offs) => offs[ndx as usize] as usize,
			LabelOffsets::Compressed(start, ref offs) => match offs[ndx as usize] {
				LabelOffset::LabelStart(o) => start + (o as usize),
				LabelOffset::PacketStart(o) => o as usize,
			}
		}
	}

	fn is_compressed(&self) -> bool {
		match *self {
			LabelOffsets::Uncompressed(_) => false,
			LabelOffsets::Compressed(_, _) => true,
		}
	}
}

/// A DNS name
/// 
/// Uses the "original" raw representation for storage (i.e. can share
/// memory with a parsed packet)
#[derive(Clone,Hash)]
pub struct DnsName {
	// in uncompressed form always includes terminating null octect;
	// but even in uncompressed form can include unused bytes at the
	// beginning
	//
	// may be empty for the root name (".", no labels)
	data: Bytes,
	// either uncompressed or compressed offsets
	label_offsets: LabelOffsets,
	// length of encoded form
	total_len: u8,
}

/// names that should be written in canonical form for DNSSEC according
/// to https://tools.ietf.org/html/rfc4034#section-6.2
///
/// TODO: make it a newtype.
/// 
/// DnsCompressedName always needs to be written in canonical form for
/// DNSSEC.
pub type DnsCanonicalName = DnsName;

impl DnsName {
	/// Create new name representing the DNS root (".")
	pub fn new_root() -> Self {
		DnsName{
			data: Bytes::new(),
			label_offsets: LabelOffsets::Uncompressed(SmallVec::new()),
			total_len: 1,
		}
	}

	/// Create new name representing the DNS root (".") and pre-allocate
	/// storage
	pub fn with_capacity(labels: u8, total_len: u8) -> Self {
		DnsName{
			data: Bytes::with_capacity(total_len as usize),
			label_offsets: LabelOffsets::Uncompressed(SmallVec::with_capacity(labels as usize)),
			total_len: 1,
		}
	}

	/// Returns whether name represents the DNS root (".")
	pub fn is_root(&self) -> bool {
		0 == self.label_count()
	}

	/// How many labels the name has (without the trailing empty label,
	/// at most 127)
	pub fn label_count(&self) -> u8 {
		self.label_offsets.len()
	}

	/// Iterator over the labels (in the order they are stored in memory,
	/// i.e. top-level name last).
	pub fn labels<'a>(&'a self) -> DnsNameIterator<'a> {
		DnsNameIterator{
			name: &self,
			front_label: 0,
			back_label: self.label_offsets.len(),
		}
	}

	/// Return label at index `ndx`
	///
	/// # Panics
	///
	/// panics if `ndx >= self.label_count()`.
	pub fn label_ref<'a>(&'a self, ndx: u8) -> DnsLabelRef<'a> {
		let pos = self.label_offsets.label_pos(ndx);
		let label_len = self.data[pos];
		debug_assert!(label_len < 64);
		let end = pos + 1 + label_len as usize;
		DnsLabelRef{label: &self.data[pos + 1..end]}
	}

	/// Return label at index `ndx`
	///
	/// # Panics
	///
	/// panics if `ndx >= self.label_count()`.
	pub fn label(&self, ndx: u8) -> DnsLabel {
		let pos = self.label_offsets.label_pos(ndx);
		let label_len = self.data[pos];
		debug_assert!(label_len < 64);
		let end = pos + 1 + label_len as usize;
		DnsLabel{label: self.data.slice(pos + 1, end) }
	}
}

impl<'a> IntoIterator for &'a DnsName {
	type Item = DnsLabelRef<'a>;
	type IntoIter = DnsNameIterator<'a>;

	fn into_iter(self) -> Self::IntoIter {
		self.labels()
	}
}

impl PartialEq<DnsName> for DnsName {
	fn eq(&self, rhs: &DnsName) -> bool {
		let a_labels = self.labels();
		let b_labels = rhs.labels();
		if a_labels.len() != b_labels.len() { return false; }
		a_labels.zip(b_labels).all(|(a,b)| a == b)
	}
}

impl Eq for DnsName{}

impl fmt::Debug for DnsName {
	fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
		DisplayLabels{
			labels: self,
			options: Default::default(),
		}.fmt(w)
	}
}

impl fmt::Display for DnsName {
	fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
		DisplayLabels{
			labels: self,
			options: Default::default(),
		}.fmt(w)
	}
}

impl DnsPacketData for DnsName {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
		DnsName::parse_name(data, false)
	}
}

/// Similar to `DnsName`, but allows using compressed labels in the
/// serialized form
#[derive(Clone)]
pub struct DnsCompressedName(pub DnsName);

impl Deref for DnsCompressedName {
	type Target = DnsName;

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

impl DerefMut for DnsCompressedName {
	fn deref_mut(&mut self) -> &mut Self::Target {
		&mut self.0
	}
}

impl<'a> IntoIterator for &'a DnsCompressedName {
	type Item = DnsLabelRef<'a>;
	type IntoIter = DnsNameIterator<'a>;

	fn into_iter(self) -> Self::IntoIter {
		self.labels()
	}
}

impl PartialEq<DnsCompressedName> for DnsCompressedName {
	fn eq(&self, rhs: &DnsCompressedName) -> bool {
		self.0 == rhs.0
	}
}

impl PartialEq<DnsName> for DnsCompressedName {
	fn eq(&self, rhs: &DnsName) -> bool {
		&self.0 == rhs
	}
}

impl PartialEq<DnsCompressedName> for DnsName {
	fn eq(&self, rhs: &DnsCompressedName) -> bool {
		self == &rhs.0
	}
}

impl Eq for DnsCompressedName{}

impl fmt::Debug for DnsCompressedName {
	fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
		self.0.fmt(w)
	}
}

impl fmt::Display for DnsCompressedName {
	fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
		self.0.fmt(w)
	}
}

impl DnsPacketData for DnsCompressedName {
	fn deserialize(data: &mut Cursor<Bytes>) -> Result<Self> {
		Ok(DnsCompressedName(DnsName::parse_name(data, true)?))
	}
}

/// Iterator type for [`DnsName::labels`]
///
/// [`DnsName::labels`]: struct.DnsName.html#method.labels
#[derive(Clone)]
pub struct DnsNameIterator<'a> {
	name: &'a DnsName,
	front_label: u8,
	back_label: u8,
}

impl<'a> Iterator for DnsNameIterator<'a> {
	type Item = DnsLabelRef<'a>;

	fn next(&mut self) -> Option<Self::Item> {
		if self.front_label >= self.back_label { return None }
		let label = self.name.label_ref(self.front_label);
		self.front_label += 1;
		Some(label)
	}

	fn size_hint(&self) -> (usize, Option<usize>) {
		let count = self.len();
		(count, Some(count))
	}

	fn count(self) -> usize {
		self.len()
	}
}

impl<'a> ExactSizeIterator for DnsNameIterator<'a> {
	fn len(&self) -> usize {
		(self.back_label - self.front_label) as usize
	}
}

impl<'a> DoubleEndedIterator for DnsNameIterator<'a> {
	fn next_back(&mut self) -> Option<Self::Item> {
		if self.front_label >= self.back_label { return None }
		self.back_label -= 1;
		let label = self.name.label_ref(self.back_label);
		Some(label)
	}
}

#[cfg(test)]
mod tests {
	use ser::packet;
	use super::*;

/*
	fn deserialize(bytes: &'static [u8]) -> Result<DnsName> {
		let result = packet::deserialize_with(Bytes::from_static(bytes), DnsName::deserialize)?;
		{
			let check_result = packet::deserialize_with(result.clone().encode(), DnsName::deserialize).unwrap();
			assert_eq!(check_result, result);
		}
		Ok(result)
	}
*/

	fn de_uncompressed(bytes: &'static [u8]) -> Result<DnsName> {
		let result = packet::deserialize_with(Bytes::from_static(bytes), DnsName::deserialize)?;
		assert_eq!(bytes, result.clone().encode());
		Ok(result)
	}

	fn check_uncompressed_display(bytes: &'static [u8], txt: &str, label_count: u8) {
		let name = de_uncompressed(bytes).unwrap();
		assert_eq!(
			name.labels().count(),
			label_count as usize
		);
		assert_eq!(
			format!("{}", name),
			txt
		);
	}

	fn check_uncompressed_debug(bytes: &'static [u8], txt: &str) {
		let name = de_uncompressed(bytes).unwrap();
		assert_eq!(
			format!("{:?}", name),
			txt
		);
	}

	#[test]
	fn parse_and_display_name() {
		check_uncompressed_display(
			b"\x07example\x03com\x00",
			"example.com.",
			2,
		);
		check_uncompressed_display(
			b"\x07e!am.l\\\x03com\x00",
			"e\\033am\\.l\\\\.com.",
			2,
		);
		check_uncompressed_debug(
			b"\x07e!am.l\\\x03com\x00",
			r#""e\\033am\\.l\\\\.com.""#,
		);
	}

	#[test]
	fn parse_and_reverse_name() {
		let name = de_uncompressed(b"\x03www\x07example\x03com\x00").unwrap();
		assert_eq!(
			format!(
				"{}",
				DisplayLabels{
					labels: name.labels().rev(),
					options: DisplayLabelsOptions{
						separator: " ",
						trailing: false,
					},
				}
			),
			"com example www"
		);
	}

	#[test]
	fn modifications() {
		let mut name = de_uncompressed(b"\x07example\x03com\x00").unwrap();

		name.push_front(DnsLabelRef::new(b"www").unwrap()).unwrap();
		assert_eq!(
			format!("{}", name),
			"www.example.com."
		);

		name.push_back(DnsLabelRef::new(b"org").unwrap()).unwrap();
		assert_eq!(
			format!("{}", name),
			"www.example.com.org."
		);

		name.pop_front();
		assert_eq!(
			format!("{}", name),
			"example.com.org."
		);

		name.push_front(DnsLabelRef::new(b"mx").unwrap()).unwrap();
		assert_eq!(
			format!("{}", name),
			"mx.example.com.org."
		);
		// the "mx" label should fit into the place "www" used before,
		// make sure the buffer was reused and the name not moved within
		assert_eq!(1, name.label_offsets.label_pos(0));

		name.push_back(DnsLabelRef::new(b"com").unwrap()).unwrap();
		assert_eq!(
			format!("{}", name),
			"mx.example.com.org.com."
		);
	}



	fn de_compressed(bytes: &'static [u8], offset: usize) -> Result<DnsCompressedName> {
		use bytes::Buf;

		let mut c = Cursor::new(Bytes::from_static(bytes));
		c.set_position(offset as u64);
		let result = DnsPacketData::deserialize(&mut c)?;
		if c.remaining() != 0 {
			bail!("data remaining: {}", c.remaining())
		}
		Ok(result)
	}

	fn check_compressed_display(bytes: &'static [u8], offset: usize, txt: &str, label_count: u8) {
		let name = de_compressed(bytes, offset).unwrap();
		assert_eq!(
			name.labels().count(),
			label_count as usize
		);
		assert_eq!(
			format!("{}", name),
			txt
		);
	}

	fn check_compressed_debug(bytes: &'static [u8], offset: usize, txt: &str) {
		let name = de_compressed(bytes, offset).unwrap();
		assert_eq!(
			format!("{:?}", name),
			txt
		);
	}

	#[test]
	fn parse_invalid_compressed_name() {
		de_compressed(b"\x11com\x00\x07example\xc0\x00", 5).unwrap_err();
		de_compressed(b"\x10com\x00\x07example\xc0\x00", 5).unwrap_err();
	}

	#[test]
	fn parse_and_display_compressed_name() {
		check_compressed_display(
			b"\x03com\x00\x07example\xc0\x00", 5,
			"example.com.",
			2,
		);
		check_compressed_display(
			b"\x03com\x00\x07e!am.l\\\xc0\x00", 5,
			"e\\033am\\.l\\\\.com.",
			2,
		);
		check_compressed_debug(
			b"\x03com\x00\x07e!am.l\\\xc0\x00", 5,
			r#""e\\033am\\.l\\\\.com.""#,
		);
		check_compressed_display(
			b"\x03com\x00\x07example\xc0\x00\x03www\xc0\x05", 15,
			"www.example.com.",
			3,
		);
	}

	#[test]
	fn modifications_compressed() {
		let mut name = de_compressed(b"\x03com\x00\x07example\xc0\x00\xc0\x05", 15).unwrap();

		name.push_front(DnsLabelRef::new(b"www").unwrap()).unwrap();
		assert_eq!(
			format!("{}", name),
			"www.example.com."
		);

		name.push_back(DnsLabelRef::new(b"org").unwrap()).unwrap();
		assert_eq!(
			format!("{}", name),
			"www.example.com.org."
		);

		name.pop_front();
		assert_eq!(
			format!("{}", name),
			"example.com.org."
		);

		name.push_front(DnsLabelRef::new(b"mx").unwrap()).unwrap();
		assert_eq!(
			format!("{}", name),
			"mx.example.com.org."
		);
		// the "mx" label should fit into the place "www" used before,
		// make sure the buffer was reused and the name not moved within
		assert_eq!(1, name.label_offsets.label_pos(0));

		name.push_back(DnsLabelRef::new(b"com").unwrap()).unwrap();
		assert_eq!(
			format!("{}", name),
			"mx.example.com.org.com."
		);
	}
}