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

mz_kafka_util/
client.rs
