diff --git a/Cargo.lock b/Cargo.lock index bbeb8461..be730927 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -486,6 +486,30 @@ dependencies = [ "libc", ] +[[package]] +name = "crc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" + +[[package]] +name = "crc32c" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a47af21622d091a8f0fb295b88bc886ac74efcc613efc19f5d0b21de5c89e47" +dependencies = [ + "rustc_version", +] + [[package]] name = "crc32fast" version = "1.5.2" @@ -1719,6 +1743,20 @@ dependencies = [ "snafu 0.9.2", ] +[[package]] +name = "kafka-protocol" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "099d5c2f1b40cd830cbf18ca4d2a0805f2875b811ca18f0372b2433e68fe2dda" +dependencies = [ + "anyhow", + "bytes", + "crc", + "crc32c", + "indexmap", + "uuid", +] + [[package]] name = "konst" version = "0.2.20" @@ -3196,6 +3234,23 @@ dependencies = [ "zeroize", ] +[[package]] +name = "stackable-kafka-admin" +version = "0.0.0-dev" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bytes", + "futures", + "kafka-protocol", + "rstest", + "rustls", + "snafu 0.9.2", + "tokio", + "tokio-rustls", + "tokio-util", +] + [[package]] name = "stackable-kafka-agent" version = "0.0.0-dev" diff --git a/Cargo.nix b/Cargo.nix index 116feeff..7003815a 100644 --- a/Cargo.nix +++ b/Cargo.nix @@ -43,6 +43,16 @@ rec { # You can override the features with # workspaceMembers."${crateName}".build.override { features = [ "default" "feature1" ... ]; }. workspaceMembers = { + "stackable-kafka-admin" = rec { + packageId = "stackable-kafka-admin"; + build = internal.buildRustCrateWithFeatures { + packageId = "stackable-kafka-admin"; + }; + + # Debug support which might change between releases. + # File a bug if you depend on any for non-debug work! + debug = internal.debugCrate { inherit packageId; }; + }; "stackable-kafka-agent" = rec { packageId = "stackable-kafka-agent"; build = internal.buildRustCrateWithFeatures { @@ -1516,6 +1526,50 @@ rec { } ]; + }; + "crc" = rec { + crateName = "crc"; + version = "3.4.0"; + edition = "2021"; + sha256 = "03dsq5qsv86m35ikg84l80d00wnkjm8q4pjxgac0vaqjrnhs5f2y"; + authors = [ + "Rui Hu " + "Akhil Velagapudi <4@4khil.com>" + ]; + dependencies = [ + { + name = "crc-catalog"; + packageId = "crc-catalog"; + } + ]; + + }; + "crc-catalog" = rec { + crateName = "crc-catalog"; + version = "2.5.0"; + edition = "2018"; + sha256 = "0lq8dl60g457za8la86pak6i7ydxanm2lrpkqh5kyjkbz7m9hxi1"; + libName = "crc_catalog"; + authors = [ + "Akhil Velagapudi " + ]; + + }; + "crc32c" = rec { + crateName = "crc32c"; + version = "0.6.8"; + edition = "2018"; + sha256 = "0iwyr3jivcnhylczqgk1rkpp9b46r25vi5dj1y7il29dc8hsyirs"; + authors = [ + "Zack Owens" + ]; + buildDependencies = [ + { + name = "rustc_version"; + packageId = "rustc_version"; + } + ]; + }; "crc32fast" = rec { crateName = "crc32fast"; @@ -5472,6 +5526,52 @@ rec { }; resolvedDefaultFeatures = [ "darling" ]; }; + "kafka-protocol" = rec { + crateName = "kafka-protocol"; + version = "0.18.0"; + edition = "2021"; + sha256 = "1nidzrl3whxjf81qz88wh5dqgwh510m4vjhqpw687ka03cpmr789"; + libName = "kafka_protocol"; + authors = [ + "Diggory Blake " + "Charlotte McElwain " + "belltoy " + ]; + dependencies = [ + { + name = "anyhow"; + packageId = "anyhow"; + } + { + name = "bytes"; + packageId = "bytes"; + } + { + name = "crc"; + packageId = "crc"; + } + { + name = "crc32c"; + packageId = "crc32c"; + } + { + name = "indexmap"; + packageId = "indexmap"; + } + { + name = "uuid"; + packageId = "uuid"; + } + ]; + features = { + "default" = [ "client" "broker" "gzip" "zstd" "snappy" "lz4" ]; + "gzip" = [ "dep:flate2" ]; + "lz4" = [ "dep:lz4" ]; + "snappy" = [ "dep:snap" ]; + "zstd" = [ "dep:zstd" ]; + }; + resolvedDefaultFeatures = [ "broker" "client" ]; + }; "konst" = rec { crateName = "konst"; version = "0.2.20"; @@ -10492,6 +10592,78 @@ rec { }; resolvedDefaultFeatures = [ "default" "rustls" ]; }; + "stackable-kafka-admin" = rec { + crateName = "stackable-kafka-admin"; + version = "0.0.0-dev"; + edition = "2024"; + src = lib.cleanSourceWith { filter = sourceFilter; src = ./rust/kafka-admin; }; + libName = "stackable_kafka_admin"; + authors = [ + "Stackable GmbH " + ]; + dependencies = [ + { + name = "anyhow"; + packageId = "anyhow"; + } + { + name = "base64"; + packageId = "base64 0.22.1"; + } + { + name = "bytes"; + packageId = "bytes"; + } + { + name = "futures"; + packageId = "futures"; + } + { + name = "kafka-protocol"; + packageId = "kafka-protocol"; + usesDefaultFeatures = false; + features = [ "broker" "client" ]; + } + { + name = "rustls"; + packageId = "rustls"; + usesDefaultFeatures = false; + features = [ "logging" "ring" "std" "tls12" ]; + } + { + name = "snafu"; + packageId = "snafu 0.9.2"; + } + { + name = "tokio"; + packageId = "tokio"; + features = [ "full" ]; + } + { + name = "tokio-rustls"; + packageId = "tokio-rustls"; + usesDefaultFeatures = false; + features = [ "logging" "ring" "tls12" ]; + } + { + name = "tokio-util"; + packageId = "tokio-util"; + features = [ "codec" ]; + } + ]; + devDependencies = [ + { + name = "rstest"; + packageId = "rstest"; + } + { + name = "tokio"; + packageId = "tokio"; + features = [ "full" "test-util" ]; + } + ]; + + }; "stackable-kafka-agent" = rec { crateName = "stackable-kafka-agent"; version = "0.0.0-dev"; @@ -12011,7 +12183,7 @@ rec { "tracing" = [ "dep:tracing" ]; "windows-sys" = [ "dep:windows-sys" ]; }; - resolvedDefaultFeatures = [ "bytes" "default" "fs" "full" "io-std" "io-util" "libc" "macros" "mio" "net" "parking_lot" "process" "rt" "rt-multi-thread" "signal" "signal-hook-registry" "socket2" "sync" "time" "tokio-macros" "windows-sys" ]; + resolvedDefaultFeatures = [ "bytes" "default" "fs" "full" "io-std" "io-util" "libc" "macros" "mio" "net" "parking_lot" "process" "rt" "rt-multi-thread" "signal" "signal-hook-registry" "socket2" "sync" "test-util" "time" "tokio-macros" "windows-sys" ]; }; "tokio-macros" = rec { crateName = "tokio-macros"; diff --git a/Cargo.toml b/Cargo.toml index 8f9321a8..540a2aeb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["rust/operator-binary", "rust/agent-binary"] +members = ["rust/operator-binary", "rust/agent-binary", "rust/kafka-admin"] resolver = "2" [workspace.package] @@ -13,12 +13,16 @@ repository = "https://github.com/stackabletech/kafka-operator" stackable-operator = { git = "https://github.com/stackabletech/operator-rs.git", branch = "feat/platform-access", features = ["crds", "webhook"] } anyhow = "1.0" +base64 = "0.22" built = { version = "0.8", features = ["chrono", "git2"] } +bytes = "1.12" clap = "4.6" const_format = "0.2" futures = "0.3" indoc = "2.0" +kafka-protocol = { version = "0.18", default-features = false, features = ["broker", "client"] } rstest = "0.26" +rustls = { version = "0.23", default-features = false, features = ["logging", "ring", "std", "tls12"] } semver = "1.0" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" @@ -26,6 +30,8 @@ serde_yaml = "0.9" snafu = "0.9" strum = { version = "0.28", features = ["derive"] } tokio = { version = "1.53", features = ["full"] } +tokio-rustls = { version = "0.26", default-features = false, features = ["logging", "ring", "tls12"] } +tokio-util = { version = "0.7", features = ["codec"] } tracing = "0.1" [patch."https://github.com/stackabletech/operator-rs.git"] diff --git a/rust/kafka-admin/Cargo.toml b/rust/kafka-admin/Cargo.toml new file mode 100644 index 00000000..2540fda7 --- /dev/null +++ b/rust/kafka-admin/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "stackable-kafka-admin" +description = "Kafka admin client of the Stackable agent for Apache Kafka" +version.workspace = true +authors.workspace = true +license.workspace = true +edition.workspace = true +repository.workspace = true +publish = false + +[dependencies] +anyhow.workspace = true +base64.workspace = true +bytes.workspace = true +futures.workspace = true +kafka-protocol.workspace = true +rustls.workspace = true +snafu.workspace = true +tokio.workspace = true +tokio-rustls.workspace = true +tokio-util.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/rust/kafka-admin/src/bootstrap.rs b/rust/kafka-admin/src/bootstrap.rs new file mode 100644 index 00000000..1ab97bb9 --- /dev/null +++ b/rust/kafka-admin/src/bootstrap.rs @@ -0,0 +1,106 @@ +//! The bootstrap servers of a Kafka cluster. + +use std::fmt; + +use snafu::{OptionExt, Snafu, ensure}; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("no bootstrap servers given"))] + NoBootstrapServers, + + #[snafu(display("{server:?} is not a valid bootstrap server, expected host:port"))] + InvalidBootstrapServer { server: String }, +} + +type Result = std::result::Result; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct BrokerAddress { + pub host: String, + pub port: u16, +} + +impl fmt::Display for BrokerAddress { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.host.contains(':') { + write!(f, "[{}]:{}", self.host, self.port) + } else { + write!(f, "{}:{}", self.host, self.port) + } + } +} + +/// Parses a comma-separated list of `host:port` entries. IPv6 hosts may be enclosed in brackets. +pub fn parse_bootstrap_servers(value: &str) -> Result> { + let servers = value + .split(',') + .map(str::trim) + .filter(|server| !server.is_empty()) + .map(|server| { + // Splits only at the ':' separating host and port, not inside IPv6 hosts + let (host, port) = server + .rsplit_once(':') + .context(InvalidBootstrapServerSnafu { server })?; + let host = host + .strip_prefix('[') + .and_then(|host| host.strip_suffix(']')) + .unwrap_or(host); + let port = port + .parse() + .ok() + .context(InvalidBootstrapServerSnafu { server })?; + ensure!( + !host.is_empty() && !host.contains(['[', ']']), + InvalidBootstrapServerSnafu { server } + ); + Ok(BrokerAddress { + host: host.to_owned(), + port, + }) + }) + .collect::>>()?; + ensure!(!servers.is_empty(), NoBootstrapServersSnafu); + Ok(servers) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case("kafka:9092", &[("kafka", 9092)])] + #[case("kafka-0:9092, kafka-1:9093", &[("kafka-0", 9092), ("kafka-1", 9093)])] + #[case("kafka:9092,", &[("kafka", 9092)])] + #[case("[fd00::1]:9092", &[("fd00::1", 9092)])] + #[case("fd00::1:9092", &[("fd00::1", 9092)])] + fn parse_valid_bootstrap_servers(#[case] value: &str, #[case] expected: &[(&str, u16)]) { + let expected = expected + .iter() + .map(|(host, port)| BrokerAddress { + host: host.to_string(), + port: *port, + }) + .collect::>(); + + assert_eq!( + expected, + parse_bootstrap_servers(value).expect("the bootstrap servers should be valid") + ); + } + + #[rstest] + #[case("")] + #[case(" , ")] + #[case("kafka")] + #[case("kafka:port")] + #[case("kafka:99999")] + #[case(":9092")] + #[case("[fd00::1:9092")] + #[case("fd00::1]:9092")] + fn reject_invalid_bootstrap_servers(#[case] value: &str) { + parse_bootstrap_servers(value).expect_err("the bootstrap servers should be invalid"); + } +} diff --git a/rust/kafka-admin/src/client.rs b/rust/kafka-admin/src/client.rs new file mode 100644 index 00000000..ce041188 --- /dev/null +++ b/rust/kafka-admin/src/client.rs @@ -0,0 +1,860 @@ +//! An admin client for the topics of a Kafka cluster in KRaft mode. + +use std::{ + collections::{BTreeMap, BTreeSet}, + time::Duration, +}; + +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use kafka_protocol::{ + ResponseError, + error::ParseResponseErrorCode, + messages::{ + BrokerId, CreatePartitionsRequest, CreateTopicsRequest, DeleteTopicsRequest, + DeleteTopicsResponse, DescribeConfigsRequest, DescribeConfigsResponse, + IncrementalAlterConfigsRequest, ListPartitionReassignmentsRequest, + ListPartitionReassignmentsResponse, MetadataRequest, MetadataResponse, TopicName, + create_partitions_request::CreatePartitionsTopic, + create_topics_request::{CreatableTopic, CreatableTopicConfig}, + delete_topics_request::DeleteTopicState, + describe_configs_request::DescribeConfigsResource, + incremental_alter_configs_request::{AlterConfigsResource, AlterableConfig}, + list_partition_reassignments_request::ListPartitionReassignmentsTopics, + metadata_request::MetadataRequestTopic, + }, + protocol::{Request, StrBytes}, +}; +use snafu::{OptionExt, ResultExt, Snafu}; +use tokio::time::Instant; + +use crate::{ + bootstrap::BrokerAddress, + connection::{self, Connection, Security}, +}; + +const METADATA_VERSION: i16 = 12; +const CREATE_TOPICS_VERSION: i16 = 7; +const DELETE_TOPICS_VERSION: i16 = 6; +const CREATE_PARTITIONS_VERSION: i16 = 3; +const DESCRIBE_CONFIGS_VERSION: i16 = 4; +const INCREMENTAL_ALTER_CONFIGS_VERSION: i16 = 1; +const LIST_PARTITION_REASSIGNMENTS_VERSION: i16 = 0; + +/// Must stay below the connection's request timeout. +const BROKER_TIMEOUT_MS: i32 = 25_000; + +/// Must stay below the broker's `connections.max.idle.ms` (10 minutes by default). +const CONNECTION_MAX_IDLE: Duration = Duration::from_secs(5 * 60); + +const RESOURCE_TYPE_TOPIC: i8 = 2; +const CONFIG_SOURCE_DYNAMIC_TOPIC: i8 = 1; +const CONFIG_OPERATION_SET: i8 = 0; +const CONFIG_OPERATION_DELETE: i8 = 1; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("no bootstrap servers given"))] + NoBootstrapServers, + + #[snafu(display("failed to connect to the bootstrap server {server}"))] + Connect { + source: connection::Error, + server: BrokerAddress, + }, + + #[snafu(display("failed to send the request to the broker"))] + Send { source: connection::Error }, + + #[snafu(display( + "the response to the request with API key {api_key} contains no result for the topic {topic:?}" + ))] + MissingResult { api_key: i16, topic: String }, + + #[snafu(display( + "the broker rejected the request with API key {api_key} for the topic {topic:?} with {error}{}", + message.as_deref().map(|message| format!(": {message}")).unwrap_or_default() + ))] + Kafka { + api_key: i16, + topic: String, + error: ResponseError, + message: Option, + }, +} + +type Result = std::result::Result; + +impl Error { + /// Whether the broker answered with an error code, as opposed to a connection problem. + pub fn is_rejected_by_kafka(&self) -> bool { + matches!(self, Error::Kafka { .. }) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Topic { + /// The topic id in Kafka's string format (URL-safe base64 without padding), as shown by + /// `kafka-topics.sh` + pub id: String, + + /// Sorted by the partition index + pub partitions: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Partition { + pub index: i32, + /// Broker ids, including replicas that are added or removed by a reassignment + pub replicas: Vec, +} + +/// An ongoing reassignment of one partition. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct PartitionReassignment { + pub index: i32, + pub adding_replicas: Vec, + pub removing_replicas: Vec, +} + +pub struct KafkaAdminClient { + bootstrap_servers: Vec, + security: Security, + client_id: String, + /// The connection and the time of its last successful request + connection: Option<(Connection, Instant)>, +} + +impl KafkaAdminClient { + /// Creates the client without connecting. + pub fn new(bootstrap_servers: Vec, security: Security, client_id: &str) -> Self { + Self { + bootstrap_servers, + security, + client_id: client_id.to_owned(), + connection: None, + } + } + + /// Returns `None` if the topic does not exist. + pub async fn describe_topic(&mut self, topic: &str) -> Result> { + let response = self + .send(describe_topic_request(topic), METADATA_VERSION) + .await?; + describe_topic_result(response, topic) + } + + /// Returns the configs set on the topic itself, without broker or default values. + pub async fn describe_topic_configs( + &mut self, + topic: &str, + ) -> Result> { + let response = self + .send( + describe_topic_configs_request(topic), + DESCRIBE_CONFIGS_VERSION, + ) + .await?; + describe_topic_configs_result(response, topic) + } + + pub async fn create_topic( + &mut self, + topic: &str, + partitions: i32, + replication_factor: i16, + configs: &BTreeMap, + ) -> Result<()> { + let response = self + .send( + create_topic_request(topic, partitions, replication_factor, configs), + CREATE_TOPICS_VERSION, + ) + .await?; + + let result = first(response.topics, CreateTopicsRequest::KEY, topic)?; + check( + CreateTopicsRequest::KEY, + topic, + result.error_code, + result.error_message, + ) + } + + /// Increases the number of partitions to `partitions`. + pub async fn create_partitions(&mut self, topic: &str, partitions: i32) -> Result<()> { + let response = self + .send( + create_partitions_request(topic, partitions), + CREATE_PARTITIONS_VERSION, + ) + .await?; + + let result = first(response.results, CreatePartitionsRequest::KEY, topic)?; + check( + CreatePartitionsRequest::KEY, + topic, + result.error_code, + result.error_message, + ) + } + + /// Sets the `set` configs and resets the `delete` configs to their defaults. Other configs of + /// the topic stay untouched. + pub async fn incremental_alter_configs( + &mut self, + topic: &str, + set: &BTreeMap, + delete: &BTreeSet, + ) -> Result<()> { + let response = self + .send( + incremental_alter_configs_request(topic, set, delete), + INCREMENTAL_ALTER_CONFIGS_VERSION, + ) + .await?; + + let result = first( + response.responses, + IncrementalAlterConfigsRequest::KEY, + topic, + )?; + check( + IncrementalAlterConfigsRequest::KEY, + topic, + result.error_code, + result.error_message, + ) + } + + /// Succeeds if the topic does not exist. + pub async fn delete_topic(&mut self, topic: &str) -> Result<()> { + let response = self + .send(delete_topic_request(topic), DELETE_TOPICS_VERSION) + .await?; + delete_topic_result(response, topic) + } + + /// Returns the ongoing reassignments of the given partitions, sorted by the partition index. + /// Needs `DESCRIBE` on the cluster. + pub async fn list_partition_reassignments( + &mut self, + topic: &str, + partitions: &[i32], + ) -> Result> { + let response = self + .send( + list_partition_reassignments_request(topic, partitions), + LIST_PARTITION_REASSIGNMENTS_VERSION, + ) + .await?; + list_partition_reassignments_result(response, topic) + } + + /// Sends the request on the current connection or a new one. A failed or cancelled request + /// drops the connection, and a connection idle for longer than [`CONNECTION_MAX_IDLE`] is + /// replaced. + async fn send(&mut self, request: R, version: i16) -> Result { + let mut connection = match self.connection.take() { + Some((connection, last_used)) if last_used.elapsed() < CONNECTION_MAX_IDLE => { + connection + } + _ => self.connect().await?, + }; + let response = connection.send(request, version).await.context(SendSnafu)?; + self.connection = Some((connection, Instant::now())); + Ok(response) + } + + /// Connects to the first bootstrap server that answers. Any broker works, because brokers in + /// KRaft mode forward admin requests to the active controller. + async fn connect(&self) -> Result { + let mut last_error = Error::NoBootstrapServers; + for server in &self.bootstrap_servers { + match Connection::connect(&server.host, server.port, &self.security, &self.client_id) + .await + { + Ok(connection) => return Ok(connection), + Err(source) => { + last_error = Error::Connect { + source, + server: server.clone(), + }; + } + } + } + Err(last_error) + } +} + +fn describe_topic_request(topic: &str) -> MetadataRequest { + MetadataRequest::default() + .with_topics(Some(vec![ + MetadataRequestTopic::default().with_name(Some(topic_name(topic))), + ])) + .with_allow_auto_topic_creation(false) +} + +fn describe_topic_result(response: MetadataResponse, topic: &str) -> Result> { + let result = first(response.topics, MetadataRequest::KEY, topic)?; + if result.error_code.err() == Some(ResponseError::UnknownTopicOrPartition) { + return Ok(None); + } + check(MetadataRequest::KEY, topic, result.error_code, None)?; + + let mut partitions = result + .partitions + .into_iter() + .map(|partition| Partition { + index: partition.partition_index, + replicas: broker_ids(partition.replica_nodes), + }) + .collect::>(); + partitions.sort_by_key(|partition| partition.index); + Ok(Some(Topic { + id: URL_SAFE_NO_PAD.encode(result.topic_id.as_bytes()), + partitions, + })) +} + +fn describe_topic_configs_request(topic: &str) -> DescribeConfigsRequest { + DescribeConfigsRequest::default().with_resources(vec![ + DescribeConfigsResource::default() + .with_resource_type(RESOURCE_TYPE_TOPIC) + .with_resource_name(StrBytes::from_string(topic.to_owned())), + ]) +} + +fn describe_topic_configs_result( + response: DescribeConfigsResponse, + topic: &str, +) -> Result> { + let result = first(response.results, DescribeConfigsRequest::KEY, topic)?; + check( + DescribeConfigsRequest::KEY, + topic, + result.error_code, + result.error_message, + )?; + + Ok(result + .configs + .into_iter() + .filter(|config| config.config_source == CONFIG_SOURCE_DYNAMIC_TOPIC) + .filter_map(|config| Some((config.name.to_string(), config.value?.to_string()))) + .collect()) +} + +fn create_topic_request( + topic: &str, + partitions: i32, + replication_factor: i16, + configs: &BTreeMap, +) -> CreateTopicsRequest { + let configs = configs + .iter() + .map(|(name, value)| { + CreatableTopicConfig::default() + .with_name(StrBytes::from_string(name.clone())) + .with_value(Some(StrBytes::from_string(value.clone()))) + }) + .collect(); + CreateTopicsRequest::default() + .with_topics(vec![ + CreatableTopic::default() + .with_name(topic_name(topic)) + .with_num_partitions(partitions) + .with_replication_factor(replication_factor) + .with_configs(configs), + ]) + .with_timeout_ms(BROKER_TIMEOUT_MS) +} + +fn create_partitions_request(topic: &str, partitions: i32) -> CreatePartitionsRequest { + CreatePartitionsRequest::default() + .with_topics(vec![ + CreatePartitionsTopic::default() + .with_name(topic_name(topic)) + .with_count(partitions) + .with_assignments(None), + ]) + .with_timeout_ms(BROKER_TIMEOUT_MS) +} + +fn incremental_alter_configs_request( + topic: &str, + set: &BTreeMap, + delete: &BTreeSet, +) -> IncrementalAlterConfigsRequest { + let set = set.iter().map(|(name, value)| { + AlterableConfig::default() + .with_name(StrBytes::from_string(name.clone())) + .with_config_operation(CONFIG_OPERATION_SET) + .with_value(Some(StrBytes::from_string(value.clone()))) + }); + let delete = delete.iter().map(|name| { + AlterableConfig::default() + .with_name(StrBytes::from_string(name.clone())) + .with_config_operation(CONFIG_OPERATION_DELETE) + .with_value(None) + }); + // We can send all config operations (set & delete) together as one list + IncrementalAlterConfigsRequest::default().with_resources(vec![ + AlterConfigsResource::default() + .with_resource_type(RESOURCE_TYPE_TOPIC) + .with_resource_name(StrBytes::from_string(topic.to_owned())) + .with_configs(set.chain(delete).collect()), + ]) +} + +fn delete_topic_request(topic: &str) -> DeleteTopicsRequest { + DeleteTopicsRequest::default() + .with_topics(vec![ + DeleteTopicState::default().with_name(Some(topic_name(topic))), + ]) + .with_timeout_ms(BROKER_TIMEOUT_MS) +} + +fn delete_topic_result(response: DeleteTopicsResponse, topic: &str) -> Result<()> { + let result = first(response.responses, DeleteTopicsRequest::KEY, topic)?; + // It's ok if topic or partition are already being deleted + if matches!( + result.error_code.err(), + Some(ResponseError::UnknownTopicOrPartition | ResponseError::UnknownTopicId) + ) { + return Ok(()); + } + check( + DeleteTopicsRequest::KEY, + topic, + result.error_code, + result.error_message, + ) +} + +/// Only the given partitions are checked; an empty list checks none. +fn list_partition_reassignments_request( + topic: &str, + partitions: &[i32], +) -> ListPartitionReassignmentsRequest { + ListPartitionReassignmentsRequest::default() + .with_topics(Some(vec![ + ListPartitionReassignmentsTopics::default() + .with_name(topic_name(topic)) + .with_partition_indexes(partitions.to_vec()), + ])) + .with_timeout_ms(BROKER_TIMEOUT_MS) +} + +/// The response only contains partitions with an ongoing reassignment. +fn list_partition_reassignments_result( + response: ListPartitionReassignmentsResponse, + topic: &str, +) -> Result> { + check( + ListPartitionReassignmentsRequest::KEY, + topic, + response.error_code, + response.error_message, + )?; + + let mut reassignments = response + .topics + .into_iter() + .filter(|result| result.name == topic_name(topic)) + .flat_map(|result| result.partitions) + .map(|partition| PartitionReassignment { + index: partition.partition_index, + adding_replicas: broker_ids(partition.adding_replicas), + removing_replicas: broker_ids(partition.removing_replicas), + }) + .collect::>(); + reassignments.sort_by_key(|reassignment| reassignment.index); + Ok(reassignments) +} + +fn topic_name(topic: &str) -> TopicName { + TopicName::from(StrBytes::from_string(topic.to_owned())) +} + +fn broker_ids(ids: Vec) -> Vec { + ids.into_iter().map(|id| *id).collect() +} + +// Used instead of [0] to correctly handle empty results & take ownership +fn first(results: Vec, api_key: i16, topic: &str) -> Result { + results + .into_iter() + .next() + .context(MissingResultSnafu { api_key, topic }) +} + +fn check(api_key: i16, topic: &str, error_code: i16, message: Option) -> Result<()> { + match error_code.err() { + None => Ok(()), + Some(error) => KafkaSnafu { + api_key, + topic, + error, + message: message.map(|message| message.to_string()), + } + .fail(), + } +} + +#[cfg(test)] +mod tests { + use kafka_protocol::messages::{ + delete_topics_response::DeletableTopicResult, + describe_configs_response::{DescribeConfigsResourceResult, DescribeConfigsResult}, + list_partition_reassignments_response::{ + OngoingPartitionReassignment, OngoingTopicReassignment, + }, + metadata_response::{MetadataResponsePartition, MetadataResponseTopic}, + }; + use rstest::rstest; + use tokio::net::TcpStream; + + use super::*; + use crate::connection::test_support::{accept, listen, read_request_body, write_response}; + + const TOPIC: &str = "orders"; + const APIS: &[(i16, i16, i16)] = &[(DeleteTopicsRequest::KEY, 0, DELETE_TOPICS_VERSION)]; + + #[tokio::test] + async fn try_the_next_bootstrap_server() { + let (unreachable, unreachable_port) = listen().await; + drop(unreachable); + let (listener, port) = listen().await; + let mut client = kafka_admin_client(&[unreachable_port, port]); + + let broker_side = async { + let mut broker = accept(&listener, APIS).await; + answer_delete_topics(&mut broker).await; + }; + let (result, ()) = tokio::join!(client.delete_topic(TOPIC), broker_side); + + result.expect("the second bootstrap server should be used"); + } + + #[tokio::test] + async fn reconnect_after_a_failed_request() { + let (listener, port) = listen().await; + let mut client = kafka_admin_client(&[port]); + + let broker_side = async { + let mut broker = accept(&listener, APIS).await; + read_request_body::(&mut broker).await; + drop(broker); + }; + let (result, ()) = tokio::join!(client.delete_topic(TOPIC), broker_side); + assert!(matches!( + result.expect_err("the broker should have closed the connection"), + Error::Send { .. } + )); + + let broker_side = async { + let mut broker = accept(&listener, APIS).await; + answer_delete_topics(&mut broker).await; + }; + let (result, ()) = tokio::join!(client.delete_topic(TOPIC), broker_side); + + result.expect("the client should reconnect"); + } + + #[test] + fn describe_topic_by_name_without_auto_creation() { + let request = describe_topic_request(TOPIC); + + let topics = request.topics.expect("the topics should be requested"); + assert_eq!(Some(topic_name(TOPIC)), topics[0].name); + assert!(!request.allow_auto_topic_creation); + } + + #[test] + fn describe_existing_topic() { + let response = MetadataResponse::default().with_topics(vec![ + MetadataResponseTopic::default() + .with_name(Some(topic_name(TOPIC))) + .with_topic_id( + "517c94f6-2e40-467e-96d3-438ddad69e58" + .parse() + .expect("the topic id should be valid"), + ) + .with_partitions(vec![ + MetadataResponsePartition::default() + .with_partition_index(1) + .with_replica_nodes(vec![2.into(), 3.into()]), + MetadataResponsePartition::default() + .with_partition_index(0) + .with_replica_nodes(vec![1.into(), 2.into()]), + ]), + ]); + + assert_eq!( + Some(Topic { + id: "UXyU9i5ARn6W00ON2taeWA".to_owned(), + partitions: vec![ + Partition { + index: 0, + replicas: vec![1, 2] + }, + Partition { + index: 1, + replicas: vec![2, 3] + }, + ] + }), + describe_topic_result(response, TOPIC).expect("the topic should be described") + ); + } + + #[test] + fn describe_missing_topic() { + let response = MetadataResponse::default().with_topics(vec![ + MetadataResponseTopic::default() + .with_error_code(ResponseError::UnknownTopicOrPartition.code()), + ]); + + assert_eq!( + None, + describe_topic_result(response, TOPIC).expect("a missing topic should not be an error") + ); + } + + #[test] + fn describe_topic_configs_of_the_topic_resource() { + let request = describe_topic_configs_request(TOPIC); + + assert_eq!(RESOURCE_TYPE_TOPIC, request.resources[0].resource_type); + assert_eq!(TOPIC, request.resources[0].resource_name.as_str()); + } + + #[test] + fn describe_only_dynamic_topic_configs() { + let config = |name: &'static str, value: &'static str, source: i8| { + DescribeConfigsResourceResult::default() + .with_name(StrBytes::from_static_str(name)) + .with_value(Some(StrBytes::from_static_str(value))) + .with_config_source(source) + }; + let response = DescribeConfigsResponse::default().with_results(vec![ + DescribeConfigsResult::default().with_configs(vec![ + config("retention.ms", "1000", CONFIG_SOURCE_DYNAMIC_TOPIC), + config("cleanup.policy", "delete", 5), + ]), + ]); + + assert_eq!( + BTreeMap::from([("retention.ms".to_owned(), "1000".to_owned())]), + describe_topic_configs_result(response, TOPIC) + .expect("the configs should be described") + ); + } + + #[test] + fn create_topic_with_configs() { + let configs = BTreeMap::from([("retention.ms".to_owned(), "1000".to_owned())]); + + let request = create_topic_request(TOPIC, 3, 2, &configs); + + let topic = &request.topics[0]; + assert_eq!(topic_name(TOPIC), topic.name); + assert_eq!(3, topic.num_partitions); + assert_eq!(2, topic.replication_factor); + assert!(topic.assignments.is_empty()); + assert_eq!("retention.ms", topic.configs[0].name.as_str()); + assert_eq!( + Some("1000"), + topic.configs[0].value.as_ref().map(|value| value.as_str()) + ); + assert_eq!(BROKER_TIMEOUT_MS, request.timeout_ms); + } + + #[test] + fn create_partitions_without_assignments() { + let request = create_partitions_request(TOPIC, 6); + + assert_eq!(topic_name(TOPIC), request.topics[0].name); + assert_eq!(6, request.topics[0].count); + assert_eq!(None, request.topics[0].assignments); + assert_eq!(BROKER_TIMEOUT_MS, request.timeout_ms); + } + + #[test] + fn set_and_delete_configs_in_one_request() { + let set = BTreeMap::from([("retention.ms".to_owned(), "1000".to_owned())]); + let delete = BTreeSet::from(["cleanup.policy".to_owned()]); + + let request = incremental_alter_configs_request(TOPIC, &set, &delete); + + let resource = &request.resources[0]; + assert_eq!(RESOURCE_TYPE_TOPIC, resource.resource_type); + assert_eq!(TOPIC, resource.resource_name.as_str()); + let configs = resource + .configs + .iter() + .map(|config| { + ( + config.name.as_str(), + config.config_operation, + config.value.as_ref().map(|value| value.as_str()), + ) + }) + .collect::>(); + assert_eq!( + vec![ + ("retention.ms", CONFIG_OPERATION_SET, Some("1000")), + ("cleanup.policy", CONFIG_OPERATION_DELETE, None), + ], + configs + ); + } + + #[test] + fn list_reassignments_of_the_given_partitions() { + let request = list_partition_reassignments_request(TOPIC, &[0, 1, 2]); + + let topics = request.topics.expect("the topics should be requested"); + assert_eq!(topic_name(TOPIC), topics[0].name); + assert_eq!(vec![0, 1, 2], topics[0].partition_indexes); + assert_eq!(BROKER_TIMEOUT_MS, request.timeout_ms); + } + + #[test] + fn read_ongoing_reassignments() { + let reassignment = |index: i32, adding: Vec, removing: Vec| { + OngoingPartitionReassignment::default() + .with_partition_index(index) + .with_adding_replicas(adding.into_iter().map(Into::into).collect()) + .with_removing_replicas(removing.into_iter().map(Into::into).collect()) + }; + let response = ListPartitionReassignmentsResponse::default().with_topics(vec![ + OngoingTopicReassignment::default() + .with_name(topic_name("other")) + .with_partitions(vec![reassignment(0, vec![9], vec![8])]), + OngoingTopicReassignment::default() + .with_name(topic_name(TOPIC)) + .with_partitions(vec![ + reassignment(2, vec![4], vec![]), + reassignment(0, vec![4, 5], vec![1, 2]), + ]), + ]); + + assert_eq!( + vec![ + PartitionReassignment { + index: 0, + adding_replicas: vec![4, 5], + removing_replicas: vec![1, 2], + }, + PartitionReassignment { + index: 2, + adding_replicas: vec![4], + removing_replicas: vec![], + }, + ], + list_partition_reassignments_result(response, TOPIC) + .expect("the reassignments should be readable") + ); + } + + #[test] + fn read_no_ongoing_reassignments() { + assert_eq!( + Vec::::new(), + list_partition_reassignments_result( + ListPartitionReassignmentsResponse::default(), + TOPIC + ) + .expect("no reassignments should not be an error") + ); + } + + #[test] + fn reject_reassignments_without_cluster_permission() { + let response = ListPartitionReassignmentsResponse::default() + .with_error_code(ResponseError::ClusterAuthorizationFailed.code()); + + let error = list_partition_reassignments_result(response, TOPIC) + .expect_err("the request should be rejected"); + + assert!(matches!( + error, + Error::Kafka { + error: ResponseError::ClusterAuthorizationFailed, + .. + } + )); + } + + #[test] + fn delete_topic_by_name() { + let request = delete_topic_request(TOPIC); + + assert_eq!(Some(topic_name(TOPIC)), request.topics[0].name); + assert!(request.topics[0].topic_id.is_nil()); + assert_eq!(BROKER_TIMEOUT_MS, request.timeout_ms); + } + + #[rstest] + #[case(ResponseError::UnknownTopicOrPartition)] + #[case(ResponseError::UnknownTopicId)] + fn delete_missing_topic(#[case] error: ResponseError) { + let response = DeleteTopicsResponse::default().with_responses(vec![ + DeletableTopicResult::default().with_error_code(error.code()), + ]); + + delete_topic_result(response, TOPIC).expect("deleting a missing topic should succeed"); + } + + #[test] + fn report_rejected_request() { + let error = check( + CreateTopicsRequest::KEY, + TOPIC, + ResponseError::TopicAlreadyExists.code(), + Some(StrBytes::from_static_str("Topic 'orders' already exists.")), + ) + .expect_err("the request should be rejected"); + + assert!(error.is_rejected_by_kafka()); + assert!(error.to_string().contains("Topic 'orders' already exists.")); + } + + #[test] + fn do_not_treat_connection_errors_as_rejected_by_kafka() { + let error = Error::NoBootstrapServers; + + assert!(!error.is_rejected_by_kafka()); + } + + #[test] + fn reject_response_without_result() { + let error = first( + Vec::::new(), + DeleteTopicsRequest::KEY, + TOPIC, + ) + .expect_err("the result should be missing"); + + assert!(matches!(error, Error::MissingResult { .. })); + } + + fn kafka_admin_client(ports: &[u16]) -> KafkaAdminClient { + let bootstrap_servers = ports + .iter() + .map(|port| BrokerAddress { + host: "127.0.0.1".to_owned(), + port: *port, + }) + .collect(); + KafkaAdminClient::new(bootstrap_servers, Security::Plaintext, "test") + } + + async fn answer_delete_topics(broker: &mut TcpStream) { + let (header, _) = read_request_body::(broker).await; + let response = + DeleteTopicsResponse::default().with_responses(vec![DeletableTopicResult::default()]); + write_response(broker, &header, header.correlation_id, response).await; + } +} diff --git a/rust/kafka-admin/src/codec.rs b/rust/kafka-admin/src/codec.rs new file mode 100644 index 00000000..0c69347d --- /dev/null +++ b/rust/kafka-admin/src/codec.rs @@ -0,0 +1,121 @@ +//! Encoding of request and decoding of response frames: a header followed by a message body. The +//! length prefix is added and removed by the connection. + +use bytes::{Bytes, BytesMut}; +use kafka_protocol::{ + messages::{RequestHeader, ResponseHeader}, + protocol::{Decodable, Encodable, HeaderVersion, Request, StrBytes}, +}; +use snafu::{ResultExt, Snafu, ensure}; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("failed to encode the request with API key {api_key}"))] + EncodeRequest { source: anyhow::Error, api_key: i16 }, + + #[snafu(display("failed to decode the response with API key {api_key}"))] + DecodeResponse { source: anyhow::Error, api_key: i16 }, + + #[snafu(display("expected correlation id {expected}, received {received}"))] + CorrelationIdMismatch { expected: i32, received: i32 }, +} + +pub fn encode_request( + request: &R, + version: i16, + correlation_id: i32, + client_id: &str, +) -> Result { + let header = RequestHeader::default() + .with_request_api_key(R::KEY) + .with_request_api_version(version) + .with_correlation_id(correlation_id) + .with_client_id(Some(StrBytes::from_string(client_id.to_owned()))); + + let mut frame = BytesMut::new(); + header + .encode(&mut frame, R::header_version(version)) + .context(EncodeRequestSnafu { api_key: R::KEY })?; + request + .encode(&mut frame, version) + .context(EncodeRequestSnafu { api_key: R::KEY })?; + Ok(frame.freeze()) +} + +pub fn decode_response( + mut frame: Bytes, + version: i16, + correlation_id: i32, +) -> Result { + let header = ResponseHeader::decode(&mut frame, R::Response::header_version(version)) + .context(DecodeResponseSnafu { api_key: R::KEY })?; + ensure!( + header.correlation_id == correlation_id, + CorrelationIdMismatchSnafu { + expected: correlation_id, + received: header.correlation_id, + } + ); + R::Response::decode(&mut frame, version).context(DecodeResponseSnafu { api_key: R::KEY }) +} + +#[cfg(test)] +mod tests { + use kafka_protocol::{ + messages::{MetadataRequest, MetadataResponse}, + protocol::decode_request_header_from_buffer, + }; + + use super::*; + + #[test] + fn encode_request_with_header() { + let mut frame = encode_request(&MetadataRequest::default(), 12, 7, "agent") + .expect("the request should be encodable"); + + let header = + decode_request_header_from_buffer(&mut frame).expect("the header should be decodable"); + assert_eq!(3, header.request_api_key); + assert_eq!(12, header.request_api_version); + assert_eq!(7, header.correlation_id); + assert_eq!(Some("agent"), header.client_id.as_deref()); + } + + fn response_frame(correlation_id: i32, version: i16, response: &MetadataResponse) -> Bytes { + let mut frame = BytesMut::new(); + ResponseHeader::default() + .with_correlation_id(correlation_id) + .encode(&mut frame, MetadataResponse::header_version(version)) + .expect("the header should be encodable"); + response + .encode(&mut frame, version) + .expect("the response should be encodable"); + frame.freeze() + } + + #[test] + fn decode_response_with_matching_correlation_id() { + let response = MetadataResponse::default().with_controller_id(2.into()); + + let decoded = decode_response::(response_frame(7, 12, &response), 12, 7) + .expect("the response should be decodable"); + + assert_eq!(response, decoded); + } + + #[test] + fn reject_response_with_other_correlation_id() { + let frame = response_frame(8, 12, &MetadataResponse::default()); + + let error = decode_response::(frame, 12, 7) + .expect_err("the correlation id should not match"); + + assert!(matches!( + error, + Error::CorrelationIdMismatch { + expected: 7, + received: 8 + } + )); + } +} diff --git a/rust/kafka-admin/src/connection.rs b/rust/kafka-admin/src/connection.rs new file mode 100644 index 00000000..b6f97c0b --- /dev/null +++ b/rust/kafka-admin/src/connection.rs @@ -0,0 +1,581 @@ +//! A connection to a single Kafka broker. + +use std::{collections::BTreeMap, io, path::PathBuf, sync::Arc, time::Duration}; + +use futures::{SinkExt, StreamExt}; +use kafka_protocol::{ + ResponseError, + error::ParseResponseErrorCode, + messages::{ApiVersionsRequest, ApiVersionsResponse}, + protocol::{Request, StrBytes, VersionRange}, +}; +use rustls::{ClientConfig, pki_types::ServerName}; +use snafu::{OptionExt, ResultExt, Snafu, ensure}; +use tokio::{ + net::TcpStream, + time::{error::Elapsed, timeout}, +}; +use tokio_rustls::{TlsConnector, client::TlsStream}; +use tokio_util::{ + codec::{Framed, LengthDelimitedCodec}, + either::Either, +}; + +use crate::{ + codec::{self, decode_response, encode_request}, + tls, +}; + +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +const MAX_FRAME_SIZE: usize = 64 * 1024 * 1024; +const API_VERSIONS_VERSION: i16 = 3; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("failed to connect to {address:?}"))] + Connect { source: io::Error, address: String }, + + #[snafu(display("timed out connecting to {address:?}"))] + ConnectTimeout { source: Elapsed, address: String }, + + #[snafu(display("failed to configure TLS"))] + Tls { source: tls::Error }, + + #[snafu(display("{host:?} is not a valid TLS server name"))] + InvalidServerName { host: String }, + + #[snafu(display("failed the TLS handshake with {address:?}"))] + TlsHandshake { source: io::Error, address: String }, + + #[snafu(display("failed to write to the broker"))] + Write { source: io::Error }, + + #[snafu(display("failed to read from the broker"))] + Read { source: io::Error }, + + #[snafu(display("the broker closed the connection"))] + ConnectionClosed, + + #[snafu(display("the connection is broken by an earlier failed or cancelled request"))] + ConnectionBroken, + + #[snafu(display("timed out waiting for the response to the request with API key {api_key}"))] + RequestTimeout { source: Elapsed, api_key: i16 }, + + #[snafu(display("failed to encode or decode a message"))] + Codec { source: codec::Error }, + + #[snafu(display("the broker returned {error:?} for the request with API key {api_key}"))] + Kafka { error: ResponseError, api_key: i16 }, + + #[snafu(display("the broker does not support the API key {api_key}"))] + UnsupportedApi { api_key: i16 }, + + #[snafu(display( + "the broker supports the versions {broker_versions} of the API key {api_key}, but not {version}" + ))] + UnsupportedVersion { + api_key: i16, + version: i16, + broker_versions: VersionRange, + }, +} + +type Result = std::result::Result; + +/// How to connect to the broker. Paths point to PEM files. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum Security { + Plaintext, + Tls { + ca: PathBuf, + }, + MutualTls { + ca: PathBuf, + certificate: PathBuf, + key: PathBuf, + }, +} + +pub struct Connection { + stream: Framed>, LengthDelimitedCodec>, + client_id: String, + correlation_id: i32, + awaiting_response: bool, + broker_api_versions: BTreeMap, +} + +impl Connection { + /// Connects to the broker and fetches the API versions it supports. + pub async fn connect( + host: &str, + port: u16, + security: &Security, + client_id: &str, + ) -> Result { + let address = format!("{host}:{port}"); + let tcp_stream = connect_tcp(host, port, &address).await?; + let stream = match tls::client_config(security).context(TlsSnafu)? { + None => Either::Left(tcp_stream), + Some(config) => Either::Right(tls_handshake(tcp_stream, host, &address, config).await?), + }; + + let mut connection = Self::new(stream, client_id); + let response = connection + .send_with_version(api_versions_request(), API_VERSIONS_VERSION) + .await?; + connection.broker_api_versions = api_versions_result(response)?; + Ok(connection) + } + + fn new(stream: Either>, client_id: &str) -> Self { + let codec = LengthDelimitedCodec::builder() + .length_field_type::() + .max_frame_length(MAX_FRAME_SIZE) + .new_codec(); + Self { + stream: Framed::new(stream, codec), + client_id: client_id.to_owned(), + correlation_id: 0, + awaiting_response: false, + broker_api_versions: BTreeMap::new(), + } + } + + /// Sends the request with the given version, if the broker supports it. + pub async fn send(&mut self, request: R, version: i16) -> Result { + // Each API that is relevant for us has a version that is supported by all of our Kafka versions. + // We use these versions and pin them in client.rs. + // We will have to add API version negotiation at some point in the future. + check_version(&self.broker_api_versions, R::KEY, version)?; + self.send_with_version(request, version).await + } + + async fn send_with_version( + &mut self, + request: R, + version: i16, + ) -> Result { + ensure!(!self.awaiting_response, ConnectionBrokenSnafu); + self.awaiting_response = true; + + self.correlation_id = self.correlation_id.wrapping_add(1); + let frame = encode_request(&request, version, self.correlation_id, &self.client_id) + .context(CodecSnafu)?; + + let response = timeout(REQUEST_TIMEOUT, async { + self.stream.send(frame).await.context(WriteSnafu)?; + self.stream + .next() + .await + .context(ConnectionClosedSnafu)? + .context(ReadSnafu) + }) + .await + .context(RequestTimeoutSnafu { api_key: R::KEY })??; + + let response = decode_response::(response.freeze(), version, self.correlation_id) + .context(CodecSnafu)?; + + self.awaiting_response = false; + Ok(response) + } +} + +async fn connect_tcp(host: &str, port: u16, address: &str) -> Result { + timeout(CONNECT_TIMEOUT, TcpStream::connect((host, port))) + .await + .context(ConnectTimeoutSnafu { address })? + .context(ConnectSnafu { address }) +} + +fn api_versions_request() -> ApiVersionsRequest { + ApiVersionsRequest::default() + .with_client_software_name(StrBytes::from_static_str("stackable-kafka-admin")) + .with_client_software_version(StrBytes::from_static_str(env!("CARGO_PKG_VERSION"))) +} + +fn api_versions_result(response: ApiVersionsResponse) -> Result> { + if let Some(error) = response.error_code.err() { + return KafkaSnafu { + error, + api_key: ApiVersionsRequest::KEY, + } + .fail(); + } + + Ok(response + .api_keys + .iter() + .map(|api| { + ( + api.api_key, + VersionRange { + min: api.min_version, + max: api.max_version, + }, + ) + }) + .collect()) +} + +fn check_version( + broker_api_versions: &BTreeMap, + api_key: i16, + version: i16, +) -> Result<()> { + let broker_versions = broker_api_versions + .get(&api_key) + .context(UnsupportedApiSnafu { api_key })?; + ensure!( + (broker_versions.min..=broker_versions.max).contains(&version), + UnsupportedVersionSnafu { + api_key, + version, + broker_versions: *broker_versions, + } + ); + Ok(()) +} + +async fn tls_handshake( + tcp_stream: TcpStream, + host: &str, + address: &str, + config: ClientConfig, +) -> Result> { + let server_name = ServerName::try_from(host.to_owned()) + .ok() + .context(InvalidServerNameSnafu { host })?; + timeout( + CONNECT_TIMEOUT, + TlsConnector::from(Arc::new(config)).connect(server_name, tcp_stream), + ) + .await + .context(ConnectTimeoutSnafu { address })? + .context(TlsHandshakeSnafu { address }) +} + +/// A mock Kafka broker for tests, speaking the wire protocol over a local TCP socket. +#[cfg(test)] +pub(crate) mod test_support { + use bytes::{BufMut, Bytes, BytesMut}; + use kafka_protocol::{ + messages::{ + ApiVersionsRequest, ApiVersionsResponse, RequestHeader, ResponseHeader, + api_versions_response::ApiVersion, + }, + protocol::{Encodable, HeaderVersion, Request, decode_request_header_from_buffer}, + }; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + }; + + use super::API_VERSIONS_VERSION; + + pub async fn listen() -> (TcpListener, u16) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("the mock broker should be able to listen"); + let port = listener + .local_addr() + .expect("the mock broker should have an address") + .port(); + (listener, port) + } + + /// Accepts a connection and answers its ApiVersions request with the given + /// `(api_key, min_version, max_version)` ranges. + pub async fn accept(listener: &TcpListener, apis: &[(i16, i16, i16)]) -> TcpStream { + let (mut broker, _) = listener + .accept() + .await + .expect("the mock broker should accept the connection"); + let (request, _) = read_request(&mut broker).await; + assert_eq!(ApiVersionsRequest::KEY, request.request_api_key); + assert_eq!(API_VERSIONS_VERSION, request.request_api_version); + write_response( + &mut broker, + &request, + request.correlation_id, + api_versions(apis), + ) + .await; + broker + } + + pub async fn read_request(broker: &mut TcpStream) -> (RequestHeader, Bytes) { + let size = broker + .read_i32() + .await + .expect("the request size should be readable"); + let mut frame = vec![0; size as usize]; + broker + .read_exact(&mut frame) + .await + .expect("the request should be readable"); + let mut frame = Bytes::from(frame); + let header = decode_request_header_from_buffer(&mut frame) + .expect("the request header should be decodable"); + (header, frame) + } + + pub async fn read_request_body(broker: &mut TcpStream) -> (RequestHeader, R) { + let (header, mut body) = read_request(broker).await; + assert_eq!(R::KEY, header.request_api_key); + let request = R::decode(&mut body, header.request_api_version) + .expect("the request should be decodable"); + (header, request) + } + + pub async fn write_response( + broker: &mut TcpStream, + request: &RequestHeader, + correlation_id: i32, + response: M, + ) { + let mut body = BytesMut::new(); + ResponseHeader::default() + .with_correlation_id(correlation_id) + .encode(&mut body, M::header_version(request.request_api_version)) + .expect("the response header should be encodable"); + response + .encode(&mut body, request.request_api_version) + .expect("the response should be encodable"); + let mut frame = BytesMut::new(); + frame.put_i32(body.len() as i32); + frame.extend_from_slice(&body); + broker + .write_all(&frame) + .await + .expect("the response should be writable"); + } + + fn api_versions(apis: &[(i16, i16, i16)]) -> ApiVersionsResponse { + ApiVersionsResponse::default().with_api_keys( + apis.iter() + .map(|(api_key, min_version, max_version)| { + ApiVersion::default() + .with_api_key(*api_key) + .with_min_version(*min_version) + .with_max_version(*max_version) + }) + .collect(), + ) + } +} + +#[cfg(test)] +mod tests { + use kafka_protocol::messages::{ + MetadataRequest, MetadataResponse, api_versions_response::ApiVersion, + }; + use tokio::io::AsyncWriteExt; + + use super::{ + test_support::{accept, listen, read_request, write_response}, + *, + }; + + /// Connects to a mock broker supporting Metadata versions 0 to 13. + async fn connect() -> (Connection, TcpStream) { + let (listener, port) = listen().await; + let (connection, broker) = tokio::join!( + Connection::connect("127.0.0.1", port, &Security::Plaintext, "test"), + accept(&listener, &[(MetadataRequest::KEY, 0, 13)]) + ); + (connection.expect("the connection should succeed"), broker) + } + + #[test] + fn read_api_versions() { + let response = ApiVersionsResponse::default().with_api_keys(vec![ + ApiVersion::default() + .with_api_key(MetadataRequest::KEY) + .with_min_version(0) + .with_max_version(13), + ]); + + assert_eq!( + BTreeMap::from([(MetadataRequest::KEY, VersionRange { min: 0, max: 13 })]), + api_versions_result(response).expect("the API versions should be valid") + ); + } + + #[test] + fn reject_api_versions_with_error() { + let response = ApiVersionsResponse::default() + .with_error_code(ResponseError::UnsupportedVersion.code()); + + let error = api_versions_result(response).expect_err("the response should be rejected"); + + assert!(matches!( + error, + Error::Kafka { + error: ResponseError::UnsupportedVersion, + .. + } + )); + } + + #[test] + fn accept_supported_version() { + let broker_api_versions = + BTreeMap::from([(MetadataRequest::KEY, VersionRange { min: 0, max: 13 })]); + + check_version(&broker_api_versions, MetadataRequest::KEY, 12) + .expect("the version should be supported"); + } + + #[test] + fn reject_api_unknown_to_the_broker() { + let error = check_version(&BTreeMap::new(), MetadataRequest::KEY, 12) + .expect_err("the API should be unsupported"); + + assert!(matches!(error, Error::UnsupportedApi { .. })); + } + + #[test] + fn reject_version_unsupported_by_the_broker() { + let broker_api_versions = + BTreeMap::from([(MetadataRequest::KEY, VersionRange { min: 0, max: 9 })]); + + let error = check_version(&broker_api_versions, MetadataRequest::KEY, 12) + .expect_err("the version should be unsupported"); + + assert!(matches!( + error, + Error::UnsupportedVersion { version: 12, .. } + )); + } + + #[tokio::test] + async fn send_request_and_receive_response() { + let (mut connection, mut broker) = connect().await; + + let broker_side = async { + let (request, _) = read_request(&mut broker).await; + assert_eq!(MetadataRequest::KEY, request.request_api_key); + assert_eq!(12, request.request_api_version); + write_response( + &mut broker, + &request, + request.correlation_id, + MetadataResponse::default().with_controller_id(2.into()), + ) + .await; + }; + let (response, ()) = + tokio::join!(connection.send(MetadataRequest::default(), 12), broker_side); + + assert_eq!( + 2, + *response.expect("the request should succeed").controller_id + ); + } + + #[tokio::test] + async fn reject_response_with_other_correlation_id() { + let (mut connection, mut broker) = connect().await; + + let broker_side = async { + let (request, _) = read_request(&mut broker).await; + write_response( + &mut broker, + &request, + request.correlation_id + 1, + MetadataResponse::default(), + ) + .await; + }; + let (response, ()) = + tokio::join!(connection.send(MetadataRequest::default(), 12), broker_side); + + assert!(matches!( + response.expect_err("the correlation id should not match"), + Error::Codec { + source: codec::Error::CorrelationIdMismatch { .. } + } + )); + } + + #[tokio::test] + async fn reject_invalid_frame_size() { + let (mut connection, mut broker) = connect().await; + + let broker_side = async { + read_request(&mut broker).await; + broker + .write_i32(-1) + .await + .expect("the frame size should be writable"); + }; + let (response, ()) = + tokio::join!(connection.send(MetadataRequest::default(), 12), broker_side); + + assert!(matches!( + response.expect_err("the frame size should be invalid"), + Error::Read { .. } + )); + } + + #[tokio::test] + async fn time_out_without_response() { + let (mut connection, mut broker) = connect().await; + tokio::time::pause(); + + let broker_side = async { + read_request(&mut broker).await; + }; + let (response, ()) = + tokio::join!(connection.send(MetadataRequest::default(), 12), broker_side); + + assert!(matches!( + response.expect_err("the request should time out"), + Error::RequestTimeout { .. } + )); + } + + #[tokio::test] + async fn reject_requests_after_a_failed_request() { + let (mut connection, mut broker) = connect().await; + + let broker_side = async { + read_request(&mut broker).await; + drop(broker); + }; + let (failed, ()) = + tokio::join!(connection.send(MetadataRequest::default(), 12), broker_side); + assert!(matches!( + failed.expect_err("the broker should have closed the connection"), + Error::ConnectionClosed + )); + + let error = connection + .send(MetadataRequest::default(), 12) + .await + .expect_err("the connection should be broken"); + + assert!(matches!(error, Error::ConnectionBroken)); + } + + #[tokio::test] + async fn reject_requests_after_a_cancelled_request() { + let (mut connection, _broker) = connect().await; + + { + let request = connection.send(MetadataRequest::default(), 12); + tokio::pin!(request); + let _ = futures::poll!(request.as_mut()); + } + + let error = connection + .send(MetadataRequest::default(), 12) + .await + .expect_err("the connection should be broken"); + + assert!(matches!(error, Error::ConnectionBroken)); + } +} diff --git a/rust/kafka-admin/src/lib.rs b/rust/kafka-admin/src/lib.rs new file mode 100644 index 00000000..cb1f00ea --- /dev/null +++ b/rust/kafka-admin/src/lib.rs @@ -0,0 +1,7 @@ +#![allow(clippy::result_large_err)] + +pub mod bootstrap; +pub mod client; +mod codec; +pub mod connection; +pub mod tls; diff --git a/rust/kafka-admin/src/tls.rs b/rust/kafka-admin/src/tls.rs new file mode 100644 index 00000000..6e23a512 --- /dev/null +++ b/rust/kafka-admin/src/tls.rs @@ -0,0 +1,106 @@ +//! TLS client configuration from PEM files. + +use std::{ + path::{Path, PathBuf}, + sync::Arc, +}; + +use rustls::{ + ClientConfig, RootCertStore, + pki_types::{CertificateDer, PrivateKeyDer, pem::PemObject}, +}; +use snafu::{ResultExt, Snafu, ensure}; + +use crate::connection::Security; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("failed to read the PEM file {path:?}"))] + ReadPem { + source: rustls::pki_types::pem::Error, + path: PathBuf, + }, + + #[snafu(display("the PEM file {path:?} contains no certificate"))] + NoCertificate { path: PathBuf }, + + #[snafu(display("failed to add the CA certificate from {path:?}"))] + AddCaCertificate { + source: rustls::Error, + path: PathBuf, + }, + + #[snafu(display("failed to build the TLS configuration"))] + BuildConfig { source: rustls::Error }, +} + +type Result = std::result::Result; + +/// Returns `None` for plaintext connections. The PEM files are read on every call. +pub(crate) fn client_config(security: &Security) -> Result> { + let (ca, client_certificate) = match security { + Security::Plaintext => return Ok(None), + Security::Tls { ca } => (ca, None), + Security::MutualTls { + ca, + certificate, + key, + } => (ca, Some((certificate, key))), + }; + + let builder = + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .context(BuildConfigSnafu)? + .with_root_certificates(root_certificates(ca)?); + + let config = match client_certificate { + None => builder.with_no_client_auth(), + Some((certificate, key)) => { + let key = PrivateKeyDer::from_pem_file(key).context(ReadPemSnafu { path: key })?; + builder + .with_client_auth_cert(certificate_chain(certificate)?, key) + .context(BuildConfigSnafu)? + } + }; + Ok(Some(config)) +} + +fn root_certificates(ca: &Path) -> Result { + let mut root_certificates = RootCertStore::empty(); + for certificate in CertificateDer::pem_file_iter(ca).context(ReadPemSnafu { path: ca })? { + root_certificates + .add(certificate.context(ReadPemSnafu { path: ca })?) + .context(AddCaCertificateSnafu { path: ca })?; + } + ensure!( + !root_certificates.is_empty(), + NoCertificateSnafu { path: ca } + ); + Ok(root_certificates) +} + +fn certificate_chain(path: &Path) -> Result>> { + let certificate_chain = CertificateDer::pem_file_iter(path) + .context(ReadPemSnafu { path })? + .collect::, _>>() + .context(ReadPemSnafu { path })?; + ensure!(!certificate_chain.is_empty(), NoCertificateSnafu { path }); + Ok(certificate_chain) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reject_ca_file_without_certificate() { + let ca = std::env::temp_dir().join(format!("empty-ca-{}.crt", std::process::id())); + std::fs::write(&ca, "").expect("the CA file should be writable"); + + let error = root_certificates(&ca).expect_err("the CA file should be rejected"); + std::fs::remove_file(&ca).expect("the CA file should be removable"); + + assert!(matches!(error, Error::NoCertificate { .. })); + } +}