Provide all resolved IPs to mysql, postgres, and ssh clients

Author: rjobanp

Cargo.lock

@@ -192,6 +192,9 @@ postgres-types = { git = "https://github.com/MaterializeInc/rust-postgres" }
 postgres-openssl = { git = "https://github.com/MaterializeInc/rust-postgres" }
 postgres_array = { git = "https://github.com/MaterializeInc/rust-postgres-array" }
 
+# Waiting on https://github.com/blackbeam/mysql_async/pull/300
+mysql_async = { git = "https://github.com/MaterializeInc/mysql_async", rev = "6afd2136181f2ec93b7ca7de524f6d02b6f10c1d" }
+
 # Waiting on https://github.com/MaterializeInc/serde-value/pull/35.
 serde-value = { git = "https://github.com/MaterializeInc/serde-value.git" }
 

Cargo.toml

@@ -22,7 +22,7 @@ use anyhow::{anyhow, Context};
 use crossbeam::channel::{unbounded, Receiver, Sender};
 use mz_ore::collections::CollectionExt;
 use mz_ore::error::ErrorExt;
-use mz_ore::netio::resolve_external_address;
+use mz_ore::netio::resolve_address;
 use mz_ssh_util::tunnel::{SshTimeoutConfig, SshTunnelConfig, SshTunnelStatus};
 use mz_ssh_util::tunnel_manager::{ManagedSshTunnelHandle, SshTunnelManager};
 use rdkafka::client::{BrokerAddr, Client, NativeClient, OAuthToken};
@@ -297,7 +297,7 @@ enum BrokerRewriteHandle {
     FailedDefaultSshTunnel(String),
     /// We store an error if DNS resolution fails when resolving
     /// a new broker host.
-    FailedDNSResolution(String),
+    FailedDnsResolution(String),
 }
 
 /// Tunneling clients
@@ -425,7 +425,7 @@ impl<C> TunnelingClientContext<C> {
                 BrokerRewriteHandle::FailedDefaultSshTunnel(e) => {
                     SshTunnelStatus::Errored(e.clone())
                 }
-                BrokerRewriteHandle::Simple(_) | BrokerRewriteHandle::FailedDNSResolution(_) => {
+                BrokerRewriteHandle::Simple(_) | BrokerRewriteHandle::FailedDnsResolution(_) => {
                     SshTunnelStatus::Running
                 }
             })
@@ -447,7 +447,7 @@ impl<C> TunnelingClientContext<C> {
         let broker_status = rewrites
             .values()
             .map(|handle| match handle {
-                BrokerRewriteHandle::FailedDNSResolution(e) => BrokerStatus::Failed(e.clone()),
+                BrokerRewriteHandle::FailedDnsResolution(e) => BrokerStatus::Failed(e.clone()),
                 _ => BrokerStatus::Nominal,
             })
             .fold(BrokerStatus::Nominal, |acc, status| match (acc, status) {
@@ -486,7 +486,7 @@ where
                     }
                 }
                 BrokerRewriteHandle::FailedDefaultSshTunnel(_)
-                | BrokerRewriteHandle::FailedDNSResolution(_) => {
+                | BrokerRewriteHandle::FailedDnsResolution(_) => {
                     unreachable!()
                 }
             };
@@ -510,28 +510,18 @@ where
         match rewrite {
             None
             | Some(BrokerRewriteHandle::FailedDefaultSshTunnel(_))
-            | Some(BrokerRewriteHandle::FailedDNSResolution(_)) => {
+            | Some(BrokerRewriteHandle::FailedDnsResolution(_)) => {
                 match &self.default_tunnel {
                     TunnelConfig::Ssh(default_tunnel) => {
                         // Multiple users could all run `connect` at the same time; only one ssh
                         // tunnel will ever be connected, and only one will be inserted into the
                         // map.
                         let ssh_tunnel = self.runtime.block_on(async {
-                            // Ensure the default tunnel host is resolved to an external address.
-                            let resolved_tunnel_addr = resolve_external_address(
-                                &default_tunnel.host,
-                                self.enforce_external_addresses,
-                            )
-                            .await?;
-                            let tunnel_config = SshTunnelConfig {
-                                host: resolved_tunnel_addr.to_string(),
-                                port: default_tunnel.port,
-                                user: default_tunnel.user.clone(),
-                                key_pair: default_tunnel.key_pair.clone(),
-                            };
                             self.ssh_tunnel_manager
                                 .connect(
-                                    tunnel_config,
+                                    // We know the default_tunnel has already been validated by the `resolve_address`
+                                    // method when it was provided to the client, so we don't need to check it again.
+                                    default_tunnel.clone(),
                                     &addr.host,
                                     addr.port.parse().unwrap(),
                                     self.ssh_timeout_config,
@@ -546,7 +536,7 @@ where
                                         if matches!(
                                             o.get(),
                                             BrokerRewriteHandle::FailedDefaultSshTunnel(_)
-                                                | BrokerRewriteHandle::FailedDNSResolution(_)
+                                                | BrokerRewriteHandle::FailedDnsResolution(_)
                                         ) =>
                                     {
                                         o.insert(BrokerRewriteHandle::SshTunnel(
@@ -594,15 +584,14 @@ where
                         // If no rewrite is specified, we still should check that this potentially
                         // new broker address is a global address.
                         self.runtime.block_on(async {
-                            match resolve_external_address(
-                                &addr.host,
-                                self.enforce_external_addresses,
-                            )
-                            .await
+                            match resolve_address(&addr.host, self.enforce_external_addresses).await
                             {
                                 Ok(resolved) => {
                                     let rewrite = BrokerRewriteHandle::Simple(BrokerRewrite {
-                                        host: resolved.to_string(),
+                                        // `resolve_address` will always return at least one address.
+                                        // TODO: Once we have a way to provide multiple hosts to rdkafka, we should
+                                        // return all resolved addresses here.
+                                        host: resolved.first().unwrap().to_string(),
                                         port: addr.port.parse().ok(),
                                     });
                                     return_rewrite(&rewrite)
@@ -616,7 +605,7 @@ where
                                     // Write an error if no one else has already written one.
                                     let mut rewrites = self.rewrites.lock().expect("poisoned");
                                     rewrites.entry(addr.clone()).or_insert_with(|| {
-                                        BrokerRewriteHandle::FailedDNSResolution(
+                                        BrokerRewriteHandle::FailedDnsResolution(
                                             e.to_string_with_causes(),
                                         )
                                     });

src/kafka-util/src/client.rs

@@ -7,6 +7,7 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::net::IpAddr;
 use std::ops::{Deref, DerefMut};
 use std::time::Duration;
 
@@ -25,7 +26,7 @@ use crate::MySqlError;
 #[derive(Debug, PartialEq, Clone)]
 pub enum TunnelConfig {
     /// Establish a direct TCP connection to the database host.
-    Direct { tcp_host_override: Option<String> },
+    Direct { resolved_ips: Option<Vec<IpAddr>> },
     /// Establish a TCP connection to the database via an SSH tunnel.
     /// This means first establishing an SSH connection to a bastion host,
     /// and then opening a separate connection from that host to the database.
@@ -225,27 +226,9 @@ impl Config {
         ssh_tunnel_manager: &SshTunnelManager,
     ) -> Result<MySqlConn, MySqlError> {
         match &self.tunnel {
-            TunnelConfig::Direct { tcp_host_override } => {
-                // Override the connection host for the actual TCP connection to point to
-                // the privatelink hostname instead.
-                let mut opts_builder = OptsBuilder::from_opts(self.inner.clone());
-
-                if let Some(tcp_override) = tcp_host_override {
-                    opts_builder = opts_builder.ip_or_hostname(tcp_override);
-
-                    if let Some(ssl_opts) = self.inner.ssl_opts() {
-                        if !ssl_opts.skip_domain_validation() {
-                            // If the TLS configuration will validate the hostname, we need to set
-                            // the TLS hostname back to the actual upstream host and not the
-                            // TCP hostname.
-                            opts_builder = opts_builder.ssl_opts(Some(
-                                ssl_opts.clone().with_danger_tls_hostname_override(Some(
-                                    self.inner.ip_or_hostname().to_string(),
-                                )),
-                            ));
-                        }
-                    }
-                };
+            TunnelConfig::Direct { resolved_ips } => {
+                let opts_builder =
+                    OptsBuilder::from_opts(self.inner.clone()).resolved_ips(resolved_ips.clone());
 
                 Ok(MySqlConn {
                     conn: Conn::new(opts_builder).await.map_err(MySqlError::from)?,

src/mysql-util/src/tunnel.rs

@@ -22,10 +22,10 @@ const DUMMY_PORT: u16 = 11111;
 
 /// Resolves a host address and ensures it is a global address when `enforce_global` is set.
 /// This parameter is useful when connecting to user-defined unverified addresses.
-pub async fn resolve_external_address(
+pub async fn resolve_address(
     mut host: &str,
     enforce_global: bool,
-) -> Result<IpAddr, io::Error> {
+) -> Result<Vec<IpAddr>, io::Error> {
     // `net::lookup_host` requires a port to be specified, but we don't care about the port.
     let mut port = DUMMY_PORT;
     // If a port is already specified, use it and remove it from the host.
@@ -37,23 +37,26 @@ pub async fn resolve_external_address(
     }
 
     let mut addrs = lookup_host((host, port)).await?;
-
-    if let Some(addr) = addrs.next() {
+    let mut ips = Vec::new();
+    while let Some(addr) = addrs.next() {
         let ip = addr.ip();
         if enforce_global && !is_global(ip) {
             Err(io::Error::new(
                 io::ErrorKind::AddrNotAvailable,
                 "address is not global",
-            ))
+            ))?
         } else {
-            Ok(ip)
+            ips.push(ip);
         }
-    } else {
+    }
+
+    if ips.len() == 0 {
         Err(io::Error::new(
             io::ErrorKind::AddrNotAvailable,
             "no addresses found",
-        ))
+        ))?
     }
+    Ok(ips)
 }
 
 fn is_global(addr: IpAddr) -> bool {

src/ore/src/netio/dns.rs

@@ -22,7 +22,7 @@ mod read_exact;
 mod socket;
 
 pub use crate::netio::async_ready::AsyncReady;
-pub use crate::netio::dns::resolve_external_address;
+pub use crate::netio::dns::resolve_address;
 pub use crate::netio::framed::{FrameTooBig, MAX_FRAME_SIZE};
 pub use crate::netio::read_exact::{read_exact_or_eof, ReadExactOrEof};
 pub use crate::netio::socket::{Listener, SocketAddr, SocketAddrType, Stream, UnixSocketAddr};

src/ore/src/netio/mod.rs

@@ -7,6 +7,7 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::net::IpAddr;
 use std::time::Duration;
 
 use mz_ore::option::OptionExt;
@@ -37,7 +38,7 @@ macro_rules! bail_generic {
 #[derive(Debug, PartialEq, Clone)]
 pub enum TunnelConfig {
     /// Establish a direct TCP connection to the database host.
-    Direct { tcp_host_override: Option<String> },
+    Direct { resolved_ips: Option<Vec<IpAddr>> },
     /// Establish a TCP connection to the database via an SSH tunnel.
     /// This means first establishing an SSH connection to a bastion host,
     /// and then opening a separate connection from that host to the database.
@@ -222,19 +223,22 @@ impl Config {
         })?;
 
         match &self.tunnel {
-            TunnelConfig::Direct { tcp_host_override } => {
-                // Override the TCP host we connect to, leaving the host used for TLS verification
-                // as the original host
-                if let Some(tcp_override) = tcp_host_override {
-                    let (host, _) = self.address()?;
-                    postgres_config.tls_verify_host(host);
-
-                    match postgres_config.get_hosts_mut() {
-                        [Host::Tcp(host)] => *host = tcp_override.clone(),
+            TunnelConfig::Direct { resolved_ips } => {
+                if let Some(ips) = resolved_ips {
+                    let host = match postgres_config.get_hosts() {
+                        [Host::Tcp(host)] => host,
                         _ => bail_generic!(
                             "only TCP connections to a single PostgreSQL server are supported"
                         ),
                     }
+                    .to_owned();
+                    // The number of 'host' and 'hostaddr' values must be the same.
+                    for (idx, ip) in ips.iter().enumerate() {
+                        if idx != 0 {
+                            postgres_config.host(&host);
+                        }
+                        postgres_config.hostaddr(ip.clone());
+                    }
                 };
 
                 let (client, connection) = postgres_config.connect(tls).await?;
@@ -276,17 +280,25 @@ impl Config {
                 // `tokio_postgres::Config` to do this is somewhat confusing, and requires we edit
                 // the singular host in place.
 
-                let (host, _) = self.address()?;
-                postgres_config.tls_verify_host(host);
-
                 let privatelink_host = mz_cloud_resources::vpc_endpoint_name(*connection_id);
+                let privatelink_addrs = tokio::net::lookup_host(privatelink_host).await?;
 
-                match postgres_config.get_hosts_mut() {
-                    [Host::Tcp(host)] => *host = privatelink_host,
+                // Override the actual IPs to connect to for the TCP connection, leaving the original host in-place
+                // for TLS verification
+                let host = match postgres_config.get_hosts() {
+                    [Host::Tcp(host)] => host,
                     _ => bail_generic!(
                         "only TCP connections to a single PostgreSQL server are supported"
                     ),
                 }
+                .to_owned();
+                // The number of 'host' and 'hostaddr' values must be the same.
+                for (idx, addr) in privatelink_addrs.enumerate() {
+                    if idx != 0 {
+                        postgres_config.host(&host);
+                    }
+                    postgres_config.hostaddr(addr.ip());
+                }
 
                 let (client, connection) = postgres_config.connect(tls).await?;
                 task::spawn(|| task_name, connection);