mz_kafka_util/
client.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Helpers for working with Kafka's client API.
11
12use anyhow::bail;
13use aws_config::SdkConfig;
14use fancy_regex::Regex;
15use std::collections::{BTreeMap, btree_map};
16use std::error::Error;
17use std::io;
18use std::net::{SocketAddr, ToSocketAddrs};
19use std::str::FromStr;
20use std::sync::Arc;
21use std::sync::Mutex;
22use std::time::Duration;
23use tokio::sync::watch;
24
25use anyhow::{Context, anyhow};
26use crossbeam::channel::{Receiver, Sender, unbounded};
27use mz_ore::collections::CollectionExt;
28use mz_ore::error::ErrorExt;
29use mz_ore::future::InTask;
30use mz_ore::netio::{DnsResolutionError, resolve_address};
31use mz_ssh_util::tunnel::{SshTimeoutConfig, SshTunnelConfig, SshTunnelStatus};
32use mz_ssh_util::tunnel_manager::{ManagedSshTunnelHandle, SshTunnelManager};
33use rdkafka::client::{Client, NativeClient, OAuthToken};
34use rdkafka::config::{
35    ClientConfig, FromClientConfig, FromClientConfigAndContext, RDKafkaLogLevel,
36};
37use rdkafka::consumer::{ConsumerContext, Rebalance};
38use rdkafka::error::{KafkaError, KafkaResult, RDKafkaErrorCode};
39use rdkafka::producer::{DefaultProducerContext, DeliveryResult, ProducerContext};
40use rdkafka::types::RDKafkaRespErr;
41use rdkafka::util::Timeout;
42use rdkafka::{ClientContext, Statistics, TopicPartitionList};
43use serde::{Deserialize, Serialize};
44use tokio::runtime::Handle;
45use tracing::{Level, debug, info, trace, warn};
46
47use crate::aws;
48
49/// A reasonable default timeout when refreshing topic metadata. This is configured
50/// at a source level.
51// 30s may seem infrequent, but the default is 5m. More frequent metadata
52// refresh rates are surprising to Kafka users, as topic partition counts hardly
53// ever change in production.
54pub const DEFAULT_TOPIC_METADATA_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
55
56/// A `ClientContext` implementation that uses `tracing` instead of `log`
57/// macros.
58///
59/// All code in Materialize that constructs Kafka clients should use this
60/// context or a custom context that delegates the `log` and `error` methods to
61/// this implementation.
62pub struct MzClientContext {
63    /// The last observed error log, if any.
64    error_tx: Sender<MzKafkaError>,
65    /// A tokio watch that retains the last statistics received by rdkafka and provides async
66    /// notifications to anyone interested in subscribing.
67    statistics_tx: watch::Sender<Statistics>,
68}
69
70impl Default for MzClientContext {
71    fn default() -> Self {
72        Self::with_errors().0
73    }
74}
75
76impl MzClientContext {
77    /// Constructs a new client context and returns an mpsc `Receiver` that can be used to learn
78    /// about librdkafka errors.
79    // `crossbeam` channel receivers can be cloned, but this is intended to be used as a mpsc,
80    // until we upgrade to `1.72` and the std mpsc sender is `Sync`.
81    pub fn with_errors() -> (Self, Receiver<MzKafkaError>) {
82        let (error_tx, error_rx) = unbounded();
83        let (statistics_tx, _) = watch::channel(Default::default());
84        let ctx = Self {
85            error_tx,
86            statistics_tx,
87        };
88        (ctx, error_rx)
89    }
90
91    /// Creates a tokio Watch subscription for statistics reported by librdkafka. It is necessary
92    /// that the `statistics.ms.interval` is set for this stream to contain any values.
93    pub fn subscribe_statistics(&self) -> watch::Receiver<Statistics> {
94        self.statistics_tx.subscribe()
95    }
96
97    fn record_error(&self, msg: &str) {
98        let err = match MzKafkaError::from_str(msg) {
99            Ok(err) => err,
100            Err(()) => {
101                warn!(original_error = msg, "failed to parse kafka error");
102                MzKafkaError::Internal(msg.to_owned())
103            }
104        };
105        // If no one cares about errors we drop them on the floor
106        let _ = self.error_tx.send(err);
107    }
108}
109
110/// Prefix of the `FAIL` log [`TunnelingClientContext`] emits when it rejects a
111/// broker, which [`MzKafkaError::from_str`] maps to
112/// [`MzKafkaError::PrivateBrokerAddress`].
113const PRIVATE_BROKER_ADDRESS_MARKER: &str = "Broker address resolved to a private IP";
114
115/// A structured error type for errors reported by librdkafka through its logs.
116#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
117pub enum MzKafkaError {
118    /// Invalid username or password
119    #[error("Invalid username or password")]
120    InvalidCredentials,
121    /// Missing CA certificate
122    #[error("Invalid CA certificate")]
123    InvalidCACertificate,
124    /// Broker might require SSL encryption
125    #[error("Disconnected during handshake; broker might require SSL encryption")]
126    SSLEncryptionMaybeRequired,
127    /// Broker does not support SSL connections
128    #[error("Broker does not support SSL connections")]
129    SSLUnsupported,
130    /// Broker did not provide a certificate
131    #[error("Broker did not provide a certificate")]
132    BrokerCertificateMissing,
133    /// Failed to verify broker certificate
134    #[error("Failed to verify broker certificate")]
135    InvalidBrokerCertificate,
136    /// Connection reset
137    #[error("Connection reset: {0}")]
138    ConnectionReset(String),
139    /// Connection timeout
140    #[error("Connection timeout")]
141    ConnectionTimeout,
142    /// Failed to resolve hostname
143    #[error("Failed to resolve hostname")]
144    HostnameResolutionFailed,
145    /// Broker resolved to a private address while external addresses are enforced
146    #[error("Broker address resolved to a private IP")]
147    PrivateBrokerAddress,
148    /// Unsupported SASL mechanism
149    #[error("Unsupported SASL mechanism")]
150    UnsupportedSASLMechanism,
151    /// Unsupported broker version
152    #[error("Unsupported broker version")]
153    UnsupportedBrokerVersion,
154    /// Connection to broker failed
155    #[error("Broker transport failure")]
156    BrokerTransportFailure,
157    /// All brokers down
158    #[error("All brokers down")]
159    AllBrokersDown,
160    /// SASL authentication required
161    #[error("SASL authentication required")]
162    SaslAuthenticationRequired,
163    /// SASL authentication required
164    #[error("SASL authentication failed")]
165    SaslAuthenticationFailed,
166    /// SSL authentication required
167    #[error("SSL authentication required")]
168    SslAuthenticationRequired,
169    /// Unknown topic or partition
170    #[error("Unknown topic or partition")]
171    UnknownTopicOrPartition,
172    /// An internal kafka error
173    #[error("Internal kafka error: {0}")]
174    Internal(String),
175}
176
177impl FromStr for MzKafkaError {
178    type Err = ();
179
180    fn from_str(s: &str) -> Result<Self, Self::Err> {
181        if s.contains("Authentication failed: Invalid username or password") {
182            Ok(Self::InvalidCredentials)
183        } else if s.contains("broker certificate could not be verified") {
184            Ok(Self::InvalidCACertificate)
185        } else if s.contains("connecting to a SSL listener?") {
186            Ok(Self::SSLEncryptionMaybeRequired)
187        } else if s.contains("client SSL authentication might be required") {
188            Ok(Self::SslAuthenticationRequired)
189        } else if s.contains("connecting to a PLAINTEXT broker listener") {
190            Ok(Self::SSLUnsupported)
191        } else if s.contains("Broker did not provide a certificate") {
192            Ok(Self::BrokerCertificateMissing)
193        } else if s.contains("Failed to verify broker certificate: ") {
194            Ok(Self::InvalidBrokerCertificate)
195        } else if let Some((_prefix, inner)) = s.split_once("Send failed: ") {
196            Ok(Self::ConnectionReset(inner.to_owned()))
197        } else if let Some((_prefix, inner)) = s.split_once("Receive failed: ") {
198            Ok(Self::ConnectionReset(inner.to_owned()))
199        } else if s.contains("request(s) timed out: disconnect") {
200            Ok(Self::ConnectionTimeout)
201        } else if s.contains(PRIVATE_BROKER_ADDRESS_MARKER) {
202            Ok(Self::PrivateBrokerAddress)
203        } else if s.contains("Failed to resolve") {
204            Ok(Self::HostnameResolutionFailed)
205        } else if s.contains("mechanism handshake failed:") {
206            Ok(Self::UnsupportedSASLMechanism)
207        } else if s.contains(
208            "verify that security.protocol is correctly configured, \
209            broker might require SASL authentication",
210        ) {
211            Ok(Self::SaslAuthenticationRequired)
212        } else if s.contains("SASL authentication error: Authentication failed") {
213            Ok(Self::SaslAuthenticationFailed)
214        } else if s
215            .contains("incorrect security.protocol configuration (connecting to a SSL listener?)")
216        {
217            Ok(Self::SslAuthenticationRequired)
218        } else if s.contains("probably due to broker version < 0.10") {
219            Ok(Self::UnsupportedBrokerVersion)
220        } else if s.contains("Disconnected while requesting ApiVersion")
221            || s.contains("Broker transport failure")
222            || s.contains("Connection refused")
223        {
224            Ok(Self::BrokerTransportFailure)
225        } else if Regex::new(r"\d+/\d+ brokers are down")
226            .unwrap()
227            .is_match(s)
228            .unwrap_or_default()
229        {
230            // The counts can differ ("0/2"): a re-bootstrap resets the broker table first.
231            Ok(Self::AllBrokersDown)
232        } else if s.contains("Unknown topic or partition") || s.contains("Unknown partition") {
233            Ok(Self::UnknownTopicOrPartition)
234        } else {
235            Err(())
236        }
237    }
238}
239
240impl ClientContext for MzClientContext {
241    fn log(&self, level: rdkafka::config::RDKafkaLogLevel, fac: &str, log_message: &str) {
242        use rdkafka::config::RDKafkaLogLevel::*;
243
244        // Sniff out log messages that indicate errors.
245        //
246        // We consider any event at error, critical, alert, or emergency level,
247        // for self explanatory reasons. We also consider any event with a
248        // facility of `FAIL`. librdkafka often uses info or warn level for
249        // these `FAIL` events, but as they always indicate a failure to connect
250        // to a broker we want to always treat them as errors.
251        if matches!(level, Emerg | Alert | Critical | Error) || fac == "FAIL" {
252            self.record_error(log_message);
253        }
254
255        // Copied from https://docs.rs/rdkafka/0.28.0/src/rdkafka/client.rs.html#58-79
256        // but using `tracing`
257        match level {
258            Emerg | Alert | Critical | Error => {
259                // We downgrade error messages to `warn!` level to avoid
260                // sending the errors to Sentry. Most errors are customer
261                // configuration problems that are not appropriate to send to
262                // Sentry.
263                warn!(target: "librdkafka", "error: {} {}", fac, log_message);
264            }
265            Warning => warn!(target: "librdkafka", "warning: {} {}", fac, log_message),
266            Notice => info!(target: "librdkafka", "{} {}", fac, log_message),
267            Info => info!(target: "librdkafka", "{} {}", fac, log_message),
268            Debug => debug!(target: "librdkafka", "{} {}", fac, log_message),
269        }
270    }
271
272    fn stats(&self, statistics: Statistics) {
273        self.statistics_tx.send_replace(statistics);
274    }
275
276    fn error(&self, error: KafkaError, reason: &str) {
277        self.record_error(reason);
278        // Refer to the comment in the `log` callback.
279        warn!(target: "librdkafka", "error: {}: {}", error, reason);
280    }
281}
282
283impl ConsumerContext for MzClientContext {}
284
285impl ProducerContext for MzClientContext {
286    type DeliveryOpaque = <DefaultProducerContext as ProducerContext>::DeliveryOpaque;
287    fn delivery(
288        &self,
289        delivery_result: &DeliveryResult<'_>,
290        delivery_opaque: Self::DeliveryOpaque,
291    ) {
292        DefaultProducerContext.delivery(delivery_result, delivery_opaque);
293    }
294}
295
296/// The address of a Kafka broker.
297#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
298pub struct BrokerAddr {
299    /// The broker's hostname.
300    pub host: String,
301    /// The broker's port.
302    pub port: u16,
303}
304
305impl BrokerAddr {
306    /// Attempt to resolve this broker address into a list of socket addresses.
307    pub fn to_socket_addrs(&self) -> Result<Vec<SocketAddr>, io::Error> {
308        Ok((self.host.as_str(), self.port).to_socket_addrs()?.collect())
309    }
310}
311
312/// Rewrites a broker address.
313///
314/// For use with [`TunnelingClientContext`].
315#[derive(Debug, Clone)]
316pub struct BrokerRewrite {
317    /// The rewritten hostname.
318    pub host: String,
319    /// The rewritten port.
320    ///
321    /// If unspecified, the broker's original port is left unchanged.
322    pub port: Option<u16>,
323}
324
325impl BrokerRewrite {
326    /// Apply the rewrite to this broker address.
327    pub fn rewrite(&self, address: &BrokerAddr) -> BrokerAddr {
328        BrokerAddr {
329            host: self.host.clone(),
330            port: self.port.unwrap_or(address.port),
331        }
332    }
333}
334
335#[derive(Clone)]
336enum BrokerRewriteHandle {
337    Simple(BrokerRewrite),
338    SshTunnel(
339        // This ensures the ssh tunnel is not shutdown.
340        ManagedSshTunnelHandle,
341    ),
342    /// For _default_ ssh tunnels, we store an error if _creation_
343    /// of the tunnel failed, so that `tunnel_status` can return it.
344    FailedDefaultSshTunnel(String),
345}
346
347#[derive(Clone)]
348/// Parsed from a string, with optional leading and trailing '*' wildcards.
349pub struct ConnectionRulePattern {
350    /// If true, allow any combination of characters before the literal match.
351    pub prefix_wildcard: bool,
352    /// We expect the broker's host:port to match these characters in their entirety.
353    pub literal_match: String,
354    /// If true, allow any combination of characters after the literal match.
355    pub suffix_wildcard: bool,
356}
357
358impl std::fmt::Display for ConnectionRulePattern {
359    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
360        if self.prefix_wildcard {
361            f.write_str("*")?;
362        }
363        f.write_str(&self.literal_match)?;
364        if self.suffix_wildcard {
365            f.write_str("*")?;
366        }
367        Ok(())
368    }
369}
370
371impl ConnectionRulePattern {
372    /// Does this "{host}:{port}" address fit the pattern?
373    pub fn matches(&self, address: &str) -> bool {
374        if self.prefix_wildcard {
375            if self.suffix_wildcard {
376                address.contains(&self.literal_match)
377            } else {
378                address.ends_with(&self.literal_match)
379            }
380        } else if self.suffix_wildcard {
381            address.starts_with(&self.literal_match)
382        } else {
383            address == self.literal_match
384        }
385    }
386}
387
388#[derive(Clone)]
389/// Given a host address, map it to a different host.
390pub struct HostMappingRules {
391    /// Map matching hosts to a different host. First applicable rule wins.
392    pub rules: Vec<(ConnectionRulePattern, BrokerRewrite)>,
393}
394
395impl HostMappingRules {
396    /// Rewrite this broker address according to the rules. Returns `None` when
397    /// no rule matches.
398    pub fn rewrite(&self, src: &BrokerAddr) -> Option<BrokerAddr> {
399        let address = format!("{}:{}", src.host, src.port);
400        for (pattern, dst) in &self.rules {
401            if pattern.matches(&address) {
402                let result = dst.rewrite(src);
403                info!(
404                    "HostMappingRules: broker {}:{} matched pattern '{}' -> rewriting to {}:{}",
405                    src.host, src.port, pattern, result.host, result.port,
406                );
407                return Some(result);
408            }
409        }
410
411        warn!(
412            "HostMappingRules: broker {}:{} matched no rules, using original address",
413            src.host, src.port,
414        );
415        None
416    }
417}
418
419/// Tunneling clients
420/// used for re-writing ports / hosts
421#[derive(Clone)]
422pub enum TunnelConfig {
423    /// Tunnel config option for SSH tunnels
424    Ssh(SshTunnelConfig),
425    /// Re-writes internal hosts using the value, used for privatelink
426    StaticHost(String),
427    /// Re-writes internal hosts according to an ordered list of rules, also used for privatelink
428    Rules(HostMappingRules),
429    /// Performs no re-writes
430    None,
431}
432
433/// A client context that supports rewriting broker addresses.
434#[derive(Clone)]
435pub struct TunnelingClientContext<C> {
436    inner: C,
437    rewrites: Arc<Mutex<BTreeMap<BrokerAddr, BrokerRewriteHandle>>>,
438    default_tunnel: TunnelConfig,
439    in_task: InTask,
440    ssh_tunnel_manager: SshTunnelManager,
441    ssh_timeout_config: SshTimeoutConfig,
442    aws_config: Option<SdkConfig>,
443    runtime: Handle,
444    enforce_external_addresses: bool,
445}
446
447impl<C> TunnelingClientContext<C> {
448    /// Constructs a new context that wraps `inner`.
449    pub fn new(
450        inner: C,
451        runtime: Handle,
452        ssh_tunnel_manager: SshTunnelManager,
453        ssh_timeout_config: SshTimeoutConfig,
454        aws_config: Option<SdkConfig>,
455        in_task: InTask,
456        enforce_external_addresses: bool,
457    ) -> TunnelingClientContext<C> {
458        TunnelingClientContext {
459            inner,
460            rewrites: Arc::new(Mutex::new(BTreeMap::new())),
461            default_tunnel: TunnelConfig::None,
462            in_task,
463            ssh_tunnel_manager,
464            ssh_timeout_config,
465            aws_config,
466            runtime,
467            enforce_external_addresses,
468        }
469    }
470
471    /// Adds the default broker rewrite rule.
472    ///
473    /// Connections to brokers that aren't specified in other rewrites will be rewritten to connect to
474    /// `rewrite_host` and `rewrite_port` instead.
475    pub fn set_default_tunnel(&mut self, tunnel: TunnelConfig) {
476        self.default_tunnel = tunnel;
477    }
478
479    /// Adds an SSH tunnel for a specific broker.
480    ///
481    /// Overrides the existing SSH tunnel or rewrite for this broker, if any.
482    ///
483    /// This tunnel allows the rewrite to evolve over time, for example, if
484    /// the ssh tunnel's address changes if it fails and restarts.
485    pub async fn add_ssh_tunnel(
486        &self,
487        broker: BrokerAddr,
488        tunnel: SshTunnelConfig,
489    ) -> Result<(), anyhow::Error> {
490        let ssh_tunnel = self
491            .ssh_tunnel_manager
492            .connect(
493                tunnel,
494                &broker.host,
495                broker.port,
496                self.ssh_timeout_config,
497                self.in_task,
498            )
499            .await
500            .context("creating ssh tunnel")?;
501
502        let mut rewrites = self.rewrites.lock().expect("poisoned");
503        rewrites.insert(broker, BrokerRewriteHandle::SshTunnel(ssh_tunnel));
504        Ok(())
505    }
506
507    /// Adds a broker rewrite rule.
508    ///
509    /// Overrides the existing SSH tunnel or rewrite for this broker, if any.
510    ///
511    /// `rewrite` is `BrokerRewrite` that specifies how to rewrite the address for `broker`.
512    pub fn add_broker_rewrite(&self, broker: BrokerAddr, rewrite: BrokerRewrite) {
513        let mut rewrites = self.rewrites.lock().expect("poisoned");
514        rewrites.insert(broker, BrokerRewriteHandle::Simple(rewrite));
515    }
516
517    /// Returns a reference to the wrapped context.
518    pub fn inner(&self) -> &C {
519        &self.inner
520    }
521
522    /// Returns a _consolidated_ `SshTunnelStatus` that communicates the status
523    /// of all active ssh tunnels `self` knows about.
524    pub fn tunnel_status(&self) -> SshTunnelStatus {
525        self.rewrites
526            .lock()
527            .expect("poisoned")
528            .values()
529            .map(|handle| match handle {
530                BrokerRewriteHandle::SshTunnel(s) => s.check_status(),
531                BrokerRewriteHandle::FailedDefaultSshTunnel(e) => {
532                    SshTunnelStatus::Errored(e.clone())
533                }
534                BrokerRewriteHandle::Simple(_) => SshTunnelStatus::Running,
535            })
536            .fold(SshTunnelStatus::Running, |acc, status| {
537                match (acc, status) {
538                    (SshTunnelStatus::Running, SshTunnelStatus::Errored(e))
539                    | (SshTunnelStatus::Errored(e), SshTunnelStatus::Running) => {
540                        SshTunnelStatus::Errored(e)
541                    }
542                    (SshTunnelStatus::Errored(err), SshTunnelStatus::Errored(e)) => {
543                        SshTunnelStatus::Errored(format!("{}, {}", err, e))
544                    }
545                    (SshTunnelStatus::Running, SshTunnelStatus::Running) => {
546                        SshTunnelStatus::Running
547                    }
548                }
549            })
550    }
551}
552
553impl<C: ClientContext> TunnelingClientContext<C> {
554    /// Resolves a broker address that the user controls directly, i.e. one not
555    /// rewritten to an SSH tunnel or PrivateLink endpoint. With
556    /// `enforce_external_addresses`, fails unless every address is global.
557    // NOTE: librdkafka connects only to addresses returned here, and every
558    // re-resolution (e.g. after `broker.address.ttl`) comes back through this
559    // callback, so checking here also defeats DNS rebinding. TLS SNI and
560    // certificate verification still use the broker's hostname.
561    fn resolve_direct(&self, addr: &BrokerAddr) -> Result<Vec<SocketAddr>, io::Error> {
562        if !self.enforce_external_addresses {
563            return addr.to_socket_addrs();
564        }
565        match self.runtime.block_on(resolve_address(&addr.host, true)) {
566            Ok(ips) => Ok(ips
567                .into_iter()
568                .map(|ip| SocketAddr::new(ip, addr.port))
569                .collect()),
570            Err(e) => {
571                // librdkafka reports a callback failure as a generic resolution
572                // failure, so surface the reason through the inner context.
573                if let DnsResolutionError::PrivateAddress = e {
574                    self.inner.log(
575                        RDKafkaLogLevel::Error,
576                        "FAIL",
577                        &format!(
578                            "{PRIVATE_BROKER_ADDRESS_MARKER}: {}:{} is not routable on the public internet",
579                            addr.host, addr.port
580                        ),
581                    );
582                }
583                Err(io::Error::other(e))
584            }
585        }
586    }
587}
588
589impl<C> ClientContext for TunnelingClientContext<C>
590where
591    C: ClientContext,
592{
593    const ENABLE_REFRESH_OAUTH_TOKEN: bool = true;
594
595    fn generate_oauth_token(
596        &self,
597        _oauthbearer_config: Option<&str>,
598    ) -> Result<OAuthToken, Box<dyn Error>> {
599        // NOTE(benesch): We abuse the `TunnelingClientContext` to handle AWS
600        // IAM authentication because it's used in exactly the right places and
601        // already has a handle to the Tokio runtime. It might be slightly
602        // cleaner to have a separate `AwsIamAuthenticatingClientContext`, but
603        // that would be quite a bit of additional plumbing.
604
605        // NOTE(benesch): at the moment, the only OAUTHBEARER authentication we
606        // support is AWS IAM, so we can assume that if this method is invoked
607        // AWS IAM is desired. We may need to generalize this in the future.
608
609        info!(target: "librdkafka", "generating OAuth token");
610
611        let generate = || {
612            let Some(sdk_config) = &self.aws_config else {
613                bail!("internal error: AWS configuration missing");
614            };
615
616            self.runtime.block_on(aws::generate_auth_token(sdk_config))
617        };
618
619        match generate() {
620            Ok((token, lifetime_ms)) => {
621                info!(target: "librdkafka", %lifetime_ms, "successfully generated OAuth token");
622                trace!(target: "librdkafka", %token);
623                Ok(OAuthToken {
624                    token,
625                    lifetime_ms,
626                    principal_name: "".to_string(),
627                })
628            }
629            Err(e) => {
630                warn!(target: "librdkafka", "failed to generate OAuth token: {e:#}");
631                Err(e.into())
632            }
633        }
634    }
635
636    /// Look up the broker's address in our book of rewrites.
637    /// If we've already rewritten it before, reuse the existing rewrite.
638    /// Otherwise, use our "default tunnel" rewriting strategy to attempt to rewrite this broker's address
639    /// and record it in the book of rewrites.
640    fn resolve_broker_addr(&self, host: &str, port: u16) -> Result<Vec<SocketAddr>, io::Error> {
641        info!("kafka: resolve_broker_addr called for {}:{}", host, port);
642        let return_rewrite = |rewrite: &BrokerRewriteHandle| -> Result<Vec<SocketAddr>, io::Error> {
643            let rewrite = match rewrite {
644                BrokerRewriteHandle::Simple(rewrite) => rewrite.clone(),
645                BrokerRewriteHandle::SshTunnel(ssh_tunnel) => {
646                    // The port for this can change over time, as the ssh tunnel is maintained through
647                    // errors.
648                    let addr = ssh_tunnel.local_addr();
649                    BrokerRewrite {
650                        host: addr.ip().to_string(),
651                        port: Some(addr.port()),
652                    }
653                }
654                BrokerRewriteHandle::FailedDefaultSshTunnel(_) => {
655                    unreachable!()
656                }
657            };
658            let rewrite_port = rewrite.port.unwrap_or(port);
659
660            info!(
661                "rewriting broker {}:{} to {}:{}",
662                host, port, rewrite.host, rewrite_port
663            );
664
665            (rewrite.host, rewrite_port)
666                .to_socket_addrs()
667                .map(|addrs| addrs.collect())
668        };
669
670        let addr = BrokerAddr {
671            host: host.into(),
672            port,
673        };
674        let rewrite = self.rewrites.lock().expect("poisoned").get(&addr).cloned();
675
676        match rewrite {
677            // No (successful) broker address rewrite exists yet.
678            None | Some(BrokerRewriteHandle::FailedDefaultSshTunnel(_)) => {
679                // "Default tunnel" is actually the configured rewriting strategy used for brokers we haven't already rewritten.
680                match &self.default_tunnel {
681                    // This "default tunnel" is actually a default tunnel.
682                    // Try connecting so we have a valid rewrite for thsi broker address.
683                    TunnelConfig::Ssh(default_tunnel) => {
684                        // Multiple users could all run `connect` at the same time; only one ssh
685                        // tunnel will ever be connected, and only one will be inserted into the
686                        // map.
687                        let ssh_tunnel = self.runtime.block_on(async {
688                            self.ssh_tunnel_manager
689                                .connect(
690                                    default_tunnel.clone(),
691                                    host,
692                                    port,
693                                    self.ssh_timeout_config,
694                                    self.in_task,
695                                )
696                                .await
697                        });
698                        match ssh_tunnel {
699                            // Use the tunnel we just created, but only if nobody beat us in the race.
700                            Ok(ssh_tunnel) => {
701                                let mut rewrites = self.rewrites.lock().expect("poisoned");
702                                let rewrite = match rewrites.entry(addr.clone()) {
703                                    btree_map::Entry::Occupied(mut o)
704                                        if matches!(
705                                            o.get(),
706                                            BrokerRewriteHandle::FailedDefaultSshTunnel(_)
707                                        ) =>
708                                    {
709                                        o.insert(BrokerRewriteHandle::SshTunnel(
710                                            ssh_tunnel.clone(),
711                                        ));
712                                        o.into_mut()
713                                    }
714                                    btree_map::Entry::Occupied(o) => o.into_mut(),
715                                    btree_map::Entry::Vacant(v) => {
716                                        v.insert(BrokerRewriteHandle::SshTunnel(ssh_tunnel.clone()))
717                                    }
718                                };
719
720                                return_rewrite(rewrite)
721                            }
722                            // We couldn't connect. Someone else will have to try again.
723                            Err(e) => {
724                                warn!(
725                                    "failed to create ssh tunnel for {:?}: {}",
726                                    addr,
727                                    e.display_with_causes()
728                                );
729
730                                // Write an error if no one else has already written one.
731                                let mut rewrites = self.rewrites.lock().expect("poisoned");
732                                rewrites.entry(addr.clone()).or_insert_with(|| {
733                                    BrokerRewriteHandle::FailedDefaultSshTunnel(
734                                        e.to_string_with_causes(),
735                                    )
736                                });
737
738                                Err(io::Error::new(
739                                    io::ErrorKind::Other,
740                                    "creating SSH tunnel failed",
741                                ))
742                            }
743                        }
744                    }
745                    // Our rewrite strategy is to use a specific host, e.g. a PrivateLink endpoint.
746                    TunnelConfig::StaticHost(host) => (host.as_str(), port)
747                        .to_socket_addrs()
748                        .map(|addrs| addrs.collect()),
749                    // Rewrite according to the routing rules.
750                    TunnelConfig::Rules(rules) => {
751                        // If no rules match, just use the address as-is.
752                        let (resolved, result) = match rules.rewrite(&addr) {
753                            Some(rewritten) => {
754                                let result = rewritten.to_socket_addrs();
755                                (rewritten, result)
756                            }
757                            None => (addr.clone(), self.resolve_direct(&addr)),
758                        };
759                        match result {
760                            Ok(addrs) => {
761                                info!(
762                                    "kafka: resolve_broker_addr {}:{} -> {}:{} resolved to {:?}",
763                                    host, port, resolved.host, resolved.port, addrs,
764                                );
765                                Ok(addrs)
766                            }
767                            Err(e) => {
768                                warn!(
769                                    "kafka: resolve_broker_addr {}:{} -> {}:{} DNS resolution FAILED: {e}",
770                                    host, port, resolved.host, resolved.port,
771                                );
772                                Err(e)
773                            }
774                        }
775                    }
776                    // We leave the broker's address as it is.
777                    TunnelConfig::None => self.resolve_direct(&addr),
778                }
779            }
780            // This broker's address was already rewritten. Reuse the existing rewrite.
781            Some(rewrite) => {
782                info!(
783                    "kafka: resolve_broker_addr {}:{} using cached rewrite",
784                    host, port
785                );
786                return_rewrite(&rewrite)
787            }
788        }
789    }
790
791    fn log(&self, level: RDKafkaLogLevel, fac: &str, log_message: &str) {
792        self.inner.log(level, fac, log_message)
793    }
794
795    fn error(&self, error: KafkaError, reason: &str) {
796        self.inner.error(error, reason)
797    }
798
799    fn stats(&self, statistics: Statistics) {
800        self.inner.stats(statistics)
801    }
802
803    fn stats_raw(&self, statistics: &[u8]) {
804        self.inner.stats_raw(statistics)
805    }
806}
807
808impl<C> ConsumerContext for TunnelingClientContext<C>
809where
810    C: ConsumerContext,
811{
812    fn rebalance(
813        &self,
814        native_client: &NativeClient,
815        err: RDKafkaRespErr,
816        tpl: &mut TopicPartitionList,
817    ) {
818        self.inner.rebalance(native_client, err, tpl)
819    }
820
821    fn pre_rebalance<'a>(&self, rebalance: &Rebalance<'a>) {
822        self.inner.pre_rebalance(rebalance)
823    }
824
825    fn post_rebalance<'a>(&self, rebalance: &Rebalance<'a>) {
826        self.inner.post_rebalance(rebalance)
827    }
828
829    fn commit_callback(&self, result: KafkaResult<()>, offsets: &TopicPartitionList) {
830        self.inner.commit_callback(result, offsets)
831    }
832
833    fn main_queue_min_poll_interval(&self) -> Timeout {
834        self.inner.main_queue_min_poll_interval()
835    }
836}
837
838impl<C> ProducerContext for TunnelingClientContext<C>
839where
840    C: ProducerContext,
841{
842    type DeliveryOpaque = C::DeliveryOpaque;
843
844    fn delivery(
845        &self,
846        delivery_result: &DeliveryResult<'_>,
847        delivery_opaque: Self::DeliveryOpaque,
848    ) {
849        self.inner.delivery(delivery_result, delivery_opaque)
850    }
851}
852
853/// Id of a partition in a topic.
854pub type PartitionId = i32;
855
856/// The error returned by [`get_partitions`].
857#[derive(Debug, thiserror::Error)]
858pub enum GetPartitionsError {
859    /// The specified topic does not exist.
860    #[error("Topic does not exist")]
861    TopicDoesNotExist,
862    /// A Kafka error.
863    #[error(transparent)]
864    Kafka(#[from] KafkaError),
865    /// An unstructured error.
866    #[error(transparent)]
867    Other(#[from] anyhow::Error),
868}
869
870/// Retrieve number of partitions for a given `topic` using the given `client`
871pub fn get_partitions<C: ClientContext>(
872    client: &Client<C>,
873    topic: &str,
874    timeout: Duration,
875) -> Result<Vec<PartitionId>, GetPartitionsError> {
876    let meta = client.fetch_metadata(Some(topic), timeout)?;
877    if meta.topics().len() != 1 {
878        Err(anyhow!(
879            "topic {} has {} metadata entries; expected 1",
880            topic,
881            meta.topics().len()
882        ))?;
883    }
884
885    fn check_err(err: Option<RDKafkaRespErr>) -> Result<(), GetPartitionsError> {
886        match err.map(RDKafkaErrorCode::from) {
887            Some(RDKafkaErrorCode::UnknownTopic | RDKafkaErrorCode::UnknownTopicOrPartition) => {
888                Err(GetPartitionsError::TopicDoesNotExist)
889            }
890            Some(code) => Err(anyhow!(code))?,
891            None => Ok(()),
892        }
893    }
894
895    let meta_topic = meta.topics().into_element();
896    check_err(meta_topic.error())?;
897
898    if meta_topic.name() != topic {
899        Err(anyhow!(
900            "got results for wrong topic {} (expected {})",
901            meta_topic.name(),
902            topic
903        ))?;
904    }
905
906    let mut partition_ids = Vec::with_capacity(meta_topic.partitions().len());
907    for partition_meta in meta_topic.partitions() {
908        check_err(partition_meta.error())?;
909
910        partition_ids.push(partition_meta.id());
911    }
912
913    if partition_ids.len() == 0 {
914        Err(GetPartitionsError::TopicDoesNotExist)?;
915    }
916
917    Ok(partition_ids)
918}
919
920/// Default to true as they have no downsides <https://github.com/confluentinc/librdkafka/issues/283>.
921pub const DEFAULT_KEEPALIVE: bool = true;
922/// The `rdkafka` default.
923/// - <https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md>
924pub const DEFAULT_SOCKET_TIMEOUT: Duration = Duration::from_secs(60);
925/// Increased from the rdkafka default
926/// - <https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md>
927pub const DEFAULT_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(600);
928/// The `rdkafka` default.
929/// - <https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md>
930pub const DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT: Duration = Duration::from_secs(30);
931/// A reasonable default timeout when fetching metadata or partitions.
932pub const DEFAULT_FETCH_METADATA_TIMEOUT: Duration = Duration::from_secs(10);
933/// The timeout for reading records from the progress topic. Set to something slightly longer than
934/// the idle transaction timeout (60s) to wait out any stuck producers.
935pub const DEFAULT_PROGRESS_RECORD_FETCH_TIMEOUT: Duration = Duration::from_secs(90);
936
937/// Configurable timeouts for Kafka connections.
938#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
939pub struct TimeoutConfig {
940    /// Whether or not to enable
941    pub keepalive: bool,
942    /// The timeout for network requests. Can't be more than 100ms longer than
943    /// `transaction_timeout.
944    pub socket_timeout: Duration,
945    /// The timeout for transactions.
946    pub transaction_timeout: Duration,
947    /// The timeout for setting up network connections.
948    pub socket_connection_setup_timeout: Duration,
949    /// The timeout for fetching metadata from upstream.
950    pub fetch_metadata_timeout: Duration,
951    /// The timeout for reading records from the progress topic.
952    pub progress_record_fetch_timeout: Duration,
953}
954
955impl Default for TimeoutConfig {
956    fn default() -> Self {
957        TimeoutConfig {
958            keepalive: DEFAULT_KEEPALIVE,
959            socket_timeout: DEFAULT_SOCKET_TIMEOUT,
960            transaction_timeout: DEFAULT_TRANSACTION_TIMEOUT,
961            socket_connection_setup_timeout: DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT,
962            fetch_metadata_timeout: DEFAULT_FETCH_METADATA_TIMEOUT,
963            progress_record_fetch_timeout: DEFAULT_PROGRESS_RECORD_FETCH_TIMEOUT,
964        }
965    }
966}
967
968impl TimeoutConfig {
969    /// Build a `TcpTimeoutConfig` from the given parameters. Parameters outside the supported
970    /// range are defaulted and cause an error log.
971    pub fn build(
972        keepalive: bool,
973        socket_timeout: Option<Duration>,
974        transaction_timeout: Duration,
975        socket_connection_setup_timeout: Duration,
976        fetch_metadata_timeout: Duration,
977        progress_record_fetch_timeout: Option<Duration>,
978    ) -> TimeoutConfig {
979        // Constrain values based on ranges here:
980        // <https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md>
981        //
982        // Note we error log but do not fail as this is called in a non-fallible
983        // LD-sync in the adapter.
984
985        let transaction_timeout = if transaction_timeout.as_millis() > i32::MAX.try_into().unwrap()
986        {
987            warn!(
988                "transaction_timeout ({transaction_timeout:?}) greater than max \
989                of {}, defaulting to the default of {DEFAULT_TRANSACTION_TIMEOUT:?}",
990                i32::MAX
991            );
992            DEFAULT_TRANSACTION_TIMEOUT
993        } else if transaction_timeout.as_millis() < 1000 {
994            warn!(
995                "transaction_timeout ({transaction_timeout:?}) less than max \
996                of 1000ms, defaulting to the default of {DEFAULT_TRANSACTION_TIMEOUT:?}"
997            );
998            DEFAULT_TRANSACTION_TIMEOUT
999        } else {
1000            transaction_timeout
1001        };
1002
1003        let progress_record_fetch_timeout_derived_default =
1004            std::cmp::max(transaction_timeout, DEFAULT_PROGRESS_RECORD_FETCH_TIMEOUT);
1005        let progress_record_fetch_timeout =
1006            progress_record_fetch_timeout.unwrap_or(progress_record_fetch_timeout_derived_default);
1007        let progress_record_fetch_timeout = if progress_record_fetch_timeout < transaction_timeout {
1008            warn!(
1009                "progress record fetch ({progress_record_fetch_timeout:?}) less than transaction \
1010                timeout ({transaction_timeout:?}), defaulting to transaction timeout {transaction_timeout:?}",
1011            );
1012            transaction_timeout
1013        } else {
1014            progress_record_fetch_timeout
1015        };
1016
1017        // The documented max here is `300000`, but rdkafka bans `socket.timeout.ms` being more
1018        // than `transaction.timeout.ms` + 100ms.
1019        let max_socket_timeout = std::cmp::min(
1020            transaction_timeout + Duration::from_millis(100),
1021            Duration::from_secs(300),
1022        );
1023        let socket_timeout_derived_default =
1024            std::cmp::min(max_socket_timeout, DEFAULT_SOCKET_TIMEOUT);
1025        let socket_timeout = socket_timeout.unwrap_or(socket_timeout_derived_default);
1026        let socket_timeout = if socket_timeout > max_socket_timeout {
1027            warn!(
1028                "socket_timeout ({socket_timeout:?}) greater than max \
1029                of min(30000, transaction.timeout.ms + 100 ({})), \
1030                defaulting to the maximum of {max_socket_timeout:?}",
1031                transaction_timeout.as_millis() + 100
1032            );
1033            max_socket_timeout
1034        } else if socket_timeout.as_millis() < 10 {
1035            warn!(
1036                "socket_timeout ({socket_timeout:?}) less than min \
1037                of 10ms, defaulting to the default of {socket_timeout_derived_default:?}"
1038            );
1039            socket_timeout_derived_default
1040        } else {
1041            socket_timeout
1042        };
1043
1044        let socket_connection_setup_timeout =
1045            if socket_connection_setup_timeout.as_millis() > i32::MAX.try_into().unwrap() {
1046                warn!(
1047                    "socket_connection_setup_timeout ({socket_connection_setup_timeout:?}) \
1048                    greater than max of {}ms, defaulting to the default \
1049                    of {DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT:?}",
1050                    i32::MAX,
1051                );
1052                DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT
1053            } else if socket_connection_setup_timeout.as_millis() < 10 {
1054                warn!(
1055                    "socket_connection_setup_timeout ({socket_connection_setup_timeout:?}) \
1056                    less than max of 10ms, defaulting to the default of \
1057                {DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT:?}"
1058                );
1059                DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT
1060            } else {
1061                socket_connection_setup_timeout
1062            };
1063
1064        TimeoutConfig {
1065            keepalive,
1066            socket_timeout,
1067            transaction_timeout,
1068            socket_connection_setup_timeout,
1069            fetch_metadata_timeout,
1070            progress_record_fetch_timeout,
1071        }
1072    }
1073}
1074
1075/// A simpler version of [`create_new_client_config`] that defaults
1076/// the `log_level` to `INFO` and should only be used in tests.
1077pub fn create_new_client_config_simple() -> ClientConfig {
1078    create_new_client_config(tracing::Level::INFO, Default::default())
1079}
1080
1081/// Build a new [`rdkafka`] [`ClientConfig`] with its `log_level` set correctly
1082/// based on the passed through [`tracing::Level`]. This level should be
1083/// determined for `target: "librdkafka"`.
1084pub fn create_new_client_config(
1085    tracing_level: Level,
1086    timeout_config: TimeoutConfig,
1087) -> ClientConfig {
1088    #[allow(clippy::disallowed_methods)]
1089    let mut config = ClientConfig::new();
1090
1091    let level = if tracing_level >= Level::DEBUG {
1092        RDKafkaLogLevel::Debug
1093    } else if tracing_level >= Level::INFO {
1094        RDKafkaLogLevel::Info
1095    } else if tracing_level >= Level::WARN {
1096        RDKafkaLogLevel::Warning
1097    } else {
1098        RDKafkaLogLevel::Error
1099    };
1100    // WARNING WARNING WARNING
1101    //
1102    // For whatever reason, if you change this `target` to something else, this
1103    // log line might break. I (guswynn) did some extensive investigation with
1104    // the tracing folks, and we learned that this edge case only happens with
1105    // 1. a different target
1106    // 2. only this file (so far as I can tell)
1107    // 3. only in certain subscriber combinations
1108    // 4. only if the `tracing-log` feature is on.
1109    //
1110    // Our conclusion was that one of our dependencies is doing something
1111    // problematic with `log`.
1112    //
1113    // For now, this works, and prints a nice log line exactly when we want it.
1114    //
1115    // TODO(guswynn): when we can remove `tracing-log`, remove this warning
1116    tracing::debug!(target: "librdkafka", level = ?level, "Determined log level for librdkafka");
1117    config.set_log_level(level);
1118
1119    // Patch the librdkafka debug log system into the Rust `log` ecosystem. This
1120    // is a very simple integration at the moment; enabling `debug`-level logs
1121    // for the `librdkafka` target enables the full firehouse of librdkafka
1122    // debug logs. We may want to investigate finer-grained control.
1123    if tracing_level >= Level::DEBUG {
1124        tracing::debug!(target: "librdkafka", "Enabling debug logs for rdkafka");
1125        config.set("debug", "all");
1126    }
1127
1128    if timeout_config.keepalive {
1129        config.set("socket.keepalive.enable", "true");
1130    }
1131
1132    config.set(
1133        "socket.timeout.ms",
1134        timeout_config.socket_timeout.as_millis().to_string(),
1135    );
1136    config.set(
1137        "socket.connection.setup.timeout.ms",
1138        timeout_config
1139            .socket_connection_setup_timeout
1140            .as_millis()
1141            .to_string(),
1142    );
1143
1144    config
1145}
1146
1147/// Creates a client from `config` using [`ClientConfig::create`].
1148///
1149/// Leaves the calling thread's OpenSSL error queue empty, see
1150/// [`create_with_context`].
1151pub fn create<T: FromClientConfig>(config: &ClientConfig) -> KafkaResult<T> {
1152    #[allow(clippy::disallowed_methods)]
1153    let client = config.create();
1154    drop(openssl::error::ErrorStack::get());
1155    client
1156}
1157
1158/// Creates a client from `config` using [`ClientConfig::create_with_context`].
1159///
1160/// Leaves the calling thread's OpenSSL error queue empty.
1161pub fn create_with_context<C, T>(config: &ClientConfig, context: C) -> KafkaResult<T>
1162where
1163    C: ClientContext,
1164    T: FromClientConfigAndContext<C>,
1165{
1166    #[allow(clippy::disallowed_methods)]
1167    let client = config.create_with_context(context);
1168    // NOTE: librdkafka loads `ssl.ca.pem` by reading certificates until
1169    // `PEM_read_bio_X509` fails, and leaves that final `PEM routines:get_name:no
1170    // start line` error on the calling thread's OpenSSL error queue. OpenSSL's
1171    // `SSL_get_error` reports any queued error as `SSL_ERROR_SSL`, so a later
1172    // TLS read that would merely block on this thread fails with the stale
1173    // error instead. On a tokio worker this breaks unrelated connections, for
1174    // example HTTPS requests to a schema registry. librdkafka v2.15.1 clears
1175    // the queue itself (confluentinc/librdkafka#5561).
1176    drop(openssl::error::ErrorStack::get());
1177    client
1178}
1179
1180#[cfg(test)]
1181mod tests {
1182    use super::*;
1183
1184    use openssl::asn1::Asn1Time;
1185    use openssl::hash::MessageDigest;
1186    use openssl::pkey::PKey;
1187    use openssl::rsa::Rsa;
1188    use openssl::x509::{X509, X509NameBuilder};
1189    use rdkafka::consumer::BaseConsumer;
1190
1191    fn self_signed_cert_pem() -> String {
1192        let key = PKey::from_rsa(Rsa::generate(2048).unwrap()).unwrap();
1193        let mut name = X509NameBuilder::new().unwrap();
1194        name.append_entry_by_text("CN", "test").unwrap();
1195        let name = name.build();
1196        let mut cert = X509::builder().unwrap();
1197        cert.set_version(2).unwrap();
1198        cert.set_subject_name(&name).unwrap();
1199        cert.set_issuer_name(&name).unwrap();
1200        cert.set_pubkey(&key).unwrap();
1201        cert.set_not_before(&Asn1Time::days_from_now(0).unwrap())
1202            .unwrap();
1203        cert.set_not_after(&Asn1Time::days_from_now(1).unwrap())
1204            .unwrap();
1205        cert.sign(&key, MessageDigest::sha256()).unwrap();
1206        String::from_utf8(cert.build().to_pem().unwrap()).unwrap()
1207    }
1208
1209    #[mz_ore::test]
1210    fn all_brokers_down_parses_regardless_of_counts() {
1211        for msg in ["1/1 brokers are down", "0/2 brokers are down"] {
1212            assert!(
1213                matches!(
1214                    MzKafkaError::from_str(msg),
1215                    Ok(MzKafkaError::AllBrokersDown)
1216                ),
1217                "{msg}"
1218            );
1219        }
1220    }
1221
1222    #[mz_ore::test]
1223    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function
1224    fn resolve_direct_broker_enforces_external_addresses() {
1225        let runtime = tokio::runtime::Runtime::new().unwrap();
1226        let loopback = BrokerAddr {
1227            host: "127.0.0.1".into(),
1228            port: 9092,
1229        };
1230        let no_match = HostMappingRules {
1231            rules: vec![(
1232                ConnectionRulePattern {
1233                    prefix_wildcard: false,
1234                    literal_match: "other:9092".into(),
1235                    suffix_wildcard: false,
1236                },
1237                BrokerRewrite {
1238                    host: "rewritten".into(),
1239                    port: None,
1240                },
1241            )],
1242        };
1243        for tunnel in [TunnelConfig::None, TunnelConfig::Rules(no_match)] {
1244            for enforce in [false, true] {
1245                let (inner, errors) = MzClientContext::with_errors();
1246                let mut cx = TunnelingClientContext::new(
1247                    inner,
1248                    runtime.handle().clone(),
1249                    SshTunnelManager::default(),
1250                    SshTimeoutConfig::default(),
1251                    None,
1252                    InTask::No,
1253                    enforce,
1254                );
1255                cx.set_default_tunnel(tunnel.clone());
1256                let result = cx.resolve_broker_addr(&loopback.host, loopback.port);
1257                if enforce {
1258                    assert!(result.is_err(), "loopback broker should be rejected");
1259                    assert_eq!(
1260                        errors.try_iter().collect::<Vec<_>>(),
1261                        vec![MzKafkaError::PrivateBrokerAddress],
1262                    );
1263                } else {
1264                    assert_eq!(
1265                        result.expect("loopback broker should resolve"),
1266                        loopback.to_socket_addrs().unwrap(),
1267                    );
1268                }
1269            }
1270        }
1271    }
1272
1273    #[mz_ore::test]
1274    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function
1275    fn test_create_with_context_clears_openssl_error_queue() {
1276        let mut config = create_new_client_config_simple();
1277        // Client creation does not connect, so the broker need not exist.
1278        config.set("bootstrap.servers", "localhost:1");
1279        config.set("security.protocol", "ssl");
1280        config.set("ssl.ca.pem", self_signed_cert_pem());
1281        let _consumer: BaseConsumer =
1282            create_with_context(&config, rdkafka::consumer::DefaultConsumerContext).unwrap();
1283        let errors = openssl::error::ErrorStack::get();
1284        assert!(
1285            errors.errors().is_empty(),
1286            "stale OpenSSL errors after client creation: {errors}"
1287        );
1288    }
1289
1290    #[mz_ore::test]
1291    fn test_connection_rule_pattern_matches() {
1292        let p = ConnectionRulePattern {
1293            prefix_wildcard: false,
1294            literal_match: "broker:9092".to_string(),
1295            suffix_wildcard: false,
1296        };
1297        assert!(p.matches("broker:9092"));
1298        assert!(!p.matches("other:9092"));
1299
1300        let p = ConnectionRulePattern {
1301            prefix_wildcard: true,
1302            literal_match: ":9092".to_string(),
1303            suffix_wildcard: false,
1304        };
1305        assert!(p.matches("any-host:9092"));
1306        assert!(!p.matches("broker:9093"));
1307
1308        let p = ConnectionRulePattern {
1309            prefix_wildcard: false,
1310            literal_match: "broker:".to_string(),
1311            suffix_wildcard: true,
1312        };
1313        assert!(p.matches("broker:9092"));
1314        assert!(!p.matches("other:9092"));
1315
1316        let p = ConnectionRulePattern {
1317            prefix_wildcard: true,
1318            literal_match: "broker".to_string(),
1319            suffix_wildcard: true,
1320        };
1321        assert!(p.matches("my-broker-host:1234"));
1322        assert!(!p.matches("other:9092"));
1323    }
1324}
mz_kafka_util/client.rs

mz_kafka_util/
client.rs
