use crate::protocol::name::DnsName;
use super::zone::Zone;
use std::collections::HashMap;
#[derive(Debug, Default)]
pub struct ZoneTree {
zones: HashMap<DnsName, Zone>,
}
impl ZoneTree {
pub fn new() -> Self {
Self {
zones: HashMap::new(),
}
}
pub fn insert(&mut self, zone: Zone) {
self.zones.insert(zone.origin.clone(), zone);
}
pub fn remove(&mut self, origin: &DnsName) -> Option<Zone> {
self.zones.remove(origin)
}
pub fn find_zone(&self, name: &DnsName) -> Option<&Zone> {
let labels = name.labels();
for i in 0..labels.len() {
let candidate = match DnsName::from_labels(&labels[i..]) {
Ok(name) => name,
Err(_) => continue,
};
if let Some(zone) = self.zones.get(&candidate) {
return Some(zone);
}
}
self.zones.get(&DnsName::root())
}
pub fn get_zone(&self, origin: &DnsName) -> Option<&Zone> {
self.zones.get(origin)
}
pub fn get_zone_mut(&mut self, origin: &DnsName) -> Option<&mut Zone> {
self.zones.get_mut(origin)
}
pub fn zone_names(&self) -> Vec<&DnsName> {
self.zones.keys().collect()
}
pub fn len(&self) -> usize {
self.zones.len()
}
pub fn is_empty(&self) -> bool {
self.zones.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::protocol::rdata::SoaData;
fn make_zone(origin: &str) -> Zone {
let origin = DnsName::from_str(origin).unwrap();
Zone::new(
origin.clone(),
SoaData {
mname: DnsName::from_str("ns1.example.com").unwrap(),
rname: DnsName::from_str("admin.example.com").unwrap(),
serial: 1,
refresh: 3600,
retry: 900,
expire: 604800,
minimum: 300,
},
3600,
)
}
#[test]
fn test_zone_tree_find() {
let mut tree = ZoneTree::new();
tree.insert(make_zone("example.com"));
tree.insert(make_zone("sub.example.com"));
let zone = tree.find_zone(&DnsName::from_str("www.sub.example.com").unwrap());
assert!(zone.is_some());
assert_eq!(zone.unwrap().origin.to_dotted(), "sub.example.com.");
let zone = tree.find_zone(&DnsName::from_str("www.example.com").unwrap());
assert!(zone.is_some());
assert_eq!(zone.unwrap().origin.to_dotted(), "example.com.");
let zone = tree.find_zone(&DnsName::from_str("www.other.com").unwrap());
assert!(zone.is_none());
}
#[test]
fn test_zone_tree_insert_remove() {
let mut tree = ZoneTree::new();
tree.insert(make_zone("example.com"));
assert_eq!(tree.len(), 1);
tree.remove(&DnsName::from_str("example.com").unwrap());
assert_eq!(tree.len(), 0);
}
}