Maestro: add in-process runtime generation reload

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/transport/middle_proxy/http_fetch.rs

@@ -66,33 +66,13 @@ fn extract_host_port_path(url: &str) -> Result<(String, u16, String)> {
     Ok((host, port, path))
 }
 
-async fn resolve_target_addr(host: &str, port: u16) -> Result<std::net::SocketAddr> {
-    if let Some(addr) = resolve_socket_addr(host, port) {
-        return Ok(addr);
-    }
-
-    let addrs: Vec<std::net::SocketAddr> = tokio::net::lookup_host((host, port))
-        .await
-        .map_err(|e| ProxyError::Proxy(format!("DNS resolve failed for {host}:{port}: {e}")))?
-        .collect();
-
-    if let Some(addr) = addrs.iter().copied().find(|addr| addr.is_ipv4()) {
-        return Ok(addr);
-    }
-
-    addrs
-        .first()
-        .copied()
-        .ok_or_else(|| ProxyError::Proxy(format!("DNS returned no addresses for {host}:{port}")))
-}
-
 async fn connect_https_transport(
     host: &str,
     port: u16,
     upstream: Option<Arc<UpstreamManager>>,
 ) -> Result<UpstreamStream> {
     if let Some(manager) = upstream {
-        let target = resolve_target_addr(host, port).await?;
+        let target = manager.resolve_hostname(host, port).await?;
         return timeout(HTTP_CONNECT_TIMEOUT, manager.connect(target, None, None))
             .await
             .map_err(|_| ProxyError::Proxy(format!("upstream connect timeout for {host}:{port}")))?

src/transport/upstream.rs

@@ -4,6 +4,7 @@
 
 #![allow(deprecated)]
 
+use arc_swap::ArcSwap;
 use rand::RngExt;
 use std::collections::{BTreeSet, HashMap};
 use std::net::{IpAddr, SocketAddr};
@@ -20,7 +21,7 @@ use tracing::{debug, info, trace, warn};
 
 use crate::config::{UpstreamConfig, UpstreamType};
 use crate::error::{ProxyError, Result};
-use crate::network::dns_overrides::{resolve_socket_addr, split_host_port};
+use crate::network::dns_overrides::{DnsOverrides, split_host_port};
 use crate::protocol::constants::{TG_DATACENTER_PORT, TG_DATACENTERS_V4, TG_DATACENTERS_V6};
 use crate::stats::Stats;
 use crate::transport::shadowsocks::{
@@ -333,6 +334,7 @@ pub struct UpstreamManager {
     no_upstreams_warn_epoch_ms: Arc<AtomicU64>,
     no_healthy_warn_epoch_ms: Arc<AtomicU64>,
     stats: Arc<Stats>,
+    dns_overrides: Arc<ArcSwap<DnsOverrides>>,
 }
 
 impl UpstreamManager {
@@ -374,9 +376,37 @@ impl UpstreamManager {
             no_upstreams_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
             no_healthy_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
             stats,
+            dns_overrides: Arc::new(ArcSwap::from_pointee(DnsOverrides::default())),
         }
     }
 
+    pub(crate) fn with_dns_overrides(mut self, entries: &[String]) -> Result<Self> {
+        self.dns_overrides = Arc::new(ArcSwap::from_pointee(DnsOverrides::from_entries(entries)?));
+        Ok(self)
+    }
+
+    pub(crate) fn update_dns_overrides(&self, entries: &[String]) -> Result<()> {
+        let snapshot = DnsOverrides::from_entries(entries)?;
+        self.dns_overrides.store(Arc::new(snapshot));
+        Ok(())
+    }
+
+    pub(crate) async fn resolve_hostname(&self, host: &str, port: u16) -> Result<SocketAddr> {
+        if let Some(addr) = self.dns_overrides.load().resolve_socket_addr(host, port) {
+            return Ok(addr);
+        }
+        let addrs: Vec<SocketAddr> = tokio::net::lookup_host((host, port))
+            .await
+            .map_err(ProxyError::Io)?
+            .collect();
+        if let Some(addr) = addrs.iter().copied().find(SocketAddr::is_ipv4) {
+            return Ok(addr);
+        }
+        addrs.first().copied().ok_or_else(|| {
+            ProxyError::Proxy(format!("DNS returned no addresses for {host}:{port}"))
+        })
+    }
+
     fn now_epoch_ms() -> u64 {
         std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
@@ -709,11 +739,12 @@ impl UpstreamManager {
     }
 
     async fn connect_hostname_with_dns_override(
+        &self,
         address: &str,
         connect_timeout: Duration,
     ) -> Result<TcpStream> {
         if let Some((host, port)) = split_host_port(address)
-            && let Some(addr) = resolve_socket_addr(&host, port)
+            && let Some(addr) = self.dns_overrides.load().resolve_socket_addr(&host, port)
         {
             return match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
                 Ok(Ok(stream)) => Ok(stream),
@@ -1203,7 +1234,8 @@ impl UpstreamManager {
                             "SOCKS4 interface binding is not supported for hostname addresses, ignoring"
                         );
                     }
-                    Self::connect_hostname_with_dns_override(address, connect_timeout).await?
+                    self.connect_hostname_with_dns_override(address, connect_timeout)
+                        .await?
                 };
 
                 // replace socks user_id with config.selected_scope, if set
@@ -1291,7 +1323,8 @@ impl UpstreamManager {
                             "SOCKS5 interface binding is not supported for hostname addresses, ignoring"
                         );
                     }
-                    Self::connect_hostname_with_dns_override(address, connect_timeout).await?
+                    self.connect_hostname_with_dns_override(address, connect_timeout)
+                        .await?
                 };
 
                 debug!(config = ?config, "Socks5 connection");