Move LocalUpstream to global config

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<LruCache<IpAddr, bool>>,
+
+    // 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::<IpAddr, bool>::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() {

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<Local, Remote>(
+async fn acl_lookup<Remote>(
     acl: &Option<AccessControl>,
-    local: Arc<Local>,
+    local: &upstream::LocalUpstream,
     remote: Arc<Remote>,
     query: &Query
 ) -> (io::Result<Message>, 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() {

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<Message> {
+        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<Message>;