diff --git a/easytier/src/dns/node.rs b/easytier/src/dns/node.rs index 77d04f09..1cffc697 100644 --- a/easytier/src/dns/node.rs +++ b/easytier/src/dns/node.rs @@ -12,12 +12,10 @@ use crate::proto::rpc_impl::standalone::{StandAloneClient, StandAloneServer}; use crate::proto::rpc_types::controller::BaseController; use crate::tunnel::tcp::{TcpTunnelConnector, TcpTunnelListener}; use crate::utils::AsyncRuntime; -use derivative::Derivative; -use futures::task::SpawnExt; use std::sync::Arc; use std::time::Duration; -use tokio::sync::{broadcast, Mutex, Notify}; -use tokio::task::{JoinHandle, JoinSet}; +use tokio::sync::{broadcast, Notify}; +use tokio::task::JoinSet; use tokio::time::{sleep, sleep_until, Instant}; use tokio_util::sync::CancellationToken; use uuid::Uuid; @@ -108,11 +106,12 @@ impl DnsNode { let server = Arc::new(DnsServer::new( self.peer_mgr.clone(), self.global_ctx.clone(), - rpc, #[cfg(feature = "tun")] self.nic_ctx.clone(), )); + server.register(&rpc); + self.global_ctx.set_dns(Some(server.clone())); tokio::join!( self.peer_mgr diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs index 4e3ab4f5..b6484ef1 100644 --- a/easytier/src/dns/server.rs +++ b/easytier/src/dns/server.rs @@ -140,16 +140,10 @@ impl DnsServer { pub fn new( peer_mgr: Arc, global_ctx: ArcGlobalCtx, - rpc: StandAloneServer, #[cfg(feature = "tun")] nic_ctx: ArcNicCtx, // TODO: REMOVE THIS ) -> Self { - let mgr = Arc::new(DnsNodeMgr::new()); - - rpc.registry() - .register(DnsNodeMgrRpcServer::new_arc(mgr.clone()), ""); - Self { - mgr, + mgr: Arc::new(DnsNodeMgr::new()), nic_ctx, peer_mgr, global_ctx, @@ -158,6 +152,11 @@ impl DnsServer { } } + pub fn register(&self, rpc: &StandAloneServer) { + rpc.registry() + .register(DnsNodeMgrRpcServer::new_arc(self.mgr.clone()), ""); + } + pub fn addresses(&self) -> HashSet { self.addresses.read().iter().map(|a| a.addr).collect() }