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>>(&mut self, label: L) -> Result<()> { let label = label.into(); if label.len() > 254 - self.total_len { failure::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>>(&mut self, label: L) -> Result<()> { let label = label.into(); if label.len() > 254 - self.total_len { failure::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)); if data.len() < new_len { let add = new_len - data.len(); data.reserve(add); } 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 } } }