use std::{
collections::HashMap,
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
str::FromStr,
sync::Arc,
time::Duration,
};
use async_trait::async_trait;
use ipnetwork::IpNetwork;
use rand::prelude::{IteratorRandom, SliceRandom};
use tracing::info;
use trust_dns_resolver::config::{NameServerConfig, ResolverConfig, ResolverOpts};
use zeronsd::{
addresses::Calculator,
authority::{find_members, RecordAuthority, ZTAuthority},
server::Server,
traits::{ToHostname, ToPointerSOA},
utils::{authtoken_path, domain_or_default, get_listen_ips, parse_ip_from_cidr},
};
use self::{
context::TestContext,
network::TestNetwork,
resolver::{Lookup, Resolver, Resolvers},
utils::{format_hosts_file, HostsType},
};
pub mod context;
pub mod member;
pub mod network;
pub mod resolver;
pub mod to_ip;
pub mod utils;
pub struct ServiceConfig {
hosts: HostsType,
update_interval: Option<Duration>,
ips: Option<Vec<&'static str>>,
wildcard_everything: bool,
network_filename: Option<&'static str>,
}
impl Default for ServiceConfig {
fn default() -> Self {
Self {
network_filename: None,
hosts: HostsType::None,
update_interval: None,
ips: None,
wildcard_everything: false,
}
}
}
impl ServiceConfig {
pub fn network_filename(mut self, n: &'static str) -> Self {
self.network_filename = Some(n);
self
}
pub fn hosts(mut self, h: HostsType) -> Self {
self.hosts = h;
self
}
pub fn update_interval(mut self, u: Option<Duration>) -> Self {
self.update_interval = u;
self
}
pub fn ips(mut self, ips: Option<Vec<&'static str>>) -> Self {
self.ips = ips;
self
}
pub fn wildcard_everything(mut self, w: bool) -> Self {
self.wildcard_everything = w;
self
}
}
#[derive(Clone)]
pub struct Service {
tn: Arc<TestNetwork>,
resolvers: Resolvers,
update_interval: Duration,
pub listen_ips: Vec<SocketAddr>,
}
impl Service {
pub async fn new(sc: ServiceConfig) -> Self {
let network_filename = sc.network_filename.unwrap_or("basic-ipv4");
let tn = if let Some(ips) = sc.ips {
TestNetwork::new_multi_ip(network_filename, &mut TestContext::default().await, ips)
.await
.unwrap()
} else {
TestNetwork::new(network_filename, &mut TestContext::default().await)
.await
.unwrap()
};
let mut update_interval = Duration::new(4, 0);
if let Some(interval) = sc.update_interval {
update_interval += interval
}
let listen_ips =
Self::create_listeners(&tn, sc.hosts, update_interval, sc.wildcard_everything).await;
Self {
tn: Arc::new(tn),
resolvers: Self::create_resolvers(listen_ips.clone()),
listen_ips,
update_interval,
}
}
fn create_resolvers(sockets: Vec<SocketAddr>) -> Resolvers {
let mut resolvers = Vec::new();
for socket in sockets {
let mut resolver_config = ResolverConfig::new();
resolver_config.add_search(domain_or_default(None).unwrap());
resolver_config.add_name_server(NameServerConfig {
bind_addr: None,
socket_addr: socket,
protocol: trust_dns_resolver::config::Protocol::Udp,
tls_dns_name: None,
trust_nx_responses: true,
});
let mut opts = ResolverOpts::default();
opts.attempts = 10;
opts.cache_size = 0;
opts.rotate = true;
opts.use_hosts_file = false;
opts.positive_min_ttl = Some(Duration::new(0, 0));
opts.positive_max_ttl = Some(Duration::new(0, 0));
opts.negative_min_ttl = Some(Duration::new(0, 0));
opts.negative_max_ttl = Some(Duration::new(0, 0));
resolvers.push(Arc::new(
trust_dns_resolver::TokioAsyncResolver::tokio(resolver_config, opts).unwrap(),
));
}
resolvers
}
async fn create_listeners(
tn: &TestNetwork,
hosts: HostsType,
update_interval: Duration,
wildcard_everything: bool,
) -> Vec<SocketAddr> {
let listen_cidrs = get_listen_ips(&authtoken_path(None), &tn.network.clone().id.unwrap())
.await
.unwrap();
let mut listen_ips = Vec::new();
let mut ipmap = HashMap::new();
let mut authority_map = HashMap::new();
for cidr in listen_cidrs.clone() {
let listen_ip = parse_ip_from_cidr(cidr.clone());
let socket_addr = SocketAddr::new(listen_ip.clone(), 53);
listen_ips.push(socket_addr);
let cidr = IpNetwork::from_str(&cidr.clone()).unwrap();
if !ipmap.contains_key(&listen_ip) {
ipmap.insert(listen_ip, cidr.network());
}
if !authority_map.contains_key(&cidr) {
let ptr_authority = RecordAuthority::new(
cidr.to_ptr_soa_name().unwrap(),
cidr.to_ptr_soa_name().unwrap(),
)
.await
.unwrap();
authority_map.insert(cidr, ptr_authority.clone());
}
}
if let Some(v6assign) = tn.network.config.clone().unwrap().v6_assign_mode {
if v6assign.rfc4193.unwrap_or(false) {
let cidr = tn.network.clone().rfc4193().unwrap();
if !authority_map.contains_key(&cidr) {
let ptr_authority = RecordAuthority::new(
cidr.to_ptr_soa_name().unwrap(),
cidr.to_ptr_soa_name().unwrap(),
)
.await
.unwrap();
authority_map.insert(cidr, ptr_authority);
}
}
}
let authority = RecordAuthority::new(
domain_or_default(None).unwrap().into(),
tn.member()
.to_fqdn(domain_or_default(None).unwrap().into())
.unwrap()
.into(),
)
.await
.unwrap();
let ztauthority = ZTAuthority {
network_id: tn.network.clone().id.unwrap(),
client: tn.central(),
hosts_file: format_hosts_file(hosts),
reverse_authority_map: authority_map,
update_interval,
forward_authority: authority.clone(),
wildcard: wildcard_everything,
hosts: None,
};
tokio::spawn(find_members(ztauthority.clone()));
tokio::time::sleep(update_interval).await;
for ip in listen_ips.clone() {
let server = Server::new(ztauthority.to_owned());
info!("Serving {}", ip.clone());
tokio::spawn(server.listen(ip.ip(), Duration::new(1, 0), None, None, None));
}
listen_ips
}
pub fn any_listen_ip(self) -> IpAddr {
self.listen_ips
.clone()
.into_iter()
.choose(&mut rand::thread_rng())
.unwrap()
.clone()
.ip()
}
pub fn network(&self) -> Arc<TestNetwork> {
self.tn.clone()
}
pub fn resolvers(&self) -> Resolvers {
self.resolvers.clone()
}
pub fn any_resolver(&self) -> Arc<Resolver> {
self.resolvers()
.choose(&mut rand::thread_rng())
.to_owned()
.unwrap()
.clone()
}
pub fn member_record(&self) -> String {
format!("zt-{}.home.arpa.", self.network().identity().clone())
}
pub async fn change_name(&self, name: &'static str) {
let mut member = self
.network()
.central()
.get_network_member(
&self.network().network.clone().id.unwrap(),
&self.network().identity(),
)
.await
.unwrap();
member.name = Some(name.to_string());
let _ = &self
.network()
.central()
.update_network_member(
&self.network().network.clone().id.unwrap(),
&self.network().identity(),
&member.to_owned(),
)
.await
.unwrap();
tokio::time::sleep(self.update_interval).await; }
pub fn test_network(&self) -> Arc<TestNetwork> {
self.tn.clone()
}
}
#[async_trait]
impl Lookup for Service {
async fn lookup_a(&self, record: String) -> Vec<Ipv4Addr> {
self.any_resolver().lookup_a(record).await
}
async fn lookup_aaaa(&self, record: String) -> Vec<Ipv6Addr> {
self.any_resolver().lookup_aaaa(record).await
}
async fn lookup_ptr(&self, record: String) -> Vec<String> {
self.any_resolver().lookup_ptr(record).await
}
}