From c19477a9988bec94445904e82c5899ff16461574 Mon Sep 17 00:00:00 2001 From: Alexandre Picard-Lemieux Date: Thu, 16 Jul 2026 21:15:11 -0400 Subject: [PATCH 1/2] feat(kafka): add AWS MSK IAM authentication via SASL/OAUTHBEARER Parseable's Kafka connector enumerated the OAUTHBEARER SASL mechanism but never implemented token generation, so it could not authenticate to AWS MSK clusters that use IAM. The connector was limited to PLAIN, SCRAM, and GSSAPI, and the `generate_oauth_token` callback returned "not implemented". This wires up SASL/OAUTHBEARER for AWS MSK IAM: - Mint IAM auth tokens with the `aws-msk-iam-sasl-signer` crate, driven from the (synchronous) librdkafka OAuth refresh callback via a captured Tokio runtime handle. The callback thread is not a runtime worker, so blocking on the async signer there is safe. - Enable `ClientContext::ENABLE_REFRESH_OAUTH_TOKEN`. librdkafka only invokes the callback for the OAUTHBEARER mechanism, so other mechanisms are unaffected. - Add an `--aws-region` / `P_KAFKA_AWS_REGION` option, falling back to the standard `AWS_REGION` / `AWS_DEFAULT_REGION` environment variables. - Make security validation mechanism-aware: OAUTHBEARER requires only a resolvable region (credentials come from the ambient AWS provider chain), and SASL_SSL no longer requires client certificates, which are a mutual-TLS concern rather than a SASL one. Credentials are resolved from the default AWS provider chain (IRSA, EC2 instance profile, env vars, etc.), so no secrets are added to config. Add unit tests covering the new validation branches and the OAUTHBEARER rdkafka config mapping. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 388 ++++++++++++++++++++++++++++++++- Cargo.toml | 5 + src/connectors/kafka/config.rs | 189 ++++++++++++++-- src/connectors/kafka/mod.rs | 63 +++++- 4 files changed, 607 insertions(+), 38 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8dde4cb56..a8ecd44d0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -775,6 +775,43 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "aws-config" +version = "1.8.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a8fc176d53d6fe85017f230405e3255cedb4a02221cb55ed6d76dccbbb099b2" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-sdk-sts", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand 2.3.0", + "http 1.4.0", + "time", + "tokio", + "tracing", + "url", +] + +[[package]] +name = "aws-credential-types" +version = "1.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d203b0bf2626dcba8665f5cd0871d7c2c0930223d6b6be9097592fea21242d0" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "zeroize", +] + [[package]] name = "aws-lc-rs" version = "1.16.2" @@ -797,6 +834,271 @@ dependencies = [ "fs_extra", ] +[[package]] +name = "aws-msk-iam-sasl-signer" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e62f2c07adf3f0bd1ad40414dd30fcf34c595d5f613e36011d5b62eeff14f490" +dependencies = [ + "aws-config", + "aws-credential-types", + "aws-sdk-sts", + "aws-sigv4", + "aws-types", + "base64", + "chrono", + "futures", + "thiserror 1.0.69", + "url", +] + +[[package]] +name = "aws-runtime" +version = "1.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ede2ddc593e6c8acc6ce3358c28d6677a6dc49b65ba4b37a2befe14a11297e75" +dependencies = [ + "aws-credential-types", + "aws-sigv4", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "bytes-utils", + "fastrand 2.3.0", + "http 1.4.0", + "http-body 1.0.1", + "percent-encoding", + "pin-project-lite", + "tracing", + "uuid", +] + +[[package]] +name = "aws-sdk-sts" +version = "1.99.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9acba7c62f3d4e2408fa998a3a8caacd8b9a5b5549cf36e2372fbdae329d5449" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-query", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-smithy-xml", + "aws-types", + "fastrand 2.3.0", + "http 0.2.12", + "http 1.4.0", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sigv4" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37411f8e0f4bea0c3ca0958ce7f18f6439db24d555dbd809787262cd00926aa9" +dependencies = [ + "aws-credential-types", + "aws-smithy-http", + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "form_urlencoded", + "hex", + "hmac", + "http 0.2.12", + "http 1.4.0", + "percent-encoding", + "sha2", + "time", + "tracing", +] + +[[package]] +name = "aws-smithy-async" +version = "1.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc50d0f63e714784b84223abd7abbc8577de8c35d699e0edd19f0a88a08ae13" +dependencies = [ + "futures-util", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "aws-smithy-http" +version = "0.63.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d619373d490ad70966994801bc126846afaa0d1ee920697a031f0cf63f2568e7" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "bytes-utils", + "futures-core", + "futures-util", + "http 1.4.0", + "http-body 1.0.1", + "http-body-util", + "percent-encoding", + "pin-project-lite", + "pin-utils", + "tracing", +] + +[[package]] +name = "aws-smithy-http-client" +version = "1.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00ccbb08c10f6bcf912f398188e42ee2eab5f1767ce215a02a73bc5df1bbdd95" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "h2 0.4.12", + "http 1.4.0", + "hyper", + "hyper-rustls", + "hyper-util", + "pin-project-lite", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tower", + "tracing", +] + +[[package]] +name = "aws-smithy-json" +version = "0.62.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27b3a779093e18cad88bbae08dc4261e1d95018c4c5b9356a52bcae7c0b6e9bb" +dependencies = [ + "aws-smithy-types", +] + +[[package]] +name = "aws-smithy-observability" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d3f39d5bb871aaf461d59144557f16d5927a5248a983a40654d9cf3b9ba183b" +dependencies = [ + "aws-smithy-runtime-api", +] + +[[package]] +name = "aws-smithy-query" +version = "0.60.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f76a580e3d8f8961e5d48763214025a2af65c2fa4cd1fb7f270a0e107a71b0" +dependencies = [ + "aws-smithy-types", + "urlencoding", +] + +[[package]] +name = "aws-smithy-runtime" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ccf7f6eba8b2dcf8ce9b74806c6c185659c311665c4bf8d6e71ebd454db6bf" +dependencies = [ + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-http-client", + "aws-smithy-observability", + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "fastrand 2.3.0", + "http 0.2.12", + "http 1.4.0", + "http-body 0.4.6", + "http-body 1.0.1", + "http-body-util", + "pin-project-lite", + "pin-utils", + "tokio", + "tracing", +] + +[[package]] +name = "aws-smithy-runtime-api" +version = "1.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4af6e5def28be846479bbeac55aa4603d6f7986fc5da4601ba324dd5d377516" +dependencies = [ + "aws-smithy-async", + "aws-smithy-types", + "bytes", + "http 0.2.12", + "http 1.4.0", + "pin-project-lite", + "tokio", + "tracing", + "zeroize", +] + +[[package]] +name = "aws-smithy-types" +version = "1.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ca2734c16913a45343b37313605d84e7d8b34a4611598ce1d25b35860a2bed3" +dependencies = [ + "base64-simd", + "bytes", + "bytes-utils", + "futures-core", + "http 0.2.12", + "http 1.4.0", + "http-body 0.4.6", + "http-body 1.0.1", + "http-body-util", + "itoa", + "num-integer", + "pin-project-lite", + "pin-utils", + "ryu", + "serde", + "time", + "tokio", + "tokio-util", +] + +[[package]] +name = "aws-smithy-xml" +version = "0.60.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b53543b4b86ed43f051644f704a98c7291b3618b67adf057ee77a366fa52fcaa" +dependencies = [ + "xmlparser", +] + +[[package]] +name = "aws-types" +version = "1.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0470cc047657c6e286346bdf10a8719d26efd6a91626992e0e64481e44323e96" +dependencies = [ + "aws-credential-types", + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "rustc_version", + "tracing", +] + [[package]] name = "axum" version = "0.8.8" @@ -807,7 +1109,7 @@ dependencies = [ "bytes", "futures-util", "http 1.4.0", - "http-body", + "http-body 1.0.1", "http-body-util", "itoa", "matchit", @@ -831,7 +1133,7 @@ dependencies = [ "bytes", "futures-core", "http 1.4.0", - "http-body", + "http-body 1.0.1", "http-body-util", "mime", "pin-project-lite", @@ -861,6 +1163,16 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64-simd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "339abbe78e73178762e23bea9dfd08e697eb3f3301cd4be981c0f78ba5859195" +dependencies = [ + "outref", + "vsimd", +] + [[package]] name = "base64ct" version = "1.8.1" @@ -987,6 +1299,16 @@ version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +[[package]] +name = "bytes-utils" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dafe3a8757b027e2be6e4e5601ed563c55989fcf1546e933c66c8eb3a058d35" +dependencies = [ + "bytes", + "either", +] + [[package]] name = "bytestring" version = "1.5.0" @@ -2773,6 +3095,15 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest 0.10.7", +] + [[package]] name = "hmac-sha256" version = "1.1.12" @@ -2873,6 +3204,17 @@ dependencies = [ "base64", ] +[[package]] +name = "http-body" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" +dependencies = [ + "bytes", + "http 0.2.12", + "pin-project-lite", +] + [[package]] name = "http-body" version = "1.0.1" @@ -2892,7 +3234,7 @@ dependencies = [ "bytes", "futures-core", "http 1.4.0", - "http-body", + "http-body 1.0.1", "pin-project-lite", ] @@ -2951,7 +3293,7 @@ dependencies = [ "futures-core", "h2 0.4.12", "http 1.4.0", - "http-body", + "http-body 1.0.1", "httparse", "httpdate", "itoa", @@ -3005,7 +3347,7 @@ dependencies = [ "futures-core", "futures-util", "http 1.4.0", - "http-body", + "http-body 1.0.1", "hyper", "ipnet", "libc", @@ -3929,6 +4271,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "outref" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" + [[package]] name = "parking" version = "2.2.1" @@ -4014,6 +4362,8 @@ dependencies = [ "arrow-schema", "arrow-select", "async-trait", + "aws-msk-iam-sasl-signer", + "aws-types", "base64", "byteorder", "bytes", @@ -4821,7 +5171,7 @@ dependencies = [ "futures-util", "h2 0.4.12", "http 1.4.0", - "http-body", + "http-body 1.0.1", "http-body-util", "hyper", "hyper-rustls", @@ -4864,7 +5214,7 @@ dependencies = [ "futures-core", "futures-util", "http 1.4.0", - "http-body", + "http-body 1.0.1", "http-body-util", "hyper", "hyper-util", @@ -5832,7 +6182,7 @@ dependencies = [ "flate2", "h2 0.4.12", "http 1.4.0", - "http-body", + "http-body 1.0.1", "http-body-util", "hyper", "hyper-timeout", @@ -5883,7 +6233,7 @@ dependencies = [ "base64", "bytes", "http 1.4.0", - "http-body", + "http-body 1.0.1", "pin-project", "tokio-stream", "tonic", @@ -5923,7 +6273,7 @@ dependencies = [ "futures-core", "futures-util", "http 1.4.0", - "http-body", + "http-body 1.0.1", "http-body-util", "iri-string", "pin-project-lite", @@ -6160,6 +6510,12 @@ dependencies = [ "serde", ] +[[package]] +name = "urlencoding" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" + [[package]] name = "utf8_iter" version = "1.0.4" @@ -6273,6 +6629,12 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "vsimd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" + [[package]] name = "waker-fn" version = "1.2.0" @@ -6977,6 +7339,12 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" +[[package]] +name = "xmlparser" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" + [[package]] name = "xxhash-rust" version = "0.8.15" diff --git a/Cargo.toml b/Cargo.toml index 61b162954..538b83e8e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -56,6 +56,9 @@ rdkafka = { version = "0.37", optional = true, features = [ "libz-static", ] } 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 } # Authentication and Security argon2 = "0.5.0" @@ -239,6 +242,8 @@ kafka = [ "rdkafka/sasl", "sasl2-sys", "sasl2-sys/vendored", + "aws-msk-iam-sasl-signer", + "aws-types", ] [profile.release-lto] diff --git a/src/connectors/kafka/config.rs b/src/connectors/kafka/config.rs index ddb7f6fee..4ec95c6cc 100644 --- a/src/connectors/kafka/config.rs +++ b/src/connectors/kafka/config.rs @@ -504,6 +504,15 @@ pub struct SecurityConfig { )] pub sasl_password: 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." + )] + pub aws_region: Option, + #[arg( long = "ssl-key-password", env = "P_KAFKA_SSL_KEY_PASSWORD", @@ -791,39 +800,74 @@ impl SecurityConfig { } } - // TODO: Implement OAuthBearer mechanism for SASL when needed. This depends on the vendor (on-prem, Confluent Kafka, AWS MSK, etc.). + // 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. } } + /// Resolves the AWS region to use for MSK IAM authentication. + /// + /// Precedence: the explicit `--aws-region` flag, then the standard + /// `AWS_REGION` / `AWS_DEFAULT_REGION` environment variables. + pub fn resolved_aws_region(&self) -> Option { + self.aws_region + .clone() + .or_else(|| std::env::var("AWS_REGION").ok()) + .or_else(|| std::env::var("AWS_DEFAULT_REGION").ok()) + .filter(|region| !region.is_empty()) + } + fn validate(&self) -> anyhow::Result<()> { - match self.protocol { - SecurityProtocol::Ssl | SecurityProtocol::SaslSsl => { - if self.ssl_ca_location.is_none() { - anyhow::bail!("CA certificate location is required for SSL"); - } - if self.ssl_certificate_location.is_none() { - anyhow::bail!("Client certificate location is required for SSL"); - } - if self.ssl_key_location.is_none() { - anyhow::bail!("Client key location is required for SSL"); - } + // Pure TLS (mutual auth) requires the full client certificate material. + // For SASL_SSL, TLS is only used for server authentication and channel + // encryption, so client certificates are not required — the caller + // authenticates through the SASL mechanism instead. + if matches!(self.protocol, SecurityProtocol::Ssl) { + if self.ssl_ca_location.is_none() { + anyhow::bail!("CA certificate location is required for SSL"); } - SecurityProtocol::SaslPlaintext => { - if self.sasl_mechanism.is_none() { - anyhow::bail!("SASL mechanism is required when SASL is enabled"); + if self.ssl_certificate_location.is_none() { + anyhow::bail!("Client certificate location is required for SSL"); + } + if self.ssl_key_location.is_none() { + anyhow::bail!("Client key location is required for SSL"); + } + } + + // SASL-bearing protocols must specify a mechanism and satisfy the + // credential requirements specific to that mechanism. + if matches!( + self.protocol, + SecurityProtocol::SaslSsl | SecurityProtocol::SaslPlaintext + ) { + match self.sasl_mechanism { + None => anyhow::bail!("SASL mechanism is required when SASL is enabled"), + Some(SaslMechanism::OAuthBearer) => { + // AWS MSK IAM mints tokens from the ambient AWS credential + // chain, so only a resolvable region is required here. + if self.resolved_aws_region().is_none() { + anyhow::bail!( + "AWS region is required for SASL/OAUTHBEARER (MSK IAM); set --aws-region or the AWS_REGION environment variable" + ); + } } - if self.sasl_username.is_none() || self.sasl_password.is_none() { - anyhow::bail!("SASL username and password are required"); + Some(SaslMechanism::Gssapi) => { + if self.kerberos_service_name.is_none() { + anyhow::bail!("Kerberos service name is required for GSSAPI"); + } } - - if matches!(self.sasl_mechanism, Some(SaslMechanism::Gssapi)) - && self.kerberos_service_name.is_none() - { - anyhow::bail!("Kerberos service name is required for GSSAPI"); + Some( + SaslMechanism::Plain | SaslMechanism::ScramSha256 | SaslMechanism::ScramSha512, + ) => { + if self.sasl_username.is_none() || self.sasl_password.is_none() { + anyhow::bail!("SASL username and password are required"); + } } } - SecurityProtocol::Plaintext => {} // No validation needed for PLAINTEXT } + Ok(()) } } @@ -965,6 +1009,7 @@ impl Default for SecurityConfig { sasl_mechanism: None, sasl_username: None, sasl_password: None, + aws_region: None, ssl_key_password: None, kerberos_service_name: None, kerberos_principal: None, @@ -1024,3 +1069,101 @@ impl SourceOffset { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn oauthbearer_requires_a_region() { + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + aws_region: None, + ..Default::default() + }; + + // Region resolution also consults AWS_REGION / AWS_DEFAULT_REGION, so + // only assert the failure when neither is present in the environment. + if security.resolved_aws_region().is_none() { + assert!(security.validate().is_err()); + } + } + + #[test] + fn oauthbearer_with_explicit_region_is_valid() { + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + aws_region: Some("us-east-1".to_string()), + ..Default::default() + }; + + assert!(security.validate().is_ok()); + assert_eq!(security.resolved_aws_region().as_deref(), Some("us-east-1")); + } + + #[test] + fn oauthbearer_does_not_require_username_or_password() { + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + aws_region: Some("eu-west-1".to_string()), + sasl_username: None, + sasl_password: None, + ..Default::default() + }; + + assert!(security.validate().is_ok()); + } + + #[test] + fn sasl_ssl_does_not_require_client_certificates() { + // SASL_SSL uses TLS only for the server side; client certs are a + // mutual-TLS concern and must not be required here. + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::ScramSha512), + sasl_username: Some("user".to_string()), + sasl_password: Some("pass".to_string()), + ssl_ca_location: None, + ssl_certificate_location: None, + ssl_key_location: None, + ..Default::default() + }; + + assert!(security.validate().is_ok()); + } + + #[test] + fn scram_still_requires_username_and_password() { + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::ScramSha512), + sasl_username: None, + sasl_password: None, + ..Default::default() + }; + + assert!(security.validate().is_err()); + } + + #[test] + fn oauthbearer_applies_mechanism_without_credentials() { + let security = SecurityConfig { + protocol: SecurityProtocol::SaslSsl, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + aws_region: Some("us-east-2".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")); + // No static credentials should be set for OAUTHBEARER. + assert_eq!(config.get("sasl.username"), None); + assert_eq!(config.get("sasl.password"), None); + } +} diff --git a/src/connectors/kafka/mod.rs b/src/connectors/kafka/mod.rs index 43a7796b3..5960a96a3 100644 --- a/src/connectors/kafka/mod.rs +++ b/src/connectors/kafka/mod.rs @@ -16,8 +16,10 @@ * */ -use crate::connectors::kafka::config::KafkaConfig; +use crate::connectors::kafka::config::{KafkaConfig, SaslMechanism}; use crate::handlers::TENANT_ID; +use aws_msk_iam_sasl_signer::generate_auth_token; +use aws_types::region::Region; use derive_more::Constructor; use rdkafka::client::OAuthToken; use rdkafka::consumer::{ConsumerContext, Rebalance}; @@ -30,6 +32,7 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::error::Error; use std::sync::{Arc, RwLock}; +use tokio::runtime::Handle; use tokio::sync::mpsc; use tokio::sync::mpsc::Receiver; use tracing::{error, info, warn}; @@ -53,6 +56,13 @@ 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, } impl KafkaContext { @@ -64,6 +74,7 @@ impl KafkaContext { config, statistics, rebalance_tx, + runtime_handle: Handle::current(), }, rebalance_rx, ) @@ -251,8 +262,10 @@ impl ProducerContext for KafkaContext { } impl ClientContext for KafkaContext { - // TODO: when implementing OAuth, set this to true - const ENABLE_REFRESH_OAUTH_TOKEN: bool = false; + // 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. + const ENABLE_REFRESH_OAUTH_TOKEN: bool = true; fn stats(&self, new_stats: Statistics) { match self.statistics.write() { @@ -265,11 +278,51 @@ 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. fn generate_oauth_token( &self, _oauthbearer_config: Option<&str>, ) -> Result> { - // TODO Implement OAuth token generation when needed - Err("OAuth token generation is not implemented".into()) + 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) + ) { + return Err( + "OAuth token generation is only supported for the OAUTHBEARER (AWS MSK IAM) SASL mechanism" + .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}"); + + // 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() + })?; + + Ok(OAuthToken { + // MSK does not consume the principal name, but librdkafka requires a + // non-empty value. + principal_name: "parseable".to_string(), + token, + lifetime_ms: expiration_time_ms, + }) } } From fd22621a08f6ab6d36b78bfb094076ceaffb95b4 Mon Sep 17 00:00:00 2001 From: Alexandre Picard-Lemieux Date: Fri, 17 Jul 2026 10:36:10 -0400 Subject: [PATCH 2/2] fix(kafka): address MSK IAM review feedback - Reject SASL/OAUTHBEARER over SASL_PLAINTEXT. MSK IAM tokens are signed bearer credentials and must only travel over the encrypted SASL_SSL protocol, per AWS guidance for non-Java clients. - Normalize each AWS region source (trim + reject empty) before applying precedence, so an explicitly-empty --aws-region no longer shadows a valid AWS_REGION / AWS_DEFAULT_REGION fallback. - Replace inline env-var string literals with named constants (AWS_REGION_ENV, AWS_DEFAULT_REGION_ENV) to guard against typos (DeepSource RS-W1015). - Split validation into a pure `validate_with_region` helper so the missing-region failure path is tested deterministically instead of depending on the ambient environment. Add tests for the SASL_PLAINTEXT rejection and region normalization. Co-Authored-By: Claude Opus 4.8 --- src/connectors/kafka/config.rs | 78 +++++++++++++++++++++++++++++----- 1 file changed, 67 insertions(+), 11 deletions(-) diff --git a/src/connectors/kafka/config.rs b/src/connectors/kafka/config.rs index 4ec95c6cc..738819b74 100644 --- a/src/connectors/kafka/config.rs +++ b/src/connectors/kafka/config.rs @@ -23,6 +23,12 @@ use serde::{Deserialize, Serialize}; use std::path::PathBuf; use std::time::Duration; +/// Standard AWS environment variables consulted when resolving the region for +/// MSK IAM authentication. Kept as named constants to avoid typo-prone string +/// literals in `std::env` calls. +const AWS_REGION_ENV: &str = "AWS_REGION"; +const AWS_DEFAULT_REGION_ENV: &str = "AWS_DEFAULT_REGION"; + #[derive(Debug, Clone, Parser)] pub struct KafkaConfig { #[arg( @@ -810,16 +816,33 @@ impl SecurityConfig { /// Resolves the AWS region to use for MSK IAM authentication. /// /// Precedence: the explicit `--aws-region` flag, then the standard - /// `AWS_REGION` / `AWS_DEFAULT_REGION` environment variables. + /// `AWS_REGION` / `AWS_DEFAULT_REGION` environment variables. Each source is + /// 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.aws_region - .clone() - .or_else(|| std::env::var("AWS_REGION").ok()) - .or_else(|| std::env::var("AWS_DEFAULT_REGION").ok()) + 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())) + } + + /// Trims a candidate region value and discards it if it is empty. + fn normalize_region(region: Option) -> Option { + region + .map(|region| region.trim().to_owned()) .filter(|region| !region.is_empty()) } fn validate(&self) -> anyhow::Result<()> { + // Resolve the region (including environment fallbacks) once, then defer + // to the pure validation logic so it can be tested deterministically. + self.validate_with_region(self.resolved_aws_region().as_deref()) + } + + /// Validates the security configuration against an already-resolved AWS + /// region. Kept separate from region resolution (which reads process + /// environment variables) so the validation rules can be exercised without + /// depending on the ambient environment. + fn validate_with_region(&self, resolved_region: Option<&str>) -> anyhow::Result<()> { // Pure TLS (mutual auth) requires the full client certificate material. // For SASL_SSL, TLS is only used for server authentication and channel // encryption, so client certificates are not required — the caller @@ -845,9 +868,17 @@ 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" + ); + } // AWS MSK IAM mints tokens from the ambient AWS credential // chain, so only a resolvable region is required here. - if self.resolved_aws_region().is_none() { + 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" ); @@ -1083,11 +1114,36 @@ mod tests { ..Default::default() }; - // Region resolution also consults AWS_REGION / AWS_DEFAULT_REGION, so - // only assert the failure when neither is present in the environment. - if security.resolved_aws_region().is_none() { - assert!(security.validate().is_err()); - } + // 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_err()); + } + + #[test] + fn oauthbearer_rejects_sasl_plaintext() { + // The bearer token must never traverse an unencrypted connection. + let security = SecurityConfig { + protocol: SecurityProtocol::SaslPlaintext, + sasl_mechanism: Some(SaslMechanism::OAuthBearer), + aws_region: Some("us-east-1".to_string()), + ..Default::default() + }; + + assert!(security.validate_with_region(Some("us-east-1")).is_err()); + } + + #[test] + fn normalize_region_trims_and_rejects_empty() { + assert_eq!( + SecurityConfig::normalize_region(Some(" us-east-1 ".to_string())).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(None), None); } #[test]