266 lines
7.7 KiB
Rust
266 lines
7.7 KiB
Rust
|
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));
|
||
|
data.put_u8(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
|
||
|
}
|
||
|
}
|
||
|
}
|