Always return a broker address, even in an error state

Author: rjobanp

src/kafka-util/src/client.rs

@@ -470,9 +470,21 @@ where
                         // 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(
-                                    default_tunnel.clone(),
+                                    tunnel_config,
                                     &addr.host,
                                     addr.port.parse().unwrap(),
                                     self.ssh_timeout_config,
@@ -533,19 +545,37 @@ where
                     TunnelConfig::None => {
                         // If no rewrite is specified, we still should check that this potentially
                         // new broker address is a global address.
-                        let rewrite = self.runtime.block_on(async {
-                            let resolved = resolve_external_address(
+                        self.runtime.block_on(async {
+                            match resolve_external_address(
                                 &addr.host,
                                 self.enforce_external_addresses,
                             )
                             .await
-                            .unwrap();
-                            BrokerRewriteHandle::Simple(BrokerRewrite {
-                                host: resolved.to_string(),
-                                port: addr.port.parse().ok(),
-                            })
-                        });
-                        return_rewrite(&rewrite)
+                            {
+                                Ok(resolved) => {
+                                    let rewrite = BrokerRewriteHandle::Simple(BrokerRewrite {
+                                        host: resolved.to_string(),
+                                        port: addr.port.parse().ok(),
+                                    });
+                                    return_rewrite(&rewrite)
+                                }
+                                Err(e) => {
+                                    warn!(
+                                        "failed to resolve external address for {:?}: {}",
+                                        addr,
+                                        e.display_with_causes()
+                                    );
+                                    // We have to give rdkafka an address, as this callback can't fail,
+                                    // we just give it a random one that will never resolve, for safety
+                                    // in-case `addr.host` suddenly starts resolving again.
+                                    BrokerAddr {
+                                        host: "failed-broker-address.dev.materialize.com"
+                                            .to_string(),
+                                        port: 1337.to_string(),
+                                    }
+                                }
+                            }
+                        })
                     }
                 }
             }

src/storage-types/src/connections.rs

@@ -547,8 +547,14 @@ impl KafkaConnection {
                     .await?;
                 let key_pair = SshKeyPair::from_bytes(&secret)?;
 
+                // Ensure any ssh-bastion address we connect to is resolved to an external address.
+                let resolved = resolve_external_address(
+                    &ssh_tunnel.connection.host,
+                    ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
+                )
+                .await?;
                 context.set_default_tunnel(TunnelConfig::Ssh(SshTunnelConfig {
-                    host: ssh_tunnel.connection.host.clone(),
+                    host: resolved.to_string(),
                     port: ssh_tunnel.connection.port,
                     user: ssh_tunnel.connection.user.clone(),
                     key_pair,
@@ -567,25 +573,21 @@ impl KafkaConnection {
             };
             match &broker.tunnel {
                 Tunnel::Direct => {
+                    // By default, don't override broker address lookup.
+                    //
+                    // N.B.
+                    //
                     // We _could_ pre-setup the default ssh tunnel for all known brokers here, but
                     // we avoid doing because:
                     // - Its not necessary.
                     // - Not doing so makes it easier to test the `FailedDefaultSshTunnel` path
                     // in the `TunnelingClientContext`.
-
-                    // Ensure any broker address we connect to is resolved to an external address.
-                    let resolved = resolve_external_address(
-                        &addr.host,
-                        ENFORCE_EXTERNAL_ADDRESSES.get(storage_configuration.config_set()),
-                    )
-                    .await?;
-                    context.add_broker_rewrite(
-                        addr,
-                        BrokerRewrite {
-                            host: resolved.to_string(),
-                            port: None,
-                        },
-                    )
+                    //
+                    // NOTE that we do not need to use the `resolve_external_address` method to
+                    // validate the broker address here since it will be validated when the
+                    // connection is established in `src/kafka-util/src/client.rs`, and we do not
+                    // want to specify any BrokerRewrite that would override any default-tunnel
+                    // settings.
                 }
                 Tunnel::AwsPrivatelink(aws_privatelink) => {
                     let host = mz_cloud_resources::vpc_endpoint_host(

test/ssh-connection/kafka-source-after-ssh-failure.td

@@ -12,8 +12,8 @@
 # Ensure they all are marked as broken for ssh reasons
 > SELECT status FROM mz_internal.mz_source_statuses st
   JOIN mz_sources s ON st.id = s.id
-  WHERE
-  s.name in ('fixed_text', 'dynamic_text', 'fixed_plus_csr', 'dynamic_plus_csr')
+  WHERE error LIKE 'ssh:%'
+  AND s.name in ('fixed_text', 'dynamic_text', 'fixed_plus_csr', 'dynamic_plus_csr')
 stalled
 stalled
 stalled

test/testdrive/connection-validation.td

@@ -37,10 +37,10 @@ ALTER SYSTEM SET enable_connection_validation_syntax = true
 > CREATE CONNECTION invalid_tunnel TO SSH TUNNEL (HOST 'invalid', USER 'invalid', PORT 22)
 
 ! CREATE CONNECTION invalid_kafka_conn TO KAFKA (BROKERS ('${testdrive.kafka-addr}' USING SSH TUNNEL invalid_tunnel), SECURITY PROTOCOL PLAINTEXT)
-contains:failed to lookup address information: Name or service not known
+contains:failed to lookup address information
 
 # Create the connection without validation and validate later
 > CREATE CONNECTION invalid_kafka_conn TO KAFKA (BROKERS ('${testdrive.kafka-addr}' USING SSH TUNNEL invalid_tunnel), SECURITY PROTOCOL PLAINTEXT) WITH (VALIDATE = false)
 
 ! VALIDATE CONNECTION invalid_kafka_conn
-contains:failed to lookup address information: Name or service not known
+contains:failed to lookup address information