diff --git a/src/context.rs b/src/context.rs index 8c1fb02d..3fcf7ad4 100644 --- a/src/context.rs +++ b/src/context.rs @@ -24,6 +24,8 @@ use trust_dns_resolver::TokioAsyncResolver; use crate::relay::dns_resolver::create_resolver; #[cfg(feature = "local-flow-stat")] use crate::relay::flow::ServerFlowStatistic; +#[cfg(feature = "local-dns-relay")] +use crate::relay::dnsrelay::upstream::LocalUpstream; use crate::{ acl::AccessControl, config::{Config, ConfigType, ServerConfig}, @@ -174,6 +176,10 @@ pub struct Context { // For DNS relay's ACL domain name reverse lookup -- whether the IP shall be forwarded #[cfg(feature = "local-dns-relay")] reverse_lookup_cache: Mutex>, + + // For local DNS upstream + #[cfg(feature = "local-dns-relay")] + local_dns: LocalUpstream, } /// Unique context thw whole server @@ -187,6 +193,8 @@ impl Context { let reverse_lookup_cache = Mutex::new(LruCache::::with_expiry_duration(Duration::from_secs( 3 * 24 * 60 * 60, ))); + #[cfg(feature = "local-dns-relay")] + let local_dns = LocalUpstream::new(&config); Context { config, @@ -197,6 +205,8 @@ impl Context { local_flow_statistic: ServerFlowStatistic::new(), #[cfg(feature = "local-dns-relay")] reverse_lookup_cache, + #[cfg(feature = "local-dns-relay")] + local_dns, } } @@ -321,6 +331,10 @@ impl Context { &self.config.acl } + pub fn local_dns(&self) -> &LocalUpstream { + &self.local_dns + } + /// Check target address ACL (for client) pub async fn check_target_bypassed(&self, target: &Address) -> bool { match self.acl() { diff --git a/src/relay/dnsrelay/mod.rs b/src/relay/dnsrelay/mod.rs index c4ccee3e..543f93cf 100644 --- a/src/relay/dnsrelay/mod.rs +++ b/src/relay/dnsrelay/mod.rs @@ -24,7 +24,7 @@ use crate::{ }, }; -mod upstream; +pub mod upstream; fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool { let mut iter = name.iter().rev(); @@ -129,14 +129,13 @@ fn should_forward_by_response( } } -async fn acl_lookup( +async fn acl_lookup( acl: &Option, - local: Arc, + local: &upstream::LocalUpstream, remote: Arc, query: &Query ) -> (io::Result, bool) where - Local: upstream::Upstream, Remote: upstream::Upstream, { // Start querying name servers @@ -233,14 +232,6 @@ pub async fn run(context: SharedContext) -> io::Result<()> { }); let config = context.config(); - #[cfg(target_os = "android")] - let local_upstream = Arc::new(upstream::UnixSocketUpstream { - path: config.local_dns_path.clone().expect("local query DNS path"), - }); - #[cfg(not(target_os = "android"))] - let local_upstream = Arc::new(upstream::UdpUpstream { - server: config.local_dns_addr.clone().expect("local query DNS address"), - }); // FIXME: We use TCP to send remote queries by default, which should be configuable. let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await; let remote_upstream = Arc::new(upstream::ProxyTcpUpstream { @@ -270,8 +261,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { debug!("received src: {}, query: {:?}", src, request); let context = context.clone(); - let local_upstream = Arc::clone(&local_upstream); - let remote_upstream = Arc::clone(&remote_upstream); + let remote_upstream = remote_upstream.clone(); let mut qtx = qtx.clone(); tokio::spawn(async move { @@ -286,7 +276,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> { message.set_response_code(ResponseCode::NotImp); } else if request.query_count() > 0 { let question = &request.queries()[0]; - let (r, forward) = acl_lookup(context.acl(), local_upstream, remote_upstream, question).await; + let (r, forward) = acl_lookup(context.acl(), context.local_dns(), remote_upstream, question).await; if let Ok(result) = r { for rec in result.answers() { diff --git a/src/relay/dnsrelay/upstream.rs b/src/relay/dnsrelay/upstream.rs index daf6046d..c5b0e24e 100644 --- a/src/relay/dnsrelay/upstream.rs +++ b/src/relay/dnsrelay/upstream.rs @@ -20,7 +20,7 @@ use std::path::PathBuf; use tokio::net::UnixStream; use crate::{ - config::ServerConfig, + config::{Config, ServerConfig}, context::SharedContext, relay::{ socks5::Address, @@ -28,6 +28,36 @@ use crate::{ }, }; +#[derive(Debug)] +pub enum LocalUpstream { + Udp(UdpUpstream), + // Tcp(TcpUpstream), + #[cfg(unix)] + UnixSocket(UnixSocketUpstream), +} + +impl LocalUpstream { + pub fn new(config: &Config) -> LocalUpstream { + #[cfg(target_os = "android")] + return LocalUpstream::UnixSocket(UnixSocketUpstream { + path: config.local_dns_path.clone().expect("local query DNS path"), + }); + #[cfg(not(target_os = "android"))] + LocalUpstream::Udp(UdpUpstream { + server: config.local_dns_addr.clone().expect("local query DNS address"), + }) + } + + pub async fn lookup(&self, query: &Query) -> io::Result { + match self { + LocalUpstream::Udp(upstream) => upstream.lookup(query).await, + // LocalUpstream::Tcp(upstream) => upstream.lookup(query), + #[cfg(unix)] + LocalUpstream::UnixSocket(upstream) => upstream.lookup(query).await, + } + } +} + #[async_trait] pub trait Upstream: Debug { async fn lookup(&self, query: &Query) -> io::Result;