diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs index ac44dfe6..8f17cac3 100644 --- a/easytier/src/dns/server.rs +++ b/easytier/src/dns/server.rs @@ -19,29 +19,28 @@ use hickory_server::{ server::{Request, RequestHandler, ResponseHandler, ResponseInfo}, ServerFuture, }; -use itertools::Itertools; -use parking_lot::Mutex; +use parking_lot::{Mutex, RwLock}; use pnet::packet::icmp::{IcmpTypes, MutableIcmpPacket}; use pnet::packet::ip::IpNextHeaderProtocols; use pnet::packet::ipv4::{Ipv4Packet, MutableIpv4Packet}; use pnet::packet::udp::{MutableUdpPacket, UdpPacket}; use pnet::packet::{icmp, ipv4, udp, MutablePacket, Packet}; use std::collections::HashSet; +use std::io; use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4}; -use std::{io, iter}; use std::{sync::Arc, time::Duration}; -use tokio::{sync::RwLock, task::JoinHandle}; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; #[derive(Clone)] pub struct DynamicCatalog { - inner: Arc>, + inner: Arc>, } impl DynamicCatalog { pub fn new() -> Self { Self { - inner: Arc::new(RwLock::new(Catalog::new())), + inner: Arc::new(tokio::sync::RwLock::new(Catalog::new())), } } @@ -65,43 +64,6 @@ impl RequestHandler for DynamicCatalog { } } -struct DnsServerRuntime { - token: CancellationToken, - task: Option>, -} - -impl DnsServerRuntime { - fn start(mut server: ServerFuture) -> Self { - Self { - token: server.shutdown_token().clone(), - task: Some(tokio::spawn(async move { - server - .block_until_done() - .await - .unwrap_or_else(|e| tracing::error!("DNS server exited with error: {:?}", e)); - })), - } - } - - async fn stop(mut self) -> anyhow::Result<()> { - self.token.cancel(); - if let Some(task) = self.task.take() { - task.await?; - } - Ok(()) - } -} - -impl Drop for DnsServerRuntime { - fn drop(&mut self) { - self.token.cancel(); - if let Some(task) = self.task.take() { - task.abort(); - tracing::warn!("DNS server runtime is leaked"); - } - } -} - // ResponseWrapper for serializing DNS responses into a byte buffer. // Used by the address hijacking NIC packet filter to produce DNS replies in-place. #[derive(Debug, Clone, From, Into, Deref, DerefMut)] @@ -148,6 +110,43 @@ impl ResponseHandler for Response { } } +struct DnsServerRuntime { + token: CancellationToken, + task: Option>, +} + +impl DnsServerRuntime { + fn start(mut server: ServerFuture) -> Self { + Self { + token: server.shutdown_token().clone(), + task: Some(tokio::spawn(async move { + server + .block_until_done() + .await + .unwrap_or_else(|e| tracing::error!("DNS server exited with error: {:?}", e)); + })), + } + } + + async fn stop(mut self) -> anyhow::Result<()> { + self.token.cancel(); + if let Some(task) = self.task.take() { + task.await?; + } + Ok(()) + } +} + +impl Drop for DnsServerRuntime { + fn drop(&mut self) { + self.token.cancel(); + if let Some(task) = self.task.take() { + task.abort(); + tracing::warn!("DNS server runtime is leaked"); + } + } +} + #[derive(Derivative)] #[derivative(Debug)] pub struct DnsServer { @@ -178,20 +177,8 @@ impl DnsServer { } } - async fn reload_addresses(&self, addresses: impl IntoIterator) { - let addresses = addresses.into_iter().collect::>(); - let mut current = self.addresses.write().await; // TODO: read? - - if *current == addresses { - return; - } - - let added = addresses.difference(&*current).cloned().collect_vec(); - let removed = current.difference(&addresses).cloned().collect_vec(); - - *current = addresses; - - // TODO + pub fn routes(&self) -> HashSet { + self.addresses.read().iter().map(|a| a.addr.ip()).collect() } async fn reload_listeners( @@ -239,7 +226,7 @@ impl DnsServer { loop { dirty.addresses.notified().await; if dirty.addresses.reset() { - self.reload_addresses(self.mgr.iter_addresses()).await; + *self.addresses.write() = self.mgr.iter_addresses().collect(); } tokio::time::sleep(Duration::from_secs(1)).await; } @@ -267,7 +254,7 @@ impl DnsServer { _ = reload_listeners => {}, ); - self.reload_addresses(iter::empty()).await; + self.addresses.write().clear(); if let Some(runtime) = runtime.take() { let _ = runtime.stop().await; } @@ -295,19 +282,15 @@ impl NicPacketFilter for DnsServer { } impl DnsServer { - async fn is_hijacked_ip(&self, ip: &IpAddr) -> bool { - self.addresses - .read() - .await - .iter() - .any(|a| a.addr.ip() == *ip) + fn is_hijacked_ip(&self, ip: &IpAddr) -> bool { + self.addresses.read().iter().any(|a| a.addr.ip() == *ip) } - async fn is_hijacked_addr(&self, addr: &NameServerAddr) -> bool { - self.addresses.read().await.contains(addr) + fn is_hijacked_addr(&self, addr: &NameServerAddr) -> bool { + self.addresses.read().contains(addr) } - /// Replace content of incoming UDP DNS request and ICMP echo request packet with reply data, + /// Replace the content of an incoming UDP DNS request and ICMP echo request packet with reply data, /// and swap source and destination IP addresses to send it back. async fn handle_ip_packet(&self, zc_packet: &mut ZCPacket) -> Option<()> { let (ip_header_length, ip_protocol, src_ip, dst_ip) = { @@ -325,7 +308,7 @@ impl DnsServer { ) }; - if !self.is_hijacked_ip(&dst_ip.into()).await { + if !self.is_hijacked_ip(&dst_ip.into()) { return None; } @@ -383,10 +366,7 @@ impl DnsServer { ) }; - if !self - .is_hijacked_addr(&SocketAddr::new(dst_ip.into(), dst_port).into()) - .await - { + if !self.is_hijacked_addr(&SocketAddr::new(dst_ip.into(), dst_port).into()) { return None; }