risingwave_connector/connector_common/
mqtt_common.rs

1// Copyright 2024 RisingWave Labs
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
15use phf::{Set, phf_set};
16use risingwave_common::id::ActorId;
17use rumqttc::tokio_rustls::rustls;
18use rumqttc::v5::mqttbytes::QoS;
19use rumqttc::v5::mqttbytes::v5::ConnectProperties;
20use rumqttc::v5::{AsyncClient, EventLoop, MqttOptions};
21use serde::Deserialize;
22use serde_with::{DisplayFromStr, serde_as};
23use strum_macros::{Display, EnumString};
24use with_options::WithOptions;
25
26use super::common::{load_certs, load_private_key};
27use crate::deserialize_bool_from_string;
28use crate::enforce_secret::EnforceSecret;
29use crate::error::ConnectorResult;
30
31const DEFAULT_KEEP_ALIVE_SECS: u16 = 10;
32/// `rumqttc` panics if the keep alive interval is less than 5 seconds.
33const MIN_KEEP_ALIVE_SECS: u16 = 5;
34
35/// Parses `keep_alive_secs` and rejects values below [`MIN_KEEP_ALIVE_SECS`], so that an invalid
36/// value fails when the source or sink is created instead of panicking in `rumqttc` on connect.
37fn deserialize_keep_alive_secs<'de, D>(deserializer: D) -> Result<Option<u16>, D::Error>
38where
39    D: serde::Deserializer<'de>,
40{
41    let s: String = Deserialize::deserialize(deserializer)?;
42    let secs: u16 = s.trim().parse().map_err(|_| {
43        serde::de::Error::invalid_value(
44            serde::de::Unexpected::Str(&s),
45            &"an integer between 5 and 65535",
46        )
47    })?;
48    if secs < MIN_KEEP_ALIVE_SECS {
49        return Err(serde::de::Error::invalid_value(
50            serde::de::Unexpected::Unsigned(secs.into()),
51            &"an integer between 5 and 65535",
52        ));
53    }
54    Ok(Some(secs))
55}
56
57#[derive(Debug, Clone, PartialEq, Display, Deserialize, EnumString)]
58#[strum(serialize_all = "snake_case")]
59#[expect(clippy::enum_variant_names)]
60pub enum QualityOfService {
61    AtLeastOnce,
62    AtMostOnce,
63    ExactlyOnce,
64}
65
66#[serde_as]
67#[derive(Deserialize, Debug, Clone, WithOptions)]
68pub struct MqttCommon {
69    /// The url of the broker to connect to. e.g. <tcp://localhost>.
70    /// Must be prefixed with one of either `tcp://`, `mqtt://`, `ssl://`,`mqtts://`,
71    /// to denote the protocol for establishing a connection with the broker.
72    /// `mqtts://`, `ssl://` will use the native certificates if no ca is specified
73    pub url: String,
74
75    /// The quality of service to use when publishing messages. Defaults to `at_most_once`.
76    /// Could be `at_most_once`, `at_least_once` or `exactly_once`
77    #[serde_as(as = "Option<DisplayFromStr>")]
78    pub qos: Option<QualityOfService>,
79
80    /// Username for the mqtt broker
81    #[serde(rename = "username")]
82    pub user: Option<String>,
83
84    /// Password for the mqtt broker
85    pub password: Option<String>,
86
87    /// Prefix for the mqtt client id.
88    /// The client id will be generated as `client_prefix`_`actor_id`_`(executor_id|source_id)`. Defaults to risingwave
89    pub client_prefix: Option<String>,
90
91    /// Whether to start a new session when connecting (MQTT v5 Clean Start). Defaults to `false`.
92    ///
93    /// `clean_start = true` makes the broker discard any existing session for the client id when
94    /// connecting. `clean_start = false` resumes the session the broker still holds for the same
95    /// client id, if any, including its subscriptions and any messages queued for it.
96    ///
97    /// This flag only applies when connecting. How long the broker keeps a session after a
98    /// disconnect is controlled by `session_expiry_interval`. If that is not set, the session ends
99    /// when the connection closes, so `clean_start = false` alone does not resume anything after a
100    /// reconnect. Conversely, `clean_start = true` with `session_expiry_interval = 60` still keeps
101    /// the new session for 60 seconds after a disconnect.
102    #[serde(default, deserialize_with = "deserialize_bool_from_string")]
103    pub clean_start: bool,
104
105    /// The keep alive interval in seconds. The client sends a ping to the broker if no other
106    /// packet is exchanged within this interval, and the broker may close the connection after
107    /// 1.5 times this interval without receiving any packet from the client. Must be between 5
108    /// and 65535. Defaults to 10.
109    ///
110    /// Consider a larger value if the source can be backpressured for a long time, since no
111    /// ping is sent while the source is not consuming messages.
112    #[serde(default, deserialize_with = "deserialize_keep_alive_secs")]
113    pub keep_alive_secs: Option<u16>,
114
115    /// The session expiry interval in seconds sent to the broker on connect (MQTT v5).
116    /// If not set, the broker treats it as 0 and discards the session as soon as the connection
117    /// is closed.
118    ///
119    /// Set it together with `clean_start = false` to let the broker keep the session, including
120    /// subscriptions and queued messages, across reconnects with the same client id. Note that
121    /// the client id contains the actor id, so a session is not resumed after the job is
122    /// rescheduled.
123    #[serde_as(as = "Option<DisplayFromStr>")]
124    pub session_expiry_interval: Option<u32>,
125
126    /// The maximum number of inflight messages. Defaults to 100
127    #[serde_as(as = "Option<DisplayFromStr>")]
128    pub inflight_messages: Option<usize>,
129
130    /// The max size of messages received by the MQTT client
131    #[serde_as(as = "Option<DisplayFromStr>")]
132    pub max_packet_size: Option<u32>,
133
134    /// Path to CA certificate file for verifying the broker's key.
135    #[serde(rename = "tls.ca")]
136    pub ca: Option<String>,
137    /// Path to client's certificate file (PEM). Required for client authentication.
138    /// Can be a file path under fs:// or a string with the certificate content.
139    #[serde(rename = "tls.client_cert")]
140    pub client_cert: Option<String>,
141
142    /// Path to client's private key file (PEM). Required for client authentication.
143    /// Can be a file path under fs:// or a string with the private key content.
144    #[serde(rename = "tls.client_key")]
145    pub client_key: Option<String>,
146}
147
148impl EnforceSecret for MqttCommon {
149    const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
150        "tls.client_cert",
151        "tls.client_key",
152        "password",
153    };
154}
155
156impl MqttCommon {
157    pub(crate) fn build_client(
158        &self,
159        actor_id: ActorId,
160        id: u32,
161    ) -> ConnectorResult<(AsyncClient, EventLoop)> {
162        let client_id = format!(
163            "{}_{}_{}",
164            self.client_prefix.as_deref().unwrap_or("risingwave"),
165            actor_id,
166            id
167        );
168
169        let mut url = url::Url::parse(&self.url)?;
170
171        let ssl = matches!(url.scheme(), "mqtts" | "ssl");
172
173        url.query_pairs_mut().append_pair("client_id", &client_id);
174
175        tracing::debug!("connecting mqtt using url: {}", url.as_str());
176
177        let mut options = MqttOptions::try_from(url)?;
178        options.set_keep_alive(std::time::Duration::from_secs(
179            self.keep_alive_secs
180                .unwrap_or(DEFAULT_KEEP_ALIVE_SECS)
181                .into(),
182        ));
183
184        options.set_clean_start(self.clean_start);
185
186        let mut connect_properties = ConnectProperties::new();
187        connect_properties.max_packet_size = self.max_packet_size;
188        connect_properties.session_expiry_interval = self.session_expiry_interval;
189        options.set_connect_properties(connect_properties);
190
191        if ssl {
192            let tls_config = self.get_tls_config()?;
193            options.set_transport(rumqttc::Transport::tls_with_config(
194                rumqttc::TlsConfiguration::Rustls(std::sync::Arc::new(tls_config)),
195            ));
196        }
197
198        if let Some(user) = &self.user {
199            options.set_credentials(user, self.password.as_deref().unwrap_or_default());
200        }
201
202        Ok(rumqttc::v5::AsyncClient::new(
203            options,
204            self.inflight_messages.unwrap_or(100),
205        ))
206    }
207
208    pub(crate) fn qos(&self) -> QoS {
209        self.qos
210            .as_ref()
211            .map(|qos| match qos {
212                QualityOfService::AtMostOnce => QoS::AtMostOnce,
213                QualityOfService::AtLeastOnce => QoS::AtLeastOnce,
214                QualityOfService::ExactlyOnce => QoS::ExactlyOnce,
215            })
216            .unwrap_or(QoS::AtMostOnce)
217    }
218
219    fn get_tls_config(&self) -> ConnectorResult<rustls::ClientConfig> {
220        let mut root_cert_store = rustls::RootCertStore::empty();
221        if let Some(ca) = &self.ca {
222            let certificates = load_certs(ca)?;
223            for cert in certificates {
224                root_cert_store.add(cert).unwrap();
225            }
226        } else {
227            for cert in
228                rustls_native_certs::load_native_certs().expect("could not load platform certs")
229            {
230                root_cert_store.add(cert).unwrap();
231            }
232        }
233
234        let builder = rustls::ClientConfig::builder().with_root_certificates(root_cert_store);
235
236        let tls_config = if let (Some(client_cert), Some(client_key)) =
237            (self.client_cert.as_ref(), self.client_key.as_ref())
238        {
239            let certs = load_certs(client_cert)?;
240            let key = load_private_key(client_key)?;
241
242            builder.with_client_auth_cert(certs, key)?
243        } else {
244            builder.with_no_client_auth()
245        };
246
247        Ok(tls_config)
248    }
249}
250
251#[cfg(test)]
252mod tests {
253    use std::collections::BTreeMap;
254    use std::time::Duration;
255
256    use super::*;
257
258    fn parse(extra: &[(&str, &str)]) -> Result<MqttCommon, serde_json::Error> {
259        let mut props: BTreeMap<String, String> = BTreeMap::new();
260        props.insert("url".to_owned(), "tcp://localhost:1883".to_owned());
261        for (k, v) in extra {
262            props.insert((*k).to_owned(), (*v).to_owned());
263        }
264        serde_json::from_value(serde_json::to_value(props).unwrap())
265    }
266
267    #[test]
268    fn test_keep_alive_and_session_expiry_default() {
269        let common = parse(&[]).unwrap();
270        assert_eq!(common.keep_alive_secs, None);
271        assert_eq!(common.session_expiry_interval, None);
272
273        let (_client, eventloop) = common.build_client(1.into(), 2).unwrap();
274        assert_eq!(eventloop.options.keep_alive(), Duration::from_secs(10));
275        assert_eq!(eventloop.options.session_expiry_interval(), None);
276    }
277
278    #[test]
279    fn test_keep_alive_and_session_expiry_configured() {
280        let common = parse(&[
281            ("keep_alive_secs", "300"),
282            ("session_expiry_interval", "3600"),
283            ("clean_start", "false"),
284            ("max_packet_size", "1024"),
285        ])
286        .unwrap();
287        assert_eq!(common.keep_alive_secs, Some(300));
288        assert_eq!(common.session_expiry_interval, Some(3600));
289
290        let (_client, eventloop) = common.build_client(1.into(), 2).unwrap();
291        assert_eq!(eventloop.options.keep_alive(), Duration::from_secs(300));
292        assert_eq!(eventloop.options.session_expiry_interval(), Some(3600));
293        assert!(!eventloop.options.clean_start());
294        assert_eq!(
295            eventloop
296                .options
297                .connect_properties()
298                .and_then(|p| p.max_packet_size),
299            Some(1024)
300        );
301    }
302
303    #[test]
304    fn test_invalid_keep_alive_secs() {
305        for v in ["0", "4", "65536", "-1", "abc"] {
306            assert!(
307                parse(&[("keep_alive_secs", v)]).is_err(),
308                "keep_alive_secs = {v} should be rejected"
309            );
310        }
311        assert_eq!(
312            parse(&[("keep_alive_secs", "5")]).unwrap().keep_alive_secs,
313            Some(5)
314        );
315        assert_eq!(
316            parse(&[("keep_alive_secs", "65535")])
317                .unwrap()
318                .keep_alive_secs,
319            Some(65535)
320        );
321    }
322}
risingwave_connector/connector_common/mqtt_common.rs

risingwave_connector/connector_common/
mqtt_common.rs
