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

use bytes::{BytesMut,BufMut};
use super::*;

impl DnsName {
	/// Remove the front label
	///
	/// Returns `false` if the name was the root (".")
	pub fn try_pop_front(&mut self) -> bool {
		match self.label_offsets {
			LabelOffsets::Uncompressed(ref mut offs) => {
				if offs.is_empty() { return false; }
				self.total_len -= self.data[offs[0] as usize] + 1;
				offs.remove(0);
			},
			LabelOffsets::Compressed(ref mut start_pos, ref mut offs) => {
				if offs.is_empty() { return false; }
				match offs[0] {
					LabelOffset::LabelStart(o) => {
						let label_space = self.data[*start_pos + o as usize] + 1;
						self.total_len -= label_space;
						*start_pos += label_space as usize;
					},
					LabelOffset::PacketStart(o) => {
						self.total_len -= self.data[o as usize] + 1;
					},
				}
				offs.remove(0);
			},
		}
		true
	}

	/// Remove the front label
	///
	/// # Panics
	///
	/// Panics if the name was the root (".")
	pub fn pop_front(&mut self) {
		if !self.try_pop_front() { panic!("Cannot pop label from root name") }
	}

	/// Insert a new label at the front
	///
	/// Returns an error if the resulting name would be too long
	pub fn push_front<'a, L: Into<DnsLabelRef<'a>>>(&mut self, label: L) -> Result<()> {
		let label = label.into();
		if label.len() > 254 - self.total_len { bail!("Cannot append label, resulting name too long") }
		let (mut data, start) = self.reserve(label.len() as usize + 1, 0);

		let new_label_pos = start - (label.len() + 1) as usize;
		data[new_label_pos] = label.len();
		data[new_label_pos+1..new_label_pos+1+label.len() as usize].copy_from_slice(label.as_raw());
		self.data = data.freeze();
		self.total_len += label.len() + 1;
		match self.label_offsets {
			LabelOffsets::Uncompressed(ref mut offs) => {
				offs.insert(0, new_label_pos as u8);
			},
			LabelOffsets::Compressed(_, _) => unreachable!(),
		}

		Ok(())
	}

	/// Remove the back label
	///
	/// Returns `false` if the name was the root (".")
	pub fn try_pop_back(&mut self) -> bool {
		match self.label_offsets {
			LabelOffsets::Uncompressed(ref mut offs) => {
				if offs.is_empty() { return false; }
				self.total_len -= self.data[offs[offs.len()-1] as usize] + 1;
				offs.pop();
			},
			LabelOffsets::Compressed(ref mut start_pos, ref mut offs) => {
				if offs.is_empty() { return false; }
				match offs[offs.len()-1] {
					LabelOffset::LabelStart(o) => {
						self.total_len -= self.data[*start_pos + o as usize] + 1;
					},
					LabelOffset::PacketStart(o) => {
						self.total_len -= self.data[o as usize] + 1;
					},
				}
				offs.pop();
			},
		}
		true
	}

	/// Remove the back label
	///
	/// # Panics
	///
	/// Panics if the name was the root (".")
	pub fn pop_back(&mut self) {
		if !self.try_pop_back() { panic!("Cannot pop label from root name") }
	}

	/// Insert a new label at the back
	///
	/// Returns an error if the resulting name would be too long
	pub fn push_back<'a, L: Into<DnsLabelRef<'a>>>(&mut self, label: L) -> Result<()> {
		let label = label.into();
		if label.len() > 254 - self.total_len { bail!("Cannot append label, resulting name too long") }
		let (mut data, start) = self.reserve(0, label.len() as usize + 1);

		let new_label_pos = start + self.total_len as usize - 1;
		data[new_label_pos] = label.len();
		data[new_label_pos+1..new_label_pos+1+label.len() as usize].copy_from_slice(label.as_raw());
		data[new_label_pos+1+label.len() as usize] = 0;
		self.data = data.freeze();
		self.total_len += label.len() + 1;
		match self.label_offsets {
			LabelOffsets::Uncompressed(ref mut offs) => {
				offs.push(new_label_pos as u8);
			},
			LabelOffsets::Compressed(_, _) => unreachable!(),
		}

		Ok(())
	}

	// returns mutable buffer and position of the start of the current
	// name (null terminated).
	//
	// adjusts self.label_offsets accordingly
	fn reserve(&mut self, prefix: usize, suffix: usize) -> (BytesMut, usize) {
		use std::ptr;

		let new_len = self.total_len as usize + prefix + suffix;
		assert!(new_len < 256);
		self.uncompress(prefix, suffix);
		let label_offsets = match self.label_offsets {
			LabelOffsets::Uncompressed(ref mut offs) => offs,
			LabelOffsets::Compressed(_, _) => unreachable!(),
		};
		// steal buffer from self (so it has a change to be owned)
		let data = self.data.split_off(0).try_mut();

		if label_offsets.is_empty() {
			// root name
			let mut data = data.unwrap_or_else(|_| BytesMut::with_capacity(new_len));
			unsafe { data.set_len(new_len); }
			data[0] = 0;
			return (data, 0)
		}

		let old_start = label_offsets[0] as usize;
		// if current "prefix" space (old_start) is bigger than
		// requested but fits, just increase the prefix
		let (prefix, new_len) = if old_start > prefix && self.total_len as usize + old_start + suffix < 256 {
			(old_start, self.total_len as usize + old_start + suffix)
		} else {
			(prefix, new_len)
		};

		// check if we need to reallocate
		let data = match data {
			Ok(data) => {
				if data.capacity() < new_len {
					// need a new allocation anyway, pretend we couldn't
					// own the buffer
					Err(data.freeze())
				} else {
					Ok(data)
				}
			},
			Err(data) => Err(data),
		};


		match data {
			Ok(mut data) => {
				if data.len() < new_len {
					let add = new_len - data.len();
					data.reserve(add);
					unsafe { data.set_len(new_len); }
				}
				if old_start < prefix {
					// need more space in front, move back
					let p_old = &data[old_start as usize] as *const u8;
					let p_new = &mut data[prefix] as *mut u8;
					unsafe {
						ptr::copy(p_old, p_new, self.total_len as usize);
					}
					// adjust labels
					for o in label_offsets.iter_mut() {
						*o = *o + (prefix - old_start) as u8;
					}
					return (data, prefix);
				} else if old_start > prefix {
					// too much space in front, (suffix didn't fit into
					// length restriction, see check above), move to
					// front
					let p_old = &data[old_start as usize] as *const u8;
					let p_new = &mut data[prefix] as *mut u8;
					unsafe {
						ptr::copy(p_old, p_new, self.total_len as usize);
					}
					// adjust labels
					for o in label_offsets.iter_mut() {
						*o = *o - (old_start - prefix) as u8;
					}
					return (data, prefix);
				} else {
					return (data, old_start);
				}
			},
			Err(data) => {
				let mut new_data = BytesMut::with_capacity(new_len);
				unsafe { new_data.set_len(new_len); }
				// copy old data
				new_data[prefix..prefix + self.total_len as usize].copy_from_slice(&data[old_start..old_start+self.total_len as usize]);
				// adjust labels
				for o in label_offsets.iter_mut() {
					*o = (*o - old_start as u8) + prefix as u8;
				}
				return (new_data, prefix);
			},
		}
	}

	/// return encoded (uncompressed) representation
	///
	/// might uncompress the inner representation.
	pub fn encode(&mut self) -> Bytes {
		self.uncompress(0, 0);
		if self.is_root() {
			return Bytes::from_static(b"\x00");
		}
		// at least one label, find first one
		let from = self.label_offsets.label_pos(0);
		self.data.slice(from, from + self.total_len as usize)
	}

	fn uncompress(&mut self, prefix_capacity: usize, suffix_capacity: usize) {
		if self.label_offsets.is_compressed() {
			let name = {
				let labels = self.labels();
				let label_count = labels.len();
				let new_len = self.total_len as usize + prefix_capacity + suffix_capacity;
				assert!(new_len < 256);
				let mut data = BytesMut::with_capacity(new_len);
				let mut offsets = SmallVec::<[u8;16]>::with_capacity(label_count);
				unsafe { data.set_len(prefix_capacity) }
				let mut pos = prefix_capacity as u8;

				for label in labels {
					offsets.push(pos);
					pos += label.len() + 1;
					data.put_u8(label.len());
					data.put_slice(label.as_raw());
				}
				data.put_u8(0);

				DnsName{
					data: data.freeze(),
					label_offsets: LabelOffsets::Uncompressed(offsets),
					total_len: self.total_len,
				}
			};
			*self = name
		}
	}
}