dns: forward requests from the pod workload (istio#1282)

Fixes istio#1281

src/dns/forwarder.rs

@@ -13,24 +13,80 @@
 // limitations under the License.
 
 use crate::dns::resolver::{Answer, Resolver};
+use crate::proxy::SocketFactory;
+use hickory_proto::iocompat::AsyncIoTokioAsStd;
 use hickory_resolver::config::{ResolverConfig, ResolverOpts};
 use hickory_resolver::error::ResolveError;
-use hickory_resolver::name_server::TokioConnectionProvider;
-use hickory_resolver::TokioAsyncResolver;
+use hickory_resolver::name_server;
+use hickory_resolver::name_server::{GenericConnector, RuntimeProvider};
 use hickory_server::authority::LookupError;
 use hickory_server::server::Request;
+use std::future::Future;
+use std::net::SocketAddr;
+use std::pin::Pin;
+use std::sync::Arc;
+use tokio::net::{TcpStream, UdpSocket};
 
 /// A forwarding [Resolver] that delegates requests to an upstream [TokioAsyncResolver].
-pub struct Forwarder(TokioAsyncResolver);
+pub struct Forwarder(hickory_resolver::AsyncResolver<GenericConnector<RuntimeProviderAdaptor>>);
 
 impl Forwarder {
     /// Creates a new [Forwarder] from the provided resolver configuration.
-    pub fn new(cfg: ResolverConfig, opts: ResolverOpts) -> Result<Self, ResolveError> {
-        let resolver = TokioAsyncResolver::new(cfg, opts, TokioConnectionProvider::default());
+    pub fn new(
+        cfg: ResolverConfig,
+        socket_factory: Arc<dyn SocketFactory + Send + Sync>,
+        opts: ResolverOpts,
+    ) -> Result<Self, ResolveError> {
+        let provider = GenericConnector::new(RuntimeProviderAdaptor {
+            socket_factory,
+            handle: Default::default(),
+        });
+        let resolver = hickory_resolver::AsyncResolver::new(cfg, opts, provider);
         Ok(Self(resolver))
     }
 }
 
+#[derive(Clone)]
+struct RuntimeProviderAdaptor {
+    socket_factory: Arc<dyn SocketFactory + Send + Sync>,
+    handle: name_server::TokioHandle,
+}
+
+impl RuntimeProvider for RuntimeProviderAdaptor {
+    type Handle = name_server::TokioHandle;
+    type Timer = hickory_proto::TokioTime;
+    type Udp = UdpSocket;
+    type Tcp = AsyncIoTokioAsStd<TcpStream>;
+
+    fn create_handle(&self) -> Self::Handle {
+        self.handle.clone()
+    }
+
+    fn connect_tcp(
+        &self,
+        server_addr: SocketAddr,
+    ) -> Pin<Box<dyn Send + Future<Output = std::io::Result<Self::Tcp>>>> {
+        let sf = self.socket_factory.clone();
+        Box::pin(async move {
+            let socket = if server_addr.is_ipv4() {
+                sf.new_tcp_v4()
+            } else {
+                sf.new_tcp_v6()
+            }?;
+            socket.connect(server_addr).await.map(AsyncIoTokioAsStd)
+        })
+    }
+
+    fn bind_udp(
+        &self,
+        local_addr: SocketAddr,
+        _server_addr: SocketAddr,
+    ) -> Pin<Box<dyn Send + Future<Output = std::io::Result<Self::Udp>>>> {
+        let sf = self.socket_factory.clone();
+        Box::pin(async move { sf.udp_bind(local_addr) })
+    }
+}
+
 #[async_trait::async_trait]
 impl Resolver for Forwarder {
     async fn lookup(&self, request: &Request) -> Result<Answer, LookupError> {

src/dns/server.rs

@@ -742,14 +742,17 @@ pub trait Forwarder: Send + Sync {
 pub fn forwarder_for_mode(
     proxy_mode: ProxyMode,
     cluster_domain: String,
+    socket_factory: Arc<dyn SocketFactory + Send + Sync>,
 ) -> Result<Arc<dyn Forwarder>, Error> {
     Ok(match proxy_mode {
         ProxyMode::Shared => {
             // TODO(https://github.com/istio/ztunnel/issues/555): Use pod settings if available.
             // Today, we only support the basic namespace awareness
-            Arc::new(SystemForwarder::new(true, cluster_domain)?)
+            Arc::new(SystemForwarder::new(true, cluster_domain, socket_factory)?)
+        }
+        ProxyMode::Dedicated => {
+            Arc::new(SystemForwarder::new(false, cluster_domain, socket_factory)?)
         }
-        ProxyMode::Dedicated => Arc::new(SystemForwarder::new(false, cluster_domain)?),
     })
 }
 
@@ -768,7 +771,11 @@ enum SearchDomains {
 }
 
 impl SystemForwarder {
-    fn new(per_pod: bool, cluster_domain: String) -> Result<Self, Error> {
+    fn new(
+        per_pod: bool,
+        cluster_domain: String,
+        socket_factory: Arc<dyn SocketFactory + Send + Sync>,
+    ) -> Result<Self, Error> {
         // Get the resolver config from `ztunnel's` /etc/resolv.conf.
         let (cfg, opts) = read_system_conf().map_err(|e| Error::Generic(Box::new(e)))?;
 
@@ -780,6 +787,7 @@ impl SystemForwarder {
         Self::from_parts(
             per_pod,
             cluster_domain,
+            socket_factory,
             opts,
             domain,
             search_domains,
@@ -790,6 +798,7 @@ impl SystemForwarder {
     fn from_parts(
         per_pod: bool,
         cluster_domain: String,
+        socket_factory: Arc<dyn SocketFactory + Send + Sync>,
         opts: ResolverOpts,
         domain: Option<Name>,
         search_domains: Vec<Name>,
@@ -802,7 +811,8 @@ impl SystemForwarder {
 
         // Create the resolver.
         let resolver = Arc::new(
-            dns::forwarder::Forwarder::new(cfg, opts).map_err(|e| Error::Generic(Box::new(e)))?,
+            dns::forwarder::Forwarder::new(cfg, socket_factory, opts)
+                .map_err(|e| Error::Generic(Box::new(e)))?,
         );
         let search_domains = if per_pod {
             // Standard Kubernetes search is 'istio-system.svc.cluster.local svc.cluster.local cluster.local'
@@ -870,6 +880,7 @@ mod tests {
     use crate::xds::istio::workload::Workload as XdsWorkload;
     use crate::xds::istio::workload::{IpFamilies, NetworkAddress as XdsNetworkAddress};
 
+    use crate::proxy::DefaultSocketFactory;
     use crate::state::workload::{NetworkAddress, Workload};
     use crate::strng::Strng;
     use crate::{drain, strng};
@@ -1503,6 +1514,7 @@ mod tests {
         let f = SystemForwarder::from_parts(
             true,
             "cluster.local".to_string(),
+            Arc::new(DefaultSocketFactory),
             opts,
             None,
             search,

src/proxyfactory.rs

@@ -106,6 +106,7 @@ impl ProxyFactory {
                 dns::forwarder_for_mode(
                     self.config.proxy_mode,
                     self.config.cluster_domain.clone(),
+                    socket_factory.clone(),
                 )?,
                 self.dns_metrics.clone().unwrap(),
                 drain.clone(),

src/test_helpers/dns.rs

@@ -16,7 +16,7 @@ use crate::config::Address;
 use crate::dns::resolver::{Answer, Resolver};
 use crate::dns::Metrics;
 use crate::drain::DrainTrigger;
-use crate::proxy::Error;
+use crate::proxy::{DefaultSocketFactory, Error};
 use crate::state::workload::Workload;
 use crate::state::WorkloadInfo;
 use crate::test_helpers::new_proxy_state;
@@ -299,7 +299,8 @@ pub async fn run_dns(responses: HashMap<Name, Vec<IpAddr>>) -> anyhow::Result<Te
     let cfg = internal_resolver_config(tcp, udp);
     let opts = ResolverOpts::default();
     let resolver = Arc::new(
-        dns::forwarder::Forwarder::new(cfg, opts).map_err(|e| Error::Generic(Box::new(e)))?,
+        dns::forwarder::Forwarder::new(cfg, Arc::new(DefaultSocketFactory), opts)
+            .map_err(|e| Error::Generic(Box::new(e)))?,
     );
     Ok(TestDnsServer {
         tcp,