From dc1e1929c86d23cc87cc22ca5cd6cf10a30f565f Mon Sep 17 00:00:00 2001 From: Luna Yao <40349250+ZnqbuZ@users.noreply.github.com> Date: Sat, 21 Feb 2026 22:38:36 +0100 Subject: [PATCH] client_mgr: add format client_mgr --- easytier/src/dns/client.rs | 10 +- easytier/src/dns/client_mgr.rs | 165 ++++++++++++++++++++++ easytier/src/dns/mod.rs | 1 + easytier/src/dns/server.rs | 245 +++++++++++++-------------------- easytier/src/proto/dns.proto | 2 +- 5 files changed, 270 insertions(+), 153 deletions(-) create mode 100644 easytier/src/dns/client_mgr.rs diff --git a/easytier/src/dns/client.rs b/easytier/src/dns/client.rs index 31fe3f43..7dc44d97 100644 --- a/easytier/src/dns/client.rs +++ b/easytier/src/dns/client.rs @@ -1,7 +1,7 @@ use crate::dns::config::DNS_SERVER_RPC_ADDR; use crate::dns::peer_mgr::DnsPeerMgr; use crate::peers::peer_manager::PeerManager; -use crate::proto::dns::{DnsPeerMgrRpcServer, DnsServerRpcClientFactory, HeartbeatRequest}; +use crate::proto::dns::{DnsClientMgrRpcClientFactory, DnsPeerMgrRpcServer, HeartbeatRequest}; use crate::proto::peer_rpc::RoutePeerInfo; use crate::proto::rpc_impl::standalone::StandAloneClient; use crate::proto::rpc_types::controller::BaseController; @@ -63,22 +63,22 @@ impl DnsClient { ) -> anyhow::Result<()> { let request = if heartbeat.snapshot.is_none() || self.mgr.dirty.peers.reset() { heartbeat.update(self.mgr.snapshot()); - heartbeat.clone().into() + heartbeat.clone() } else { let snapshot = heartbeat.snapshot.take(); - let request = heartbeat.clone().into(); + let request = heartbeat.clone(); heartbeat.snapshot = snapshot; request }; let client = rpc - .scoped_client::>("".to_string()) + .scoped_client::>("".to_string()) .await?; let response = client.heartbeat(BaseController::default(), request).await?; if response.resync { client - .heartbeat(BaseController::default(), heartbeat.clone().into()) + .heartbeat(BaseController::default(), heartbeat.clone()) .await?; } diff --git a/easytier/src/dns/client_mgr.rs b/easytier/src/dns/client_mgr.rs new file mode 100644 index 00000000..085609d8 --- /dev/null +++ b/easytier/src/dns/client_mgr.rs @@ -0,0 +1,165 @@ +use crate::dns::utils::{DirtyFlag, NameServerAddr}; +use crate::dns::zone::{Zone, ZoneGroup}; +use crate::proto::dns::DnsClientMgrRpc; +use crate::proto::dns::{DnsSnapshot, HeartbeatRequest, HeartbeatResponse}; +use crate::proto::rpc_types; +use crate::proto::rpc_types::controller::BaseController; +use crate::utils::{DeterministicDigest, MapTryInto}; +use anyhow::Error; +use derive_more::{Deref, DerefMut}; +use hickory_server::authority::Catalog; +use itertools::Itertools; +use moka::future::Cache; +use std::collections::HashSet; +use std::time::Duration; +use tokio::sync::Notify; +use uuid::Uuid; + +#[derive(Debug, Clone, Default)] +pub struct DnsClientInfo { + digest: Vec, + zones: ZoneGroup, + addresses: HashSet, + listeners: HashSet, +} + +impl TryFrom<&DnsSnapshot> for DnsClientInfo { + type Error = Error; + + fn try_from(value: &DnsSnapshot) -> Result { + Ok(Self { + digest: value.digest(), + zones: (&value.zones).try_into()?, + addresses: value.addresses.iter().map_try_into().try_collect()?, + listeners: value.listeners.iter().map_try_into().try_collect()?, + }) + } +} + +const DNS_CLIENT_TTL: Duration = Duration::from_secs(5); + +// TODO: same as DnsPeerMgrDirtyState +#[derive(Debug, Default, Deref, DerefMut)] +pub struct DnsClientMgrDirtyState { + pub(super) zones: DirtyFlag, + pub(super) addresses: DirtyFlag, + pub(super) listeners: DirtyFlag, + #[deref] + #[deref_mut] + notify: Notify, +} + +#[derive(Debug)] +pub struct DnsClientMgr { + clients: Cache, + pub(super) dirty: DnsClientMgrDirtyState, +} + +impl DnsClientMgr { + pub fn new() -> Self { + Self { + clients: Cache::builder().time_to_live(DNS_CLIENT_TTL).build(), + dirty: Default::default(), + } + } + + pub fn catalog(&self) -> Catalog { + let zones = self.collect_zones(); + let mut catalog = Catalog::new(); + + for zone in zones.iter() { + catalog.upsert( + zone.origin.clone(), + zone.create_memory_authority().into_iter().collect(), + ); + } + + for zone in zones.iter() { + catalog.upsert( + zone.origin.clone(), + zone.create_forward_authority().into_iter().collect(), + ); + } + + catalog + } + + pub fn collect_zones(&self) -> ZoneGroup { + let mut zones = vec![Zone::system()]; + let mut local = HashSet::::new(); + + for (_, info) in self.clients.iter() { + zones.extend(info.zones); + local.extend(info.addresses); + local.extend(info.listeners); + } + + for zone in zones.iter_mut() { + if let Some(forward) = zone.forward.as_mut() { + forward + .name_servers + .retain(|ns| !local.contains(&ns.clone().into())); + } + } + + zones.into() + } + + pub fn iter_addresses(&self) -> impl Iterator + use<'_> { + self.clients + .iter() + .flat_map(|(_, info)| info.addresses) + .unique() + } + + pub fn iter_listeners(&self) -> impl Iterator + use<'_> { + self.clients + .iter() + .flat_map(|(_, info)| info.listeners) + .unique() + } +} + +#[async_trait::async_trait] +impl DnsClientMgrRpc for DnsClientMgr { + type Controller = BaseController; + + async fn heartbeat( + &self, + _: BaseController, + input: HeartbeatRequest, + ) -> rpc_types::error::Result { + let id = input + .id + .ok_or(anyhow::anyhow!( + "missing id in heartbeat request: {:?}", + input + ))? + .into(); + + let resync = if let Some(snapshot) = input.snapshot.as_ref() { + let new = DnsClientInfo::try_from(snapshot)?; + let old = self.clients.get(&id).await.unwrap_or_default(); + if new.digest != old.digest { + self.dirty.zones.mark(); + if new.addresses != old.addresses { + self.dirty.addresses.mark(); + } + if new.listeners != old.listeners { + self.dirty.listeners.mark(); + } + + self.clients.insert(id, new).await; + self.dirty.notify_one(); + } + false + } else { + self.clients + .get(&id) + .await + .is_none_or(|info| info.digest != input.digest) + }; + + Ok(HeartbeatResponse { resync }) + } +} diff --git a/easytier/src/dns/mod.rs b/easytier/src/dns/mod.rs index a35ca209..b7451832 100644 --- a/easytier/src/dns/mod.rs +++ b/easytier/src/dns/mod.rs @@ -1,4 +1,5 @@ mod client; +mod client_mgr; pub mod config; mod peer_mgr; mod server; diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs index 1e75811c..a9a6c798 100644 --- a/easytier/src/dns/server.rs +++ b/easytier/src/dns/server.rs @@ -1,52 +1,25 @@ use super::{utils::NameServerAddr, zone::Zone}; -use crate::dns::utils::DirtyFlag; -use crate::dns::zone::ZoneGroup; -use crate::proto::dns::DnsSnapshot; -use crate::proto::rpc_types; -use crate::proto::{ - dns::{DnsServerRpc, HeartbeatRequest, HeartbeatResponse}, - rpc_types::controller::BaseController, -}; -use crate::utils::{DeterministicDigest, MapTryInto}; -use anyhow::Error; +use crate::common::PeerId; +use crate::dns::client_mgr::DnsClientMgr; +use cidr::Ipv4Inet; +use derivative::Derivative; +use derive_more::{Deref, DerefMut, From, Into}; +use hickory_proto::rr::Record; +use hickory_proto::serialize::binary::BinEncoder; use hickory_proto::xfer::Protocol; use hickory_server::{ - authority::Catalog, + authority::{Catalog, MessageResponse}, server::{Request, RequestHandler, ResponseHandler, ResponseInfo}, ServerFuture, }; use itertools::Itertools; -use moka::future::Cache; +use parking_lot::Mutex; use std::collections::HashSet; +use std::io; use std::{sync::Arc, time::Duration}; -use derivative::Derivative; -use derive_more::{Deref, DerefMut}; use tokio::net::{TcpListener, UdpSocket}; use tokio::{sync::RwLock, task::JoinHandle}; -use tokio::sync::Notify; use tokio_util::sync::CancellationToken; -use uuid::Uuid; - -#[derive(Debug, Clone, Default)] -pub struct DnsClientInfo { - digest: Vec, - zones: ZoneGroup, - addresses: HashSet, - listeners: HashSet, -} - -impl TryFrom<&DnsSnapshot> for DnsClientInfo { - type Error = Error; - - fn try_from(value: &DnsSnapshot) -> Result { - Ok(Self { - digest: value.digest(), - zones: (&value.zones).try_into()?, - addresses: value.addresses.iter().map_try_into().try_collect()?, - listeners: value.listeners.iter().map_try_into().try_collect()?, - }) - } -} #[derive(Clone)] pub struct DynamicCatalog { @@ -80,17 +53,6 @@ impl RequestHandler for DynamicCatalog { } } -// TODO: same as DnsPeerMgrDirtyState -#[derive(Debug, Default, Deref, DerefMut)] -pub struct DnsServerDirtyState { - zones: DirtyFlag, - addresses: DirtyFlag, - listeners: DirtyFlag, - #[deref] - #[deref_mut] - notify: Notify, -} - struct DnsServerRuntime { token: CancellationToken, task: JoinHandle<()>, @@ -116,70 +78,102 @@ impl DnsServerRuntime { } } +// 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)] +struct Response(Arc>>); + +impl Response { + pub fn new(capacity: usize) -> Self { + Self(Arc::new(Mutex::new(Vec::with_capacity(capacity)))) + } + + pub fn into_inner(self) -> Option> { + Arc::into_inner(self.0).map(Mutex::into_inner) + } +} + +trait RecordIter<'r>: Iterator + Send + 'r {} +impl<'r, T> RecordIter<'r> for T where T: Iterator + Send + 'r {} + +#[async_trait::async_trait] +impl ResponseHandler for Response { + async fn send_response<'r>( + &mut self, + response: MessageResponse< + '_, + 'r, + impl RecordIter<'r>, + impl RecordIter<'r>, + impl RecordIter<'r>, + impl RecordIter<'r>, + >, + ) -> io::Result { + let max_size = if let Some(edns) = response.get_edns() { + edns.max_payload() + } else { + hickory_proto::udp::MAX_RECEIVE_BUFFER_SIZE as u16 + }; + + let mut this = self.lock(); + let mut encoder = BinEncoder::new(this.as_mut()); + encoder.set_max_size(max_size); + response + .destructive_emit(&mut encoder) + .map_err(io::Error::other) + } +} #[derive(Derivative)] #[derivative(Debug)] pub struct DnsServer { - clients: Cache, - dirty: DnsServerDirtyState, + mgr: DnsClientMgr, #[derivative(Debug = "ignore")] catalog: DynamicCatalog, + + /// Current set of hijacked addresses (only UDP protocol addresses). + hijacked: RwLock>, + + /// Tun device name, needed for adding/removing routes. + tun_dev: RwLock>, + + /// Tun device IP inet, used to check if an address is within the tun subnet. + tun_inet: RwLock>, + + /// Our peer ID, used to set the to_peer_id on response packets. + my_peer_id: RwLock, } -const DNS_CLIENT_TTL: Duration = Duration::from_secs(5); const DNS_SERVER_LISTENER_TCP_TIMEOUT: Duration = Duration::from_secs(5); impl DnsServer { async fn reload_zones(&self) { - let mut zones = vec![Zone::system()]; - let mut local = HashSet::::new(); - - for (_, info) in self.clients.iter() { - zones.extend(info.zones.iter().cloned()); - local.extend(info.addresses.iter()); - local.extend(info.listeners.iter()); - } - - let mut catalog = Catalog::new(); - - for zone in zones.iter_mut() { - if let Some(forward) = zone.forward.as_mut() { - forward - .name_servers - .retain(|ns| !local.contains(&ns.clone().into())); - } - } - - for zone in zones.iter() { - catalog.upsert( - zone.origin.clone(), - zone.create_memory_authority().into_iter().collect(), - ); - } - - for zone in zones.iter() { - catalog.upsert( - zone.origin.clone(), - zone.create_forward_authority().into_iter().collect(), - ); - } - - self.catalog.replace(catalog).await; + self.catalog.replace(self.mgr.catalog()).await; } - async fn reload_addresses(&self) { - todo!() + async fn reload_addresses(&self, addresses: impl IntoIterator) { + let addresses = addresses.into_iter().collect::>(); + + let mut active = self.hijacked.write().await; + + let added = addresses.difference(&active).cloned().collect_vec(); + let removed = active.difference(&addresses).cloned().collect_vec(); + + if added.is_empty() && removed.is_empty() { + return; + } + + *active = addresses; + + // TODO } - async fn reload_listeners(&self, runtime: &mut Option) -> anyhow::Result<()> { - let listeners = self - .clients - .iter() - .map(|(_, info)| info.listeners.into_iter()) - .flatten() - .collect_vec(); - + async fn reload_listeners( + &self, + listeners: impl IntoIterator, + runtime: &mut Option, + ) -> anyhow::Result<()> { if let Some(old) = runtime.take() { old.stop().await?; } @@ -209,24 +203,27 @@ impl DnsServer { } pub async fn run(&self) { - let dirty = &self.dirty; + let dirty = &self.mgr.dirty; let mut runtime = None; loop { dirty.notified().await; if dirty.zones.reset() { - self.reload_zones().await; + self.reload_zones(&self.mgr.collect_zones()).await; } if dirty.addresses.reset() { - self.reload_addresses().await; + self.reload_addresses(self.mgr.iter_addresses()).await; } if dirty.listeners.reset() { - if let Err(e) = self.reload_listeners(&mut runtime).await { + if let Err(e) = self + .reload_listeners(self.mgr.iter_listeners(), &mut runtime) + .await + { tracing::error!("failed to reload listeners: {:?}", e); - self.dirty.listeners.mark(); - self.dirty.notify_one(); + dirty.listeners.mark(); + dirty.notify_one(); } } @@ -234,49 +231,3 @@ impl DnsServer { } } } - -#[async_trait::async_trait] -impl DnsServerRpc for DnsServer { - type Controller = BaseController; - - async fn heartbeat( - &self, - _: BaseController, - input: HeartbeatRequest, - ) -> rpc_types::error::Result { - let id = input - .id - .ok_or(anyhow::anyhow!( - "missing id in heartbeat request: {:?}", - input - ))? - .into(); - - let resync = if let Some(snapshot) = input.snapshot.as_ref() { - let new = DnsClientInfo::try_from(snapshot)?; - let old = self.clients.get(&id).await.unwrap_or_default(); - if new.digest != old.digest { - if new.zones != old.zones { - self.dirty.zones.mark(); - } - if new.addresses != old.addresses { - self.dirty.addresses.mark(); - } - if new.listeners != old.listeners { - self.dirty.listeners.mark(); - } - - self.clients.insert(id, new).await; - self.dirty.notify_one(); - } - false - } else { - self.clients - .get(&id) - .await - .is_none_or(|info| info.digest != input.digest) - }; - - Ok(HeartbeatResponse { resync }) - } -} diff --git a/easytier/src/proto/dns.proto b/easytier/src/proto/dns.proto index fa7029a7..1487e86c 100644 --- a/easytier/src/proto/dns.proto +++ b/easytier/src/proto/dns.proto @@ -38,6 +38,6 @@ message HeartbeatResponse { bool resync = 1; } -service DnsServerRpc { +service DnsClientMgrRpc { rpc Heartbeat(HeartbeatRequest) returns (HeartbeatResponse) {} }