enable more sni tests mzcompose

Author: jubrad

Cargo.lock

@@ -19,6 +19,7 @@ chrono = { version = "0.4.39", default-features = false, features = ["std"] }
 clap = { version = "4.5.23", features = ["derive", "env"] }
 domain = { version = "0.11.1", default-features = false, features = ["resolv"] }
 futures = "0.3.31"
+hickory-resolver = "0.25.2"
 humantime = "2.2.0"
 hyper = { version = "1.4.1", features = ["http1", "server"] }
 hyper-openssl = "0.10.2"

src/balancerd/Cargo.toml

@@ -18,11 +18,13 @@ use std::path::PathBuf;
 use std::time::Duration;
 
 use anyhow::Context;
-use domain::resolv::StubResolver;
+use hickory_resolver::{
+    Resolver, config::*, name_server::TokioConnectionProvider, system_conf::read_system_conf,
+};
 use jsonwebtoken::DecodingKey;
 use mz_balancerd::{
-    BUILD_INFO, BalancerConfig, BalancerService, CancellationResolver, FronteggResolver, Resolver,
-    SniResolver,
+    BUILD_INFO, BalancerConfig, BalancerResolver, BalancerService, CancellationResolver,
+    FronteggResolver, SniResolver, create_default_resolver,
 };
 use mz_frontegg_auth::{
     Authenticator, AuthenticatorConfig, DEFAULT_REFRESH_DROP_FACTOR,
@@ -241,8 +243,10 @@ pub async fn run(args: ServiceArgs, tracing_handle: TracingHandle) -> Result<(),
             if !cancellation_resolver_dir.is_dir() {
                 anyhow::bail!("{cancellation_resolver_dir:?} is not a directory");
             }
+
             (
-                Resolver::MultiTenant(
+                BalancerResolver::MultiTenant(
+                    create_default_resolver(),
                     FronteggResolver {
                         auth,
                         addr_template,
@@ -261,11 +265,7 @@ pub async fn run(args: ServiceArgs, tracing_handle: TracingHandle) -> Result<(),
                                     )
                                 })
                                 .expect("invalid port for pgwire_sni_resolver_template");
-                            Some(SniResolver {
-                                resolver: StubResolver::new(),
-                                template,
-                                port,
-                            })
+                            Some(SniResolver { template, port })
                         }
                     },
                 ),
@@ -284,8 +284,35 @@ pub async fn run(args: ServiceArgs, tracing_handle: TracingHandle) -> Result<(),
             };
             drop(addrs);
 
+            // Create a resolver for static addresses with the same caching configuration
+            let mut resolver_opts = ResolverOpts::default();
+            resolver_opts.cache_size = 10000;
+            resolver_opts.positive_max_ttl = Some(Duration::from_secs(10));
+            resolver_opts.positive_min_ttl = Some(Duration::from_secs(9));
+            resolver_opts.negative_min_ttl = Some(Duration::from_secs(1));
+
+            // Read system DNS configuration or fall back to defaults
+            let (config, opts) = read_system_conf()
+                .map(|(config, mut opts)| {
+                    // Override specific options while keeping system DNS servers
+                    opts.cache_size = resolver_opts.cache_size;
+                    opts.positive_max_ttl = resolver_opts.positive_max_ttl;
+                    opts.positive_min_ttl = resolver_opts.positive_min_ttl;
+                    opts.negative_min_ttl = resolver_opts.negative_min_ttl;
+                    (config, opts)
+                })
+                .unwrap_or_else(|err| {
+                    eprintln!("Failed to read system DNS configuration for static resolver, using defaults: {}", err);
+                    (ResolverConfig::default(), resolver_opts)
+                });
+
             (
-                Resolver::Static(addr.clone()),
+                BalancerResolver::Static(
+                    Resolver::builder_with_config(config, TokioConnectionProvider::default())
+                        .with_options(opts)
+                        .build(),
+                    addr.clone(),
+                ),
                 CancellationResolver::Static(addr),
             )
         }