From 62ac00fe19f7a59f931ff3736f97655ed37fcfa5 Mon Sep 17 00:00:00 2001 From: Nikhil Sinha Date: Sun, 19 Jul 2026 18:06:38 +0700 Subject: [PATCH 1/3] Adds a configurable OAuth provider for Kafka SASL/OAUTHBEARER and fixes the MSK IAM token refresh, which panicked at runtime on first refresh. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit | Provider | Token source | Required config | |---|---|---| | `aws-msk` (default when a region resolves) | aws-msk-iam-sasl-signer, app-managed refresh | region via `--aws-region`/`AWS_REGION`, or AWS profile/IMDS | | `oidc` (Google Managed Kafka, generic OIDC) | librdkafka built-in OIDC handler | token endpoint URL, client id, client secret | Existing MSK IAM configs work unchanged; the provider is inferred when unset. - Nested `Handle::block_on` panic in the OAUTHBEARER refresh callback (fired inside the consumer poll loop's `block_on`) — signer now runs on a dedicated thread with its own single-use runtime, 30s timeout. - Invalid Kafka config aborted only with a `warn!` and silently skipped ingestion — now a hard startup failure when Kafka is configured. - AWS region resolution falls back to the SDK default chain (profile/IMDS), matching where credentials come from; resolved once and cached. - Consumer creation no longer logs the full `ClientConfig` (contains secrets). `rdkafka/curl` + `curl-static` added for librdkafka's OIDC handler; `aws-config` added (already transitive via the signer). --- Cargo.lock | 70 +++++++ Cargo.toml | 4 + src/connectors/kafka/config.rs | 306 +++++++++++++++++++++++++++---- src/connectors/kafka/consumer.rs | 7 +- src/connectors/kafka/mod.rs | 134 +++++++++----- src/connectors/mod.rs | 10 +- src/interactive.rs | 120 +++++++++--- 7 files changed, 544 insertions(+), 107 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a8ecd44d0..ac2e9c7ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -783,6 +783,8 @@ checksum = "8a8fc176d53d6fe85017f230405e3255cedb4a02221cb55ed6d76dccbbb099b2" dependencies = [ "aws-credential-types", "aws-runtime", + "aws-sdk-sso", + "aws-sdk-ssooidc", "aws-sdk-sts", "aws-smithy-async", "aws-smithy-http", @@ -793,11 +795,14 @@ dependencies = [ "aws-types", "bytes", "fastrand 2.3.0", + "hex", "http 1.4.0", + "ring", "time", "tokio", "tracing", "url", + "zeroize", ] [[package]] @@ -877,6 +882,54 @@ dependencies = [ "uuid", ] +[[package]] +name = "aws-sdk-sso" +version = "1.95.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00c5ff27c6ba2cbd95e6e26e2e736676fdf6bcf96495b187733f521cfe4ce448" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand 2.3.0", + "http 0.2.12", + "http 1.4.0", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sdk-ssooidc" +version = "1.97.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d186f1e5a3694a188e5a0640b3115ccc6e084d104e16fd6ba968dca072ffef8" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand 2.3.0", + "http 0.2.12", + "http 1.4.0", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-sts" version = "1.99.0" @@ -1788,6 +1841,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "curl-sys" +version = "0.4.90+curl-8.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97799a0d220bfb3361e0fe4936966ff8c4b24d65c3f06dfc70d7b680b44e7897" +dependencies = [ + "cc", + "libc", + "libz-sys", + "openssl-sys", + "pkg-config", + "vcpkg", + "windows-sys 0.59.0", +] + [[package]] name = "darling" version = "0.20.11" @@ -4362,6 +4430,7 @@ dependencies = [ "arrow-schema", "arrow-select", "async-trait", + "aws-config", "aws-msk-iam-sasl-signer", "aws-types", "base64", @@ -5061,6 +5130,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5230dca48bc354d718269f3e4353280e188b610f7af7e2fcf54b7a79d5802872" dependencies = [ "cmake", + "curl-sys", "libc", "libz-sys", "num_enum", diff --git a/Cargo.toml b/Cargo.toml index 538b83e8e..18766c2a7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,6 +59,7 @@ sasl2-sys = { version = "0.1.22", optional = true, features = ["vendored"] } # AWS MSK IAM authentication for Kafka (SASL/OAUTHBEARER token signing) aws-msk-iam-sasl-signer = { version = "1.0.1", optional = true } aws-types = { version = "1", optional = true } +aws-config = { version = "1", optional = true } # Authentication and Security argon2 = "0.5.0" @@ -240,10 +241,13 @@ kafka = [ "rdkafka/ssl-vendored", "rdkafka/ssl", "rdkafka/sasl", + "rdkafka/curl", + "rdkafka/curl-static", "sasl2-sys", "sasl2-sys/vendored", "aws-msk-iam-sasl-signer", "aws-types", + "aws-config", ] [profile.release-lto] diff --git a/src/connectors/kafka/config.rs b/src/connectors/kafka/config.rs index 738819b74..a5ed0ec63 100644 --- a/src/connectors/kafka/config.rs +++ b/src/connectors/kafka/config.rs @@ -450,7 +450,6 @@ pub struct ProducerConfig { #[derive(Debug, Clone, Args)] pub struct SecurityConfig { #[clap( - value_enum, long = "security-protocol", env = "P_KAFKA_SECURITY_PROTOCOL", required = false, @@ -486,7 +485,6 @@ pub struct SecurityConfig { // SASL Configuration #[arg( - value_enum, long = "sasl-mechanism", env = "P_KAFKA_SASL_MECHANISM", required = false, @@ -510,12 +508,45 @@ pub struct SecurityConfig { )] pub sasl_password: Option, + // SASL/OAUTHBEARER provider configuration + #[arg( + long = "oauth-provider", + env = "P_KAFKA_OAUTH_PROVIDER", + required = false, + help = "OAuth provider: aws-msk for AWS IAM signing, or oidc for a standard token endpoint such as the Google Managed Kafka local auth server" + )] + pub oauth_provider: Option, + + #[arg( + long = "oauth-token-endpoint-url", + env = "P_KAFKA_OAUTH_TOKEN_ENDPOINT_URL", + required = false, + help = "OAuth/OIDC token endpoint URL. Required for the oidc provider" + )] + pub oauth_token_endpoint_url: Option, + + #[arg( + long = "oauth-client-id", + env = "P_KAFKA_OAUTH_CLIENT_ID", + required = false, + help = "OAuth/OIDC client ID. Google Managed Kafka's local auth server accepts 'unused'" + )] + pub oauth_client_id: Option, + + #[arg( + long = "oauth-client-secret", + env = "P_KAFKA_OAUTH_CLIENT_SECRET", + required = false, + help = "OAuth/OIDC client secret. Google Managed Kafka's local auth server accepts 'unused'" + )] + pub oauth_client_secret: Option, + // AWS MSK IAM configuration (SASL/OAUTHBEARER) #[arg( long = "aws-region", env = "P_KAFKA_AWS_REGION", required = false, - help = "AWS region for MSK IAM authentication (SASL/OAUTHBEARER). Falls back to AWS_REGION / AWS_DEFAULT_REGION when unset." + help = "AWS region for MSK IAM authentication (SASL/OAUTHBEARER). Falls back to AWS_REGION / AWS_DEFAULT_REGION, then the AWS SDK default region chain (profile files, IMDS) when unset." )] pub aws_region: Option, @@ -806,13 +837,50 @@ impl SecurityConfig { } } - // OAUTHBEARER (used for AWS MSK IAM) requires no static credentials here. - // The token is minted on demand by `KafkaContext::generate_oauth_token`, - // which rdkafka invokes via the OAuth refresh callback. Setting the - // `sasl.mechanism` above is sufficient to enable that callback path. + if matches!(self.sasl_mechanism, Some(SaslMechanism::OAuthBearer)) + && matches!(self.resolved_oauth_provider(), Some(OAuthProvider::Oidc)) + { + // librdkafka's built-in OIDC handler fetches and refreshes tokens + // from the configured endpoint. Google Managed Kafka's local auth + // server implements this endpoint using Application Default + // Credentials. + config.set("sasl.oauthbearer.method", "oidc"); + if let Some(ref endpoint) = self.oauth_token_endpoint_url { + config.set("sasl.oauthbearer.token.endpoint.url", endpoint); + } + if let Some(ref client_id) = self.oauth_client_id { + config.set("sasl.oauthbearer.client.id", client_id); + } + if let Some(ref client_secret) = self.oauth_client_secret { + config.set("sasl.oauthbearer.client.secret", client_secret); + } + } } } + /// Resolves the token provider while preserving the original AWS-only + /// configuration. An explicit provider wins, followed by an OIDC endpoint, + /// then an AWS region. + pub fn resolved_oauth_provider(&self) -> Option { + let region = self.resolved_aws_region(); + self.resolved_oauth_provider_with_region(region.as_deref()) + } + + fn resolved_oauth_provider_with_region( + &self, + resolved_region: Option<&str>, + ) -> Option { + self.oauth_provider.or_else(|| { + if Self::has_value(&self.oauth_token_endpoint_url) { + Some(OAuthProvider::Oidc) + } else if resolved_region.is_some() { + Some(OAuthProvider::AwsMsk) + } else { + None + } + }) + } + /// Resolves the AWS region to use for MSK IAM authentication. /// /// Precedence: the explicit `--aws-region` flag, then the standard @@ -820,16 +888,25 @@ impl SecurityConfig { /// normalized (trimmed and rejected if empty) before falling through, so an /// explicitly-empty flag does not shadow a valid environment variable. pub fn resolved_aws_region(&self) -> Option { - Self::normalize_region(self.aws_region.clone()) - .or_else(|| Self::normalize_region(std::env::var(AWS_REGION_ENV).ok())) - .or_else(|| Self::normalize_region(std::env::var(AWS_DEFAULT_REGION_ENV).ok())) + Self::normalize_region(self.aws_region.as_deref()) + .or_else(|| Self::normalize_region(std::env::var(AWS_REGION_ENV).ok().as_deref())) + .or_else(|| { + Self::normalize_region(std::env::var(AWS_DEFAULT_REGION_ENV).ok().as_deref()) + }) } /// Trims a candidate region value and discards it if it is empty. - fn normalize_region(region: Option) -> Option { + fn normalize_region(region: Option<&str>) -> Option { region - .map(|region| region.trim().to_owned()) + .map(str::trim) .filter(|region| !region.is_empty()) + .map(str::to_owned) + } + + fn has_value(value: &Option) -> bool { + value + .as_deref() + .is_some_and(|value| !value.trim().is_empty()) } fn validate(&self) -> anyhow::Result<()> { @@ -868,20 +945,39 @@ impl SecurityConfig { match self.sasl_mechanism { None => anyhow::bail!("SASL mechanism is required when SASL is enabled"), Some(SaslMechanism::OAuthBearer) => { - // AWS MSK IAM requires SASL_SSL: the OAUTHBEARER token is a - // signed bearer credential, so transmitting it over an - // unencrypted SASL_PLAINTEXT connection would expose it. if matches!(self.protocol, SecurityProtocol::SaslPlaintext) { anyhow::bail!( - "SASL/OAUTHBEARER (MSK IAM) requires the SASL_SSL security protocol; SASL_PLAINTEXT would expose the bearer token" + "SASL/OAUTHBEARER requires the SASL_SSL security protocol; SASL_PLAINTEXT would expose the bearer token" ); } - // AWS MSK IAM mints tokens from the ambient AWS credential - // chain, so only a resolvable region is required here. - if resolved_region.is_none() { - anyhow::bail!( - "AWS region is required for SASL/OAUTHBEARER (MSK IAM); set --aws-region or the AWS_REGION environment variable" - ); + + match self.resolved_oauth_provider_with_region(resolved_region) { + Some(OAuthProvider::AwsMsk) => { + // No hard region requirement: when unset here, token + // generation falls back to the AWS SDK default + // region chain (profile files, IMDS) — the same + // sources that supply the credentials. + } + Some(OAuthProvider::Oidc) => { + if !Self::has_value(&self.oauth_token_endpoint_url) { + anyhow::bail!( + "OAuth token endpoint URL is required for the oidc OAuth provider" + ); + } + if !Self::has_value(&self.oauth_client_id) { + anyhow::bail!( + "OAuth client ID is required for the oidc OAuth provider" + ); + } + if !Self::has_value(&self.oauth_client_secret) { + anyhow::bail!( + "OAuth client secret is required for the oidc OAuth provider" + ); + } + } + None => anyhow::bail!( + "OAuth provider is required for SASL/OAUTHBEARER; set --oauth-provider to aws-msk or oidc" + ), } } Some(SaslMechanism::Gssapi) => { @@ -903,7 +999,7 @@ impl SecurityConfig { } } -#[derive(ValueEnum, Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum SecurityProtocol { Plaintext, Ssl, @@ -918,8 +1014,8 @@ impl std::str::FromStr for SecurityProtocol { match s.to_uppercase().as_str() { "PLAINTEXT" => Ok(SecurityProtocol::Plaintext), "SSL" => Ok(SecurityProtocol::Ssl), - "SASL_SSL" => Ok(SecurityProtocol::SaslSsl), - "SASL_PLAINTEXT" => Ok(SecurityProtocol::SaslPlaintext), + "SASL_SSL" | "SASL-SSL" => Ok(SecurityProtocol::SaslSsl), + "SASL_PLAINTEXT" | "SASL-PLAINTEXT" => Ok(SecurityProtocol::SaslPlaintext), _ => Err(format!("Invalid security protocol: {s}")), } } @@ -936,7 +1032,7 @@ impl std::fmt::Display for SecurityProtocol { } } -#[derive(ValueEnum, Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum SaslMechanism { Plain, ScramSha256, @@ -945,6 +1041,33 @@ pub enum SaslMechanism { OAuthBearer, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum OAuthProvider { + AwsMsk, + Oidc, +} + +impl std::str::FromStr for OAuthProvider { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_ascii_lowercase().as_str() { + "aws-msk" | "aws_msk" | "aws" => Ok(Self::AwsMsk), + "oidc" | "gcp" | "gcp-managed-kafka" => Ok(Self::Oidc), + _ => Err(format!("Invalid OAuth provider: {s}")), + } + } +} + +impl std::fmt::Display for OAuthProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::AwsMsk => write!(f, "aws-msk"), + Self::Oidc => write!(f, "oidc"), + } + } +} + impl Default for KafkaConfig { fn default() -> Self { Self { @@ -1040,6 +1163,10 @@ impl Default for SecurityConfig { sasl_mechanism: None, sasl_username: None, sasl_password: None, + oauth_provider: None, + oauth_token_endpoint_url: None, + oauth_client_id: None, + oauth_client_secret: None, aws_region: None, ssl_key_password: None, kerberos_service_name: None, @@ -1058,7 +1185,7 @@ impl std::str::FromStr for SaslMechanism { "SCRAM-SHA-256" => Ok(SaslMechanism::ScramSha256), "SCRAM-SHA-512" => Ok(SaslMechanism::ScramSha512), "GSSAPI" => Ok(SaslMechanism::Gssapi), - "OAUTHBEARER" => Ok(SaslMechanism::OAuthBearer), + "OAUTHBEARER" | "OAUTH-BEARER" | "O-AUTH-BEARER" => Ok(SaslMechanism::OAuthBearer), _ => Err(format!("Invalid SASL mechanism: {s}")), } } @@ -1105,17 +1232,65 @@ impl SourceOffset { mod tests { use super::*; + #[derive(Parser)] + struct SecurityCli { + #[command(flatten)] + security: SecurityConfig, + } + #[test] - fn oauthbearer_requires_a_region() { + fn standard_kafka_security_values_parse() { + let cli = SecurityCli::try_parse_from([ + "test", + "--security-protocol", + "SASL_SSL", + "--sasl-mechanism", + "OAUTHBEARER", + "--oauth-provider", + "oidc", + "--oauth-token-endpoint-url", + "http://localhost:14293", + "--oauth-client-id", + "unused", + "--oauth-client-secret", + "unused", + ]) + .expect("standard Kafka values should parse"); + + assert!(matches!(cli.security.protocol, SecurityProtocol::SaslSsl)); + assert!(matches!( + cli.security.sasl_mechanism, + Some(SaslMechanism::OAuthBearer) + )); + assert_eq!(cli.security.oauth_provider, Some(OAuthProvider::Oidc)); + } + + #[test] + fn aws_msk_without_region_defers_to_sdk_chain() { + // An explicit aws-msk provider without a region is valid: token + // generation resolves the region via the AWS SDK default chain + // (profile files, IMDS) at runtime. let security = SecurityConfig { protocol: SecurityProtocol::SaslSsl, sasl_mechanism: Some(SaslMechanism::OAuthBearer), + oauth_provider: Some(OAuthProvider::AwsMsk), aws_region: None, ..Default::default() }; - // Validate against an explicitly-unresolved region so the failure path - // is exercised regardless of the ambient AWS_REGION environment. + assert!(security.validate_with_region(None).is_ok()); + } + + #[test] + fn oauthbearer_requires_a_provider() { + // With no explicit provider, no OIDC endpoint, and no resolvable + // region, the provider cannot be inferred and validation must fail. + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + ..Default::default() + }; + assert!(security.validate_with_region(None).is_err()); } @@ -1135,14 +1310,11 @@ mod tests { #[test] fn normalize_region_trims_and_rejects_empty() { assert_eq!( - SecurityConfig::normalize_region(Some(" us-east-1 ".to_string())).as_deref(), + SecurityConfig::normalize_region(Some(" us-east-1 ")).as_deref(), Some("us-east-1") ); - assert_eq!( - SecurityConfig::normalize_region(Some(" ".to_string())), - None - ); - assert_eq!(SecurityConfig::normalize_region(Some(String::new())), None); + assert_eq!(SecurityConfig::normalize_region(Some(" ")), None); + assert_eq!(SecurityConfig::normalize_region(Some("")), None); assert_eq!(SecurityConfig::normalize_region(None), None); } @@ -1157,6 +1329,41 @@ mod tests { assert!(security.validate().is_ok()); assert_eq!(security.resolved_aws_region().as_deref(), Some("us-east-1")); + assert_eq!( + security.resolved_oauth_provider(), + Some(OAuthProvider::AwsMsk) + ); + } + + #[test] + fn oidc_oauthbearer_does_not_require_an_aws_region() { + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + oauth_provider: Some(OAuthProvider::Oidc), + oauth_token_endpoint_url: Some("http://localhost:14293".to_string()), + oauth_client_id: Some("unused".to_string()), + oauth_client_secret: Some("unused".to_string()), + aws_region: None, + ..Default::default() + }; + + assert!(security.validate_with_region(None).is_ok()); + } + + #[test] + fn oidc_oauthbearer_requires_complete_endpoint_configuration() { + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + oauth_provider: Some(OAuthProvider::Oidc), + oauth_token_endpoint_url: Some("http://localhost:14293".to_string()), + oauth_client_id: Some("unused".to_string()), + oauth_client_secret: None, + ..Default::default() + }; + + assert!(security.validate_with_region(None).is_err()); } #[test] @@ -1221,5 +1428,32 @@ mod tests { // No static credentials should be set for OAUTHBEARER. assert_eq!(config.get("sasl.username"), None); assert_eq!(config.get("sasl.password"), None); + assert_eq!(config.get("sasl.oauthbearer.method"), None); + } + + #[test] + fn oidc_oauthbearer_applies_librdkafka_configuration() { + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + oauth_provider: Some(OAuthProvider::Oidc), + oauth_token_endpoint_url: Some("http://localhost:14293".to_string()), + oauth_client_id: Some("unused".to_string()), + oauth_client_secret: Some("unused".to_string()), + ..Default::default() + }; + + let mut config = ClientConfig::new(); + security.apply_to_config(&mut config); + + assert_eq!(config.get("security.protocol"), Some("SASL_SSL")); + assert_eq!(config.get("sasl.mechanism"), Some("OAUTHBEARER")); + assert_eq!(config.get("sasl.oauthbearer.method"), Some("oidc")); + assert_eq!( + config.get("sasl.oauthbearer.token.endpoint.url"), + Some("http://localhost:14293") + ); + assert_eq!(config.get("sasl.oauthbearer.client.id"), Some("unused")); + assert_eq!(config.get("sasl.oauthbearer.client.secret"), Some("unused")); } } diff --git a/src/connectors/kafka/consumer.rs b/src/connectors/kafka/consumer.rs index cdb811948..9b4deb544 100644 --- a/src/connectors/kafka/consumer.rs +++ b/src/connectors/kafka/consumer.rs @@ -210,10 +210,9 @@ impl KafkaStreams { fn create_consumer(context: KafkaContext) -> Arc { let kafka_config = &context.config; let consumer_config = kafka_config.to_rdkafka_consumer_config(); - info!( - "Creating Kafka consumer from configs: {:#?}", - &consumer_config - ); + // ClientConfig can contain SASL passwords and OAuth client secrets. + // Never emit the full configuration into logs. + info!("Creating Kafka consumer from configured settings"); let consumer: StreamConsumer = consumer_config .create_with_context(context.clone()) diff --git a/src/connectors/kafka/mod.rs b/src/connectors/kafka/mod.rs index 5960a96a3..90548f547 100644 --- a/src/connectors/kafka/mod.rs +++ b/src/connectors/kafka/mod.rs @@ -16,8 +16,9 @@ * */ -use crate::connectors::kafka::config::{KafkaConfig, SaslMechanism}; +use crate::connectors::kafka::config::{KafkaConfig, OAuthProvider, SaslMechanism}; use crate::handlers::TENANT_ID; +use aws_config::default_provider::region::DefaultRegionChain; use aws_msk_iam_sasl_signer::generate_auth_token; use aws_types::region::Region; use derive_more::Constructor; @@ -31,8 +32,8 @@ use rdkafka::{ClientContext, Message, Offset, Statistics}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::error::Error; -use std::sync::{Arc, RwLock}; -use tokio::runtime::Handle; +use std::sync::{Arc, OnceLock, RwLock}; +use std::time::Duration; use tokio::sync::mpsc; use tokio::sync::mpsc::Receiver; use tracing::{error, info, warn}; @@ -51,18 +52,23 @@ type BaseConsumer = rdkafka::consumer::BaseConsumer; type FutureProducer = rdkafka::producer::FutureProducer; type StreamConsumer = rdkafka::consumer::StreamConsumer; +/// Principal name reported to librdkafka for MSK IAM tokens. MSK does not +/// consume it, but librdkafka requires a non-empty value. +const MSK_OAUTH_PRINCIPAL: &str = "parseable"; + +/// Upper bound on a single MSK IAM token generation, covering the AWS +/// credential chain resolution (which may involve IMDS/STS round-trips). +const MSK_TOKEN_GENERATION_TIMEOUT: Duration = Duration::from_secs(30); + #[derive(Clone, Debug)] pub struct KafkaContext { config: Arc, statistics: Arc>, rebalance_tx: mpsc::Sender, - /// Handle to the Tokio runtime, captured at construction time. - /// - /// librdkafka invokes the OAuth token refresh callback from one of its own - /// background threads, which is not a Tokio runtime thread. The MSK IAM - /// token signer is async, so we use this handle to drive it to completion - /// from that synchronous callback. - runtime_handle: Handle, + /// AWS region for MSK IAM token signing, resolved once on first use. + /// Resolution can involve the AWS SDK default region chain (profile + /// files, IMDS), so the result is cached for subsequent token refreshes. + msk_region: Arc>, } impl KafkaContext { @@ -74,7 +80,7 @@ impl KafkaContext { config, statistics, rebalance_tx, - runtime_handle: Handle::current(), + msk_region: Arc::new(OnceLock::new()), }, rebalance_rx, ) @@ -262,9 +268,9 @@ impl ProducerContext for KafkaContext { } impl ClientContext for KafkaContext { - // Enables librdkafka's OAuth token refresh callback. librdkafka only invokes - // `generate_oauth_token` when the configured SASL mechanism is OAUTHBEARER, - // so leaving this on has no effect for other mechanisms. + // Handles application-managed OAuth refresh events for AWS MSK. When the + // OIDC method is configured, librdkafka uses its own background token + // refresh handler instead and does not enqueue these events. const ENABLE_REFRESH_OAUTH_TOKEN: bool = true; fn stats(&self, new_stats: Statistics) { @@ -280,47 +286,89 @@ impl ClientContext for KafkaContext { /// Generates a SASL/OAUTHBEARER token for AWS MSK IAM authentication. /// - /// librdkafka calls this synchronously from one of its background threads - /// whenever it needs a new token (initially and before expiry). The AWS MSK - /// IAM signer is async, so we drive it to completion on the captured Tokio - /// runtime handle. This thread is not a runtime worker, so blocking on it is - /// safe and does not stall async tasks. + /// librdkafka calls this synchronously whenever it needs a new token + /// (initially and before expiry). The invocation can happen on a thread + /// that is already executing inside a Tokio runtime — the consumer poll + /// loop drives these events from within `block_on` — where nesting + /// `Handle::block_on` panics. The async signer therefore runs on a + /// dedicated short-lived thread with its own single-use runtime: safe from + /// any calling thread, and at one refresh per ~15 minutes the overhead is + /// negligible. fn generate_oauth_token( &self, _oauthbearer_config: Option<&str>, ) -> Result> { - let security = self.config.security(); - - // This callback is only wired up for the OAUTHBEARER mechanism, which we - // use exclusively for AWS MSK IAM. Guard against misconfiguration. - if !matches!( - security.and_then(|s| s.sasl_mechanism.clone()), - Some(SaslMechanism::OAuthBearer) - ) { + let Some(security) = self.config.security() else { + return Err("Kafka security configuration is missing".into()); + }; + + // The application-managed callback is only used by the AWS MSK provider. + // OIDC providers, including Google Managed Kafka's local auth server, + // are refreshed by librdkafka's built-in OIDC handler. + if !matches!(security.sasl_mechanism, Some(SaslMechanism::OAuthBearer)) + || !matches!( + security.resolved_oauth_provider(), + Some(OAuthProvider::AwsMsk) + ) + { return Err( - "OAuth token generation is only supported for the OAUTHBEARER (AWS MSK IAM) SASL mechanism" + "Application-managed OAuth token generation is only supported for the aws-msk provider" .into(), ); } - let region = security - .and_then(|s| s.resolved_aws_region()) - .ok_or("AWS region is not configured for MSK IAM authentication")?; - - info!("Generating AWS MSK IAM OAuth token for region {region}"); + let explicit_region = security.resolved_aws_region().map(Region::new); + let cached_region = Arc::clone(&self.msk_region); + + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::Builder::new() + .name("msk-iam-token".to_string()) + .spawn(move || { + let result = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|e| format!("Failed to build MSK IAM signer runtime: {e}")) + .and_then(|runtime| { + runtime.block_on(async { + let region = match cached_region.get() { + Some(region) => region.clone(), + None => { + let region = match explicit_region { + Some(region) => region, + // Fall back to the AWS SDK default region + // chain (profile files, IMDS) — the same + // sources that supply the credentials. + None => DefaultRegionChain::builder() + .build() + .region() + .await + .ok_or_else(|| { + "Unable to resolve an AWS region for MSK IAM; set --aws-region or AWS_REGION" + .to_string() + })?, + }; + cached_region.get_or_init(|| region).clone() + } + }; + + info!("Generating AWS MSK IAM OAuth token for region {region}"); + + generate_auth_token(region).await.map_err(|e| { + format!("Failed to generate AWS MSK IAM auth token: {e}") + }) + }) + }); + // The receiver is dropped if the caller timed out; nothing to do. + let _ = tx.send(result); + }) + .map_err(|e| format!("Failed to spawn MSK IAM token thread: {e}"))?; - // Drive the async signer from this synchronous, non-runtime thread. - let (token, expiration_time_ms) = self - .runtime_handle - .block_on(generate_auth_token(Region::new(region))) - .map_err(|e| -> Box { - format!("Failed to generate AWS MSK IAM auth token: {e}").into() - })?; + let (token, expiration_time_ms) = rx + .recv_timeout(MSK_TOKEN_GENERATION_TIMEOUT) + .map_err(|e| format!("MSK IAM token generation did not complete: {e}"))??; Ok(OAuthToken { - // MSK does not consume the principal name, but librdkafka requires a - // non-empty value. - principal_name: "parseable".to_string(), + principal_name: MSK_OAUTH_PRINCIPAL.to_string(), token, lifetime_ms: expiration_time_ms, }) diff --git a/src/connectors/mod.rs b/src/connectors/mod.rs index b53119085..094039420 100644 --- a/src/connectors/mod.rs +++ b/src/connectors/mod.rs @@ -38,7 +38,15 @@ pub async fn init(prometheus: &PrometheusMetrics) -> anyhow::Result<()> { if matches!(PARSEABLE.options.mode, Mode::Ingest | Mode::All) { match PARSEABLE.kafka_config.validate() { Err(e) => { - warn!("Kafka connector configuration invalid. {}", e); + // A misconfigured connector must abort startup loudly rather + // than silently skipping ingestion. Only a completely absent + // Kafka configuration is allowed to pass through. + let kafka_configured = PARSEABLE.kafka_config.bootstrap_servers.is_some() + || std::env::vars().any(|(key, _)| key.starts_with("P_KAFKA_")); + if kafka_configured { + anyhow::bail!("Kafka connector configuration invalid: {e}"); + } + warn!("Kafka connector not configured, skipping. ({e})"); } Ok(_) => { let config = PARSEABLE.kafka_config.clone(); diff --git a/src/interactive.rs b/src/interactive.rs index 58ebd78c6..43e45d3eb 100644 --- a/src/interactive.rs +++ b/src/interactive.rs @@ -406,9 +406,8 @@ fn get_oidc_prompts() -> Vec { /// Kafka has layered dependencies: /// - If any `P_KAFKA_*` env is set, `P_KAFKA_BOOTSTRAP_SERVERS` and /// `P_KAFKA_CONSUMER_TOPICS` are required for the server to function. -/// - If security protocol is SSL or SASL_SSL, SSL cert paths are required. -/// - If security protocol is SASL_PLAINTEXT or SASL_SSL, SASL credentials -/// are required. +/// - Pure SSL requires client certificate material; SASL_SSL does not. +/// - SASL requirements depend on the selected mechanism and OAuth provider. #[cfg(feature = "kafka")] fn get_kafka_prompts() -> Vec { // Check if any Kafka env var is set @@ -439,13 +438,13 @@ fn get_kafka_prompts() -> Vec { .unwrap_or_default() .to_uppercase(); - let needs_ssl = matches!(protocol.as_str(), "SSL" | "SASL_SSL" | "SASL-SSL"); + let needs_client_cert = protocol == "SSL"; let needs_sasl = matches!( protocol.as_str(), "SASL_PLAINTEXT" | "SASL-PLAINTEXT" | "SASL_SSL" | "SASL-SSL" ); - if needs_ssl { + if needs_client_cert { prompts.extend([ EnvPrompt { env_var: "P_KAFKA_SSL_CA_LOCATION", @@ -469,26 +468,101 @@ fn get_kafka_prompts() -> Vec { } if needs_sasl { - prompts.extend([ - EnvPrompt { - env_var: "P_KAFKA_SASL_MECHANISM", - display_name: "Kafka SASL Mechanism (PLAIN, SCRAM-SHA-256, SCRAM-SHA-512, GSSAPI)", - required: true, - is_secret: false, - }, - EnvPrompt { - env_var: "P_KAFKA_SASL_USERNAME", - display_name: "Kafka SASL Username", + prompts.push(EnvPrompt { + env_var: "P_KAFKA_SASL_MECHANISM", + display_name: + "Kafka SASL Mechanism (PLAIN, SCRAM-SHA-256, SCRAM-SHA-512, GSSAPI, OAUTHBEARER)", + required: true, + is_secret: false, + }); + + let mechanism = std::env::var("P_KAFKA_SASL_MECHANISM") + .unwrap_or_default() + .to_uppercase(); + + if matches!( + mechanism.as_str(), + "OAUTHBEARER" | "OAUTH-BEARER" | "O-AUTH-BEARER" + ) { + let explicit_provider = std::env::var("P_KAFKA_OAUTH_PROVIDER") + .unwrap_or_default() + .to_ascii_lowercase(); + let has_oidc_endpoint = std::env::var("P_KAFKA_OAUTH_TOKEN_ENDPOINT_URL").is_ok(); + let has_aws_region = ["P_KAFKA_AWS_REGION", "AWS_REGION", "AWS_DEFAULT_REGION"] + .iter() + .any(|name| std::env::var(name).is_ok()); + + let provider = if !explicit_provider.is_empty() { + explicit_provider.as_str() + } else if has_oidc_endpoint { + "oidc" + } else if has_aws_region { + "aws-msk" + } else { + "" + }; + + if provider.is_empty() { + prompts.push(EnvPrompt { + env_var: "P_KAFKA_OAUTH_PROVIDER", + display_name: "Kafka OAuth Provider (aws-msk or oidc)", + required: true, + is_secret: false, + }); + } else if matches!(provider, "aws-msk" | "aws_msk" | "aws") { + if !has_aws_region { + prompts.push(EnvPrompt { + env_var: "P_KAFKA_AWS_REGION", + display_name: "AWS Region for MSK IAM", + required: true, + is_secret: false, + }); + } + } else if matches!(provider, "oidc" | "gcp" | "gcp-managed-kafka") { + prompts.extend([ + EnvPrompt { + env_var: "P_KAFKA_OAUTH_TOKEN_ENDPOINT_URL", + display_name: "OAuth/OIDC Token Endpoint URL", + required: true, + is_secret: false, + }, + EnvPrompt { + env_var: "P_KAFKA_OAUTH_CLIENT_ID", + display_name: "OAuth/OIDC Client ID", + required: true, + is_secret: false, + }, + EnvPrompt { + env_var: "P_KAFKA_OAUTH_CLIENT_SECRET", + display_name: "OAuth/OIDC Client Secret", + required: true, + is_secret: true, + }, + ]); + } + } else if mechanism == "GSSAPI" { + prompts.push(EnvPrompt { + env_var: "P_KAFKA_KERBEROS_SERVICE_NAME", + display_name: "Kafka Kerberos Service Name", required: true, is_secret: false, - }, - EnvPrompt { - env_var: "P_KAFKA_SASL_PASSWORD", - display_name: "Kafka SASL Password", - required: true, - is_secret: true, - }, - ]); + }); + } else if !mechanism.is_empty() { + prompts.extend([ + EnvPrompt { + env_var: "P_KAFKA_SASL_USERNAME", + display_name: "Kafka SASL Username", + required: true, + is_secret: false, + }, + EnvPrompt { + env_var: "P_KAFKA_SASL_PASSWORD", + display_name: "Kafka SASL Password", + required: true, + is_secret: true, + }, + ]); + } } prompts From 91318bcffc94c88d4ef3d0344302eac3e72c212a Mon Sep 17 00:00:00 2001 From: Nikhil Sinha Date: Sun, 19 Jul 2026 19:21:39 +0700 Subject: [PATCH 2/3] fix panic --- src/connectors/kafka/config.rs | 59 ++++++++++++++++++++++++++++++++-- src/connectors/kafka/sink.rs | 50 ++++++++++++++++++++++++++-- 2 files changed, 104 insertions(+), 5 deletions(-) diff --git a/src/connectors/kafka/config.rs b/src/connectors/kafka/config.rs index a5ed0ec63..07f0c79af 100644 --- a/src/connectors/kafka/config.rs +++ b/src/connectors/kafka/config.rs @@ -185,7 +185,7 @@ pub struct ConsumerConfig { required = false, env = "P_KAFKA_CONSUMER_AUTO_OFFSET_RESET", default_value_t = SourceOffset::Earliest, - help = "Auto offset reset behavior" + help = "Where to start consuming when no committed offset exists: earliest, latest, or group. group resumes from committed offsets and falls back to librdkafka's default (latest) for a brand-new consumer group" )] pub auto_offset_reset: SourceOffset, @@ -654,7 +654,11 @@ impl KafkaConfig { } pub fn validate(&self) -> anyhow::Result<()> { - if self.bootstrap_servers.is_none() { + if self + .bootstrap_servers + .as_deref() + .is_none_or(|servers| servers.trim().is_empty()) + { anyhow::bail!("Bootstrap servers must not be empty"); } @@ -711,6 +715,13 @@ impl ConsumerConfig { .set("isolation.level", self.isolation_level.to_string()) .set("group.instance.id", self.group_instance_id.to_string()) .set("statistics.interval.ms", self.stats_interval_ms.to_string()); + + // `auto.offset.reset` only governs where consumption starts when no + // committed offset exists. `Group` means "resume from committed + // offsets", so it keeps librdkafka's default fallback. + if let Some(reset) = self.auto_offset_reset.auto_offset_reset_value() { + config.set("auto.offset.reset", reset); + } } pub fn topics(&self) -> Vec<&str> { @@ -1226,6 +1237,22 @@ impl SourceOffset { SourceOffset::Group => Offset::Stored, } } + + /// Value for librdkafka's `auto.offset.reset`, or `None` to keep the + /// library default. + /// + /// `Group` delegates to committed offsets and deliberately sets nothing: + /// for a brand-new consumer group (no committed offsets anywhere) it is + /// therefore equivalent to `latest`, librdkafka's default fallback. This + /// is documented in the CLI help; use `earliest` when a new group must + /// consume the full topic history. + pub fn auto_offset_reset_value(&self) -> Option<&'static str> { + match self { + SourceOffset::Earliest => Some("earliest"), + SourceOffset::Latest => Some("latest"), + SourceOffset::Group => None, + } + } } #[cfg(test)] @@ -1265,6 +1292,34 @@ mod tests { assert_eq!(cli.security.oauth_provider, Some(OAuthProvider::Oidc)); } + #[test] + fn empty_bootstrap_servers_are_rejected() { + for bootstrap in [None, Some(String::new()), Some(" ".to_string())] { + let config = KafkaConfig { + bootstrap_servers: bootstrap, + ..Default::default() + }; + assert!(config.validate().is_err()); + } + } + + #[test] + fn auto_offset_reset_is_applied_to_consumer_config() { + assert_eq!( + SourceOffset::Earliest.auto_offset_reset_value(), + Some("earliest") + ); + assert_eq!( + SourceOffset::Latest.auto_offset_reset_value(), + Some("latest") + ); + assert_eq!(SourceOffset::Group.auto_offset_reset_value(), None); + + // Default consumer config uses Earliest and must propagate it. + let config = KafkaConfig::default().to_rdkafka_consumer_config(); + assert_eq!(config.get("auto.offset.reset"), Some("earliest")); + } + #[test] fn aws_msk_without_region_defers_to_sdk_chain() { // An explicit aws-msk provider without a region is valid: token diff --git a/src/connectors/kafka/sink.rs b/src/connectors/kafka/sink.rs index 59563d35c..e7b01949c 100644 --- a/src/connectors/kafka/sink.rs +++ b/src/connectors/kafka/sink.rs @@ -23,17 +23,47 @@ use crate::connectors::kafka::processor::StreamWorker; use anyhow::Result; use futures_util::StreamExt; use rdkafka::consumer::Consumer; +use std::future::Future; use std::sync::Arc; use tokio::runtime::Runtime; use tracing::{error, info}; +/// Owns the sink worker runtime and shuts it down without blocking on drop. +/// +/// Dropping a `Runtime` directly inside an async context panics ("Cannot drop +/// a runtime in a context where blocking is not allowed") — which is exactly +/// where the connector future is dropped when the server shuts down or a +/// sibling future in `try_join!` fails. +struct WorkerRuntime(Option); + +impl WorkerRuntime { + fn spawn(&self, future: F) -> tokio::task::JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + self.0 + .as_ref() + .expect("worker runtime is only taken on drop") + .spawn(future) + } +} + +impl Drop for WorkerRuntime { + fn drop(&mut self) { + if let Some(runtime) = self.0.take() { + runtime.shutdown_background(); + } + } +} + pub struct KafkaSinkConnector

where P: Processor, ()>, { streams: KafkaStreams, stream_processor: Arc>, - runtime: Runtime, + runtime: WorkerRuntime, } impl

KafkaSinkConnector

@@ -52,12 +82,11 @@ where "kafka-sink-worker", ) .expect("Failed to build runtime"); - let _ = runtime.enter(); Self { streams: kafka_streams, stream_processor, - runtime, + runtime: WorkerRuntime(Some(runtime)), } } @@ -92,3 +121,18 @@ where Ok(()) } } + +#[cfg(test)] +mod tests { + use super::WorkerRuntime; + use crate::connectors::common::build_runtime; + + #[tokio::test] + async fn worker_runtime_can_be_dropped_inside_async_context() { + // A bare `Runtime` dropped here would panic ("Cannot drop a runtime in + // a context where blocking is not allowed"). This is the failure path + // hit when the connector future is cancelled by `try_join!`. + let runtime = WorkerRuntime(Some(build_runtime(1, "test-worker").unwrap())); + drop(runtime); + } +} From 1353e07432c93a47dc14205a64c50cb1d0641ffc Mon Sep 17 00:00:00 2001 From: Nikhil Sinha Date: Sun, 19 Jul 2026 19:27:53 +0700 Subject: [PATCH 3/3] deepsource fix --- src/connectors/kafka/config.rs | 2 +- src/interactive.rs | 16 ++++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/connectors/kafka/config.rs b/src/connectors/kafka/config.rs index 07f0c79af..98d9b855e 100644 --- a/src/connectors/kafka/config.rs +++ b/src/connectors/kafka/config.rs @@ -1294,7 +1294,7 @@ mod tests { #[test] fn empty_bootstrap_servers_are_rejected() { - for bootstrap in [None, Some(String::new()), Some(" ".to_string())] { + for bootstrap in [None, Some("".to_string()), Some(" ".to_string())] { let config = KafkaConfig { bootstrap_servers: bootstrap, ..Default::default() diff --git a/src/interactive.rs b/src/interactive.rs index 43e45d3eb..bd6adde2b 100644 --- a/src/interactive.rs +++ b/src/interactive.rs @@ -468,15 +468,19 @@ fn get_kafka_prompts() -> Vec { } if needs_sasl { + const KAFKA_SASL_MECHANISM_ENV: &str = "P_KAFKA_SASL_MECHANISM"; + const KAFKA_OAUTH_PROVIDER_ENV: &str = "P_KAFKA_OAUTH_PROVIDER"; + const KAFKA_OAUTH_TOKEN_ENDPOINT_URL_ENV: &str = "P_KAFKA_OAUTH_TOKEN_ENDPOINT_URL"; + prompts.push(EnvPrompt { - env_var: "P_KAFKA_SASL_MECHANISM", + env_var: KAFKA_SASL_MECHANISM_ENV, display_name: "Kafka SASL Mechanism (PLAIN, SCRAM-SHA-256, SCRAM-SHA-512, GSSAPI, OAUTHBEARER)", required: true, is_secret: false, }); - let mechanism = std::env::var("P_KAFKA_SASL_MECHANISM") + let mechanism = std::env::var(KAFKA_SASL_MECHANISM_ENV) .unwrap_or_default() .to_uppercase(); @@ -484,10 +488,10 @@ fn get_kafka_prompts() -> Vec { mechanism.as_str(), "OAUTHBEARER" | "OAUTH-BEARER" | "O-AUTH-BEARER" ) { - let explicit_provider = std::env::var("P_KAFKA_OAUTH_PROVIDER") + let explicit_provider = std::env::var(KAFKA_OAUTH_PROVIDER_ENV) .unwrap_or_default() .to_ascii_lowercase(); - let has_oidc_endpoint = std::env::var("P_KAFKA_OAUTH_TOKEN_ENDPOINT_URL").is_ok(); + let has_oidc_endpoint = std::env::var(KAFKA_OAUTH_TOKEN_ENDPOINT_URL_ENV).is_ok(); let has_aws_region = ["P_KAFKA_AWS_REGION", "AWS_REGION", "AWS_DEFAULT_REGION"] .iter() .any(|name| std::env::var(name).is_ok()); @@ -504,7 +508,7 @@ fn get_kafka_prompts() -> Vec { if provider.is_empty() { prompts.push(EnvPrompt { - env_var: "P_KAFKA_OAUTH_PROVIDER", + env_var: KAFKA_OAUTH_PROVIDER_ENV, display_name: "Kafka OAuth Provider (aws-msk or oidc)", required: true, is_secret: false, @@ -521,7 +525,7 @@ fn get_kafka_prompts() -> Vec { } else if matches!(provider, "oidc" | "gcp" | "gcp-managed-kafka") { prompts.extend([ EnvPrompt { - env_var: "P_KAFKA_OAUTH_TOKEN_ENDPOINT_URL", + env_var: KAFKA_OAUTH_TOKEN_ENDPOINT_URL_ENV, display_name: "OAuth/OIDC Token Endpoint URL", required: true, is_secret: false,