rustfs_targets/
check.rs

1// Copyright 2024 RustFS Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15/// Check if MQTT Broker is available
16///
17/// # Arguments
18/// * `broker_url` - URL of MQTT Broker, for example `mqtt://localhost:1883`
19/// * `topic` - Topic for testing connections
20/// * `username` - Optional username for authentication
21/// * `password` - Optional password for authentication
22/// # Returns
23/// * `Ok(())` - If the connection is successful
24/// * `Err(TargetError)` - If the check fails.
25///   `TargetError::Configuration` indicates a bad configuration (invalid URL, TLS settings, etc.).
26///   Other variants indicate a connectivity or runtime failure.
27///
28/// # Example
29/// ```rust,no_run
30///  #[tokio::main]
31///  async fn main() {
32///     let result = rustfs_targets::check_mqtt_broker_available(
33///         "mqtt://localhost:1883",
34///         "test/topic",
35///         Some("myuser"),
36///         Some("mypass"),
37///     ).await;
38///     if result.is_ok() {
39///         println!("MQTT Broker is available");
40///     } else {
41///         println!("MQTT Broker is not available: {}", result.err().unwrap());
42///     }
43///  }
44/// ```
45///
46pub async fn check_mqtt_broker_available(
47    broker_url: &str,
48    topic: &str,
49    username: Option<&str>,
50    password: Option<&str>,
51) -> Result<(), crate::TargetError> {
52    use crate::target::mqtt::MQTTTlsConfig;
53
54    check_mqtt_broker_available_with_tls(broker_url, topic, username, password, &MQTTTlsConfig::default()).await
55}
56
57pub async fn check_mqtt_broker_available_with_tls(
58    broker_url: &str,
59    topic: &str,
60    username: Option<&str>,
61    password: Option<&str>,
62    tls: &crate::target::mqtt::MQTTTlsConfig,
63) -> Result<(), crate::TargetError> {
64    use crate::target::mqtt::build_mqtt_options;
65    use rumqttc::{AsyncClient, QoS};
66
67    let url =
68        crate::parse_url(broker_url).map_err(|e| crate::TargetError::Configuration(format!("Broker URL parsing failed: {e}")))?;
69    let url = url.url();
70
71    // build_mqtt_options returns TargetError directly; Configuration variants propagate as-is.
72    let mqtt_options = build_mqtt_options(
73        "rustfs_check".to_string(),
74        url,
75        username,
76        password,
77        tls,
78        std::time::Duration::from_secs(5),
79        None,
80    )?;
81    let (client, mut eventloop) = AsyncClient::builder(mqtt_options).capacity(1).build();
82
83    // Try to connect and subscribe
84    client
85        .subscribe(topic, QoS::AtLeastOnce)
86        .await
87        .map_err(|e| crate::TargetError::Network(format!("MQTT subscription failed: {e}")))?;
88    // Wait for eventloop to receive at least one event
89    match tokio::time::timeout(std::time::Duration::from_secs(3), eventloop.poll()).await {
90        Ok(Ok(_)) => Ok(()),
91        Ok(Err(e)) => Err(crate::TargetError::Network(format!("MQTT connection failed: {e}"))),
92        Err(_) => Err(crate::TargetError::Timeout("MQTT connection timed out".to_string())),
93    }
94}
95
96pub async fn check_nats_server_available(args: &crate::target::nats::NATSArgs) -> Result<(), crate::TargetError> {
97    tokio::time::timeout(std::time::Duration::from_secs(5), async {
98        let client = crate::target::nats::connect_nats(args).await?;
99        client
100            .flush()
101            .await
102            .map_err(|e| crate::TargetError::Network(format!("NATS connection check failed: {e}")))?;
103        // Validate the configured stream on the live connection before the drain closes it, so a
104        // missing or non-writable stream is rejected here rather than at first publish. The lookup is
105        // read-only and never creates a stream. The drain runs regardless of the validation outcome so
106        // the check connection is closed gracefully, and the validation error then propagates.
107        let validation = if args.jetstream_enable.unwrap_or(false) {
108            let mut context = async_nats::jetstream::new(client.clone());
109            // Match every other context construction site: the ack timeout is the per-operation await
110            // bound, so the validation get_stream lookup uses it rather than the library default.
111            context.set_timeout(std::time::Duration::from_secs(
112                args.jetstream_ack_timeout_secs
113                    .unwrap_or(rustfs_config::NATS_JETSTREAM_ACK_TIMEOUT_DEFAULT_SECS),
114            ));
115            crate::target::nats::validate_jetstream_stream(&context, args, "nats-check", None).await
116        } else {
117            Ok(())
118        };
119        client
120            .drain()
121            .await
122            .map_err(|e| crate::TargetError::Network(format!("Failed to close NATS check connection: {e}")))?;
123        validation
124    })
125    .await
126    .unwrap_or_else(|_| Err(crate::TargetError::Timeout("NATS connection timed out".to_string())))
127}
128
129pub async fn check_pulsar_broker_available(args: &crate::target::pulsar::PulsarArgs) -> Result<(), crate::TargetError> {
130    tokio::time::timeout(std::time::Duration::from_secs(5), async {
131        let client = crate::target::pulsar::connect_pulsar(args).await?;
132        client
133            .lookup_partitioned_topic(args.topic.clone())
134            .await
135            .map_err(|e| crate::TargetError::Network(format!("Pulsar topic lookup failed: {e}")))?;
136        Ok(())
137    })
138    .await
139    .unwrap_or_else(|_| Err(crate::TargetError::Timeout("Pulsar connection timed out".to_string())))
140}
141
142/// Probes a MySQL server for connectivity.
143///
144/// 1. Validates `args`.
145/// 2. Parses the DSN and builds a connection pool.
146/// 3. Runs `SELECT 1` to confirm credentials work.
147pub async fn check_mysql_server_available(args: &crate::target::mysql::MySqlArgs) -> Result<(), crate::TargetError> {
148    use crate::target::ensure_rustls_provider_installed;
149    use crate::target::mysql::{MySqlDsn, map_mysql_error};
150    use mysql_async::{Opts, OptsBuilder, Pool, SslOpts, prelude::Queryable};
151    use std::path::PathBuf;
152
153    args.validate()?;
154
155    let dsn = MySqlDsn::parse(&args.dsn_string)?;
156
157    let mut builder = OptsBuilder::default()
158        .user(Some(dsn.user.clone()))
159        .pass(Some(dsn.password.clone()))
160        .ip_or_hostname(dsn.host.clone())
161        .tcp_port(dsn.port)
162        .db_name(Some(dsn.database.clone()));
163
164    if dsn.tls {
165        ensure_rustls_provider_installed();
166        let mut ssl_opts = SslOpts::default();
167        if !args.tls_ca.is_empty() {
168            ssl_opts = ssl_opts.with_root_certs(vec![PathBuf::from(args.tls_ca.clone()).into()]);
169        }
170        if !args.tls_client_cert.is_empty() && !args.tls_client_key.is_empty() {
171            let identity = mysql_async::ClientIdentity::new(
172                PathBuf::from(args.tls_client_cert.clone()).into(),
173                PathBuf::from(args.tls_client_key.clone()).into(),
174            );
175            ssl_opts = ssl_opts.with_client_identity(Some(identity));
176        }
177        builder = builder.ssl_opts(Some(ssl_opts));
178    }
179
180    let pool = Pool::new(Opts::from(builder));
181    // Pool is dropped at scope exit; pool.disconnect() is deliberately
182    // avoided — integration tests show it hangs indefinitely, exceeding
183    // the 8s timeout. Drops handle cleanup without blocking.
184
185    let timeout = std::time::Duration::from_secs(8);
186    tokio::time::timeout(timeout, async {
187        let mut conn = pool
188            .get_conn()
189            .await
190            .map_err(|err| map_mysql_error(err, "MySQL connectivity probe failed to acquire connection"))?;
191        conn.query_drop("SELECT 1")
192            .await
193            .map_err(|err| map_mysql_error(err, "MySQL connectivity probe failed"))?;
194        Ok::<(), crate::TargetError>(())
195    })
196    .await
197    .unwrap_or_else(|_| Err(crate::TargetError::Timeout("MySQL connectivity probe timed out".to_string())))
198}
199
200/// Probes a PostgreSQL server for connectivity and verifies the configured
201/// table is readable.
202///
203/// Used by both the admin validation flow (pre-flight before persisting a
204/// target) and `PostgresTarget::init()` (runtime startup check). The probe is
205/// strictly read-only:
206///
207/// 1. Build a deadpool pool from `args` (cheap, no actual connection yet).
208/// 2. Check out a single connection.
209/// 3. Run `SELECT 1` to confirm the credentials work.
210/// 4. Run `SELECT 1 FROM <schema>.<table> LIMIT 0` to confirm the relation
211///    exists and the user has read permission. `LIMIT 0` ensures no rows are
212///    actually returned and no DML side effects occur.
213///
214/// The whole flow is wrapped in an 8s `tokio::time::timeout` so a stuck DNS
215/// resolver or TLS handshake cannot exhaust the admin layer's outer 10s
216/// timeout.
217pub async fn check_postgres_server_available(args: &crate::target::postgres::PostgresArgs) -> Result<(), crate::TargetError> {
218    use crate::target::postgres::{build_pool, map_pg_error, map_pool_error, table_probe_sql};
219
220    args.validate()?;
221
222    let timeout = std::time::Duration::from_secs(8);
223    tokio::time::timeout(timeout, async {
224        let pool = build_pool(args)?;
225        let client = pool
226            .get()
227            .await
228            .map_err(|e| map_pool_error(e, "PostgreSQL connectivity probe failed to acquire connection"))?;
229        client
230            .execute("SELECT 1", &[])
231            .await
232            .map_err(|e| map_pg_error(&e, "PostgreSQL liveness probe failed"))?;
233        let probe_sql = table_probe_sql(&args.schema, &args.table);
234        client
235            .execute(probe_sql.as_str(), &[])
236            .await
237            .map_err(|e| map_pg_error(&e, "PostgreSQL table probe failed"))?;
238        pool.close();
239        Ok::<(), crate::TargetError>(())
240    })
241    .await
242    .unwrap_or_else(|_| Err(crate::TargetError::Timeout("PostgreSQL connectivity probe timed out".to_string())))
243}
244
245pub async fn check_kafka_broker_available(args: &crate::target::kafka::KafkaArgs) -> Result<(), crate::TargetError> {
246    use rustfs_kafka_async::error::{ConnectionError, Error as KafkaError};
247    use rustfs_kafka_async::{AsyncProducer, AsyncProducerConfig, RequiredAcks};
248    use std::time::Duration;
249
250    args.validate()?;
251
252    let map_kafka_error = |err: KafkaError, context: &str| match err {
253        KafkaError::Connection(ConnectionError::NoHostReachable) => crate::TargetError::NotConnected,
254        KafkaError::Connection(ConnectionError::Timeout(_)) => crate::TargetError::Timeout(format!("{context}: {err}")),
255        KafkaError::Connection(_) => crate::TargetError::Network(format!("{context}: {err}")),
256        KafkaError::Config(_) => crate::TargetError::Configuration(format!("{context}: {err}")),
257        _ => crate::TargetError::Request(format!("{context}: {err}")),
258    };
259
260    let acks = match args.acks {
261        0 => RequiredAcks::None,
262        1 => RequiredAcks::One,
263        _ => RequiredAcks::All,
264    };
265
266    let mut config = AsyncProducerConfig::new()
267        .with_ack_timeout(Duration::from_secs(5))
268        .with_required_acks(acks);
269
270    if let Some(security) = args.security_config(false)? {
271        config = config.with_security(security);
272    }
273
274    tokio::time::timeout(Duration::from_secs(5), async {
275        let _ = AsyncProducer::from_hosts_with_config(args.brokers.clone(), config)
276            .await
277            .map_err(|err| map_kafka_error(err, "Kafka broker check failed to create producer"))?;
278        Ok(())
279    })
280    .await
281    .unwrap_or_else(|_| Err(crate::TargetError::Timeout("Kafka connection timed out".to_string())))
282}
283
284pub async fn check_redis_server_available(args: &crate::target::redis::RedisArgs) -> Result<(), crate::TargetError> {
285    tokio::time::timeout(std::time::Duration::from_secs(5), async {
286        let client = crate::target::redis::build_redis_client(args)?;
287        crate::target::redis::ping_redis_server(&client, args).await
288    })
289    .await
290    .unwrap_or_else(|_| Err(crate::TargetError::Timeout("Redis connection timed out".to_string())))
291}
292
293pub async fn check_amqp_broker_available(args: &crate::target::amqp::AMQPArgs) -> Result<(), crate::TargetError> {
294    match tokio::time::timeout(std::time::Duration::from_secs(5), async {
295        let connection = crate::target::amqp::connect_amqp(args).await?;
296        if !connection.connection.status().connected() || !connection.channel.status().connected() {
297            return Err(crate::TargetError::NotConnected);
298        }
299        connection
300            .connection
301            .close(200, "OK".into())
302            .await
303            .map_err(|e| crate::TargetError::Network(format!("Failed to close AMQP check connection: {e}")))?;
304        Ok(())
305    })
306    .await
307    {
308        Ok(result) => result,
309        Err(_) => Err(crate::TargetError::Timeout("AMQP connection timed out".to_string())),
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316    use crate::{
317        TargetError,
318        target::{TargetType, kafka::KafkaArgs, mysql::MySqlArgs},
319    };
320
321    fn kafka_args() -> KafkaArgs {
322        KafkaArgs {
323            enable: true,
324            brokers: vec!["127.0.0.1:9092".to_string()],
325            topic: "rustfs-events".to_string(),
326            acks: 1,
327            tls_enable: false,
328            tls_ca: String::new(),
329            tls_client_cert: String::new(),
330            tls_client_key: String::new(),
331            sasl_enable: false,
332            sasl_mechanism: String::new(),
333            sasl_username: String::new(),
334            sasl_password: String::new(),
335            queue_dir: String::new(),
336            queue_limit: 100,
337            target_type: TargetType::NotifyEvent,
338        }
339    }
340
341    fn mysql_args() -> MySqlArgs {
342        MySqlArgs {
343            enable: true,
344            dsn_string: "rustfs:password@tcp(127.0.0.1:3306)/rustfs_events".to_string(),
345            table: "rustfs_events".to_string(),
346            format: "access".to_string(),
347            tls_ca: String::new(),
348            tls_client_cert: String::new(),
349            tls_client_key: String::new(),
350            queue_dir: String::new(),
351            queue_limit: 100,
352            max_open_connections: 2,
353            target_type: TargetType::NotifyEvent,
354        }
355    }
356
357    #[test]
358    fn check_kafka_broker_available_rejects_sasl_without_tls_before_connecting() {
359        let mut args = kafka_args();
360        args.sasl_enable = true;
361        args.sasl_username = "user".to_string();
362        args.sasl_password = "secret".to_string();
363
364        let err = tokio::runtime::Runtime::new()
365            .expect("runtime")
366            .block_on(check_kafka_broker_available(&args))
367            .expect_err("SASL without TLS should fail before opening a network connection");
368
369        match err {
370            TargetError::Configuration(msg) => assert!(msg.contains("requires tls_enable")),
371            other => panic!("expected configuration error, got {other:?}"),
372        }
373    }
374
375    #[test]
376    fn check_mysql_server_available_rejects_invalid_table_before_connecting() {
377        let mut args = mysql_args();
378        args.table = "rustfs-events".to_string();
379
380        let err = tokio::runtime::Runtime::new()
381            .expect("runtime")
382            .block_on(check_mysql_server_available(&args))
383            .expect_err("invalid table should fail before opening a network connection");
384
385        match err {
386            TargetError::Configuration(msg) => assert!(msg.contains("not a valid identifier")),
387            other => panic!("expected configuration error, got {other:?}"),
388        }
389    }
390
391    #[test]
392    fn check_mysql_server_available_rejects_unpaired_tls_client_fields_before_connecting() {
393        let mut args = mysql_args();
394        args.dsn_string = "rustfs:password@tcp(127.0.0.1:3306)/rustfs_events?tls=true".to_string();
395        args.tls_client_cert = "/etc/ssl/mysql/client.pem".to_string();
396
397        let err = tokio::runtime::Runtime::new()
398            .expect("runtime")
399            .block_on(check_mysql_server_available(&args))
400            .expect_err("unpaired TLS client fields should fail before opening a network connection");
401
402        match err {
403            TargetError::Configuration(msg) => assert!(msg.contains("must be specified together")),
404            other => panic!("expected configuration error, got {other:?}"),
405        }
406    }
407}
rustfs_targets/check.rs

rustfs_targets/
check.rs
