mz_mysql_util/
tunnel.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use aws_types::SdkConfig;
11use mysql_async::{Conn, Opts, OptsBuilder};
12use std::collections::BTreeSet;
13use std::net::IpAddr;
14use std::ops::{Deref, DerefMut};
15use std::time::Duration;
16
17use mz_ore::future::{InTask, TimeoutError};
18use mz_ore::option::OptionExt;
19use mz_ore::task::spawn;
20use mz_repr::CatalogItemId;
21use mz_ssh_util::tunnel::{SshTimeoutConfig, SshTunnelConfig};
22use mz_ssh_util::tunnel_manager::{ManagedSshTunnelHandle, SshTunnelManager};
23use serde::{Deserialize, Serialize};
24use tracing::{error, info, warn};
25
26use crate::MySqlError;
27use crate::aws_rds::rds_auth_token;
28
29/// Configures an optional tunnel for use when connecting to a MySQL
30/// database.
31#[derive(Debug, PartialEq, Clone)]
32pub enum TunnelConfig {
33    /// Establish a direct TCP connection to the database host.
34    /// If `resolved_ips` is not None, the provided IPs will be used
35    /// rather than resolving the hostname.
36    Direct {
37        resolved_ips: Option<BTreeSet<IpAddr>>,
38    },
39    /// Establish a TCP connection to the database via an SSH tunnel.
40    /// This means first establishing an SSH connection to a bastion host,
41    /// and then opening a separate connection from that host to the database.
42    /// This is commonly referred by vendors as a "direct SSH tunnel", in
43    /// opposition to "reverse SSH tunnel", which is currently unsupported.
44    Ssh { config: SshTunnelConfig },
45    /// Establish a TCP connection to the database via an AWS PrivateLink
46    /// service.
47    AwsPrivatelink {
48        /// The ID of the AWS PrivateLink service.
49        connection_id: CatalogItemId,
50    },
51}
52
53pub const DEFAULT_TCP_KEEPALIVE: Duration = Duration::from_secs(60);
54pub const DEFAULT_SNAPSHOT_MAX_EXECUTION_TIME: Duration = Duration::ZERO;
55pub const DEFAULT_SNAPSHOT_LOCK_WAIT_TIMEOUT: Duration = Duration::from_secs(3600);
56/// The `wait_timeout` to set on connections used during snapshotting, chosen
57/// to comfortably outlast even very long-running snapshots.
58pub const DEFAULT_SNAPSHOT_WAIT_TIMEOUT: Duration = Duration::from_secs(48 * 60 * 60);
59/// The minimum value MySQL accepts for `wait_timeout`.
60pub const MIN_SNAPSHOT_WAIT_TIMEOUT: Duration = Duration::from_secs(1);
61/// The maximum value MySQL accepts for `wait_timeout` on Windows, which is
62/// lower than the Unix maximum and so the portable upper bound.
63pub const MAX_SNAPSHOT_WAIT_TIMEOUT: Duration = Duration::from_secs(2147483);
64pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(60);
65
66#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
67pub struct TimeoutConfig {
68    // Snapshot-related configs
69    pub snapshot_max_execution_time: Option<Duration>,
70    pub snapshot_lock_wait_timeout: Option<Duration>,
71    pub snapshot_wait_timeout: Option<Duration>,
72
73    // Socket-related configs
74    pub tcp_keepalive: Option<Duration>,
75
76    // Connection timeout.  This timeout covers creating an authenticated connection
77    // (e.g. includes network connection, TLS handshake, authentication, etc.).
78    // If the connection has not been established in that time, it is considered an error.
79    pub connect_timeout: Option<Duration>,
80    // There are other timeout options on `mysql_async::OptsBuilder`
81    // (e.g. `conn_ttl` and `wait_timeout`) that could be exposed
82    // but they only apply to connection pools, which we are not currently using.
83}
84
85impl Default for TimeoutConfig {
86    fn default() -> Self {
87        Self {
88            snapshot_max_execution_time: Some(DEFAULT_SNAPSHOT_MAX_EXECUTION_TIME),
89            snapshot_lock_wait_timeout: Some(DEFAULT_SNAPSHOT_LOCK_WAIT_TIMEOUT),
90            snapshot_wait_timeout: Some(DEFAULT_SNAPSHOT_WAIT_TIMEOUT),
91            tcp_keepalive: Some(DEFAULT_TCP_KEEPALIVE),
92            connect_timeout: Some(DEFAULT_CONNECT_TIMEOUT),
93        }
94    }
95}
96
97impl TimeoutConfig {
98    pub fn build(
99        snapshot_max_execution_time: Duration,
100        snapshot_lock_wait_timeout: Duration,
101        snapshot_wait_timeout: Duration,
102        tcp_keepalive: Duration,
103        connect_timeout: Duration,
104    ) -> Self {
105        // Verify values are within valid ranges
106        // Note we error log but do not fail as this is called in a non-fallible
107        // LD-sync in the adapter.
108
109        // https://dev.mysql.com/doc/refman/8.0/en/server-system-variables.html#sysvar_lock_wait_timeout
110        let snapshot_lock_wait_timeout = if snapshot_lock_wait_timeout.as_secs() > 31536000 {
111            error!(
112                "snapshot_lock_wait_timeout is too large: {}. Maximum is 31536000.",
113                snapshot_lock_wait_timeout.as_secs()
114            );
115            Some(DEFAULT_SNAPSHOT_LOCK_WAIT_TIMEOUT)
116        } else {
117            Some(snapshot_lock_wait_timeout)
118        };
119
120        // https://dev.mysql.com/doc/refman/8.0/en/server-system-variables.html#sysvar_wait_timeout
121        let snapshot_wait_timeout = if snapshot_wait_timeout > MAX_SNAPSHOT_WAIT_TIMEOUT
122            || snapshot_wait_timeout < MIN_SNAPSHOT_WAIT_TIMEOUT
123        {
124            error!(
125                "snapshot_wait_timeout is out of range: {}. Must be within [{}, {}].",
126                snapshot_wait_timeout.as_secs(),
127                MIN_SNAPSHOT_WAIT_TIMEOUT.as_secs(),
128                MAX_SNAPSHOT_WAIT_TIMEOUT.as_secs(),
129            );
130            Some(DEFAULT_SNAPSHOT_WAIT_TIMEOUT)
131        } else {
132            Some(snapshot_wait_timeout)
133        };
134
135        // https://dev.mysql.com/doc/refman/8.0/en/server-system-variables.html#sysvar_max_execution_time
136        let snapshot_max_execution_time = if snapshot_max_execution_time.as_millis() > 4294967295 {
137            error!(
138                "snapshot_max_execution_time is too large: {}. Maximum is 4294967295.",
139                snapshot_max_execution_time.as_secs()
140            );
141            Some(DEFAULT_SNAPSHOT_MAX_EXECUTION_TIME)
142        } else {
143            Some(snapshot_max_execution_time)
144        };
145
146        let tcp_keepalive = match u32::try_from(tcp_keepalive.as_millis()) {
147            Err(_) => {
148                error!(
149                    "tcp_keepalive is too large: {}. Maximum is {}.",
150                    tcp_keepalive.as_millis(),
151                    u32::MAX,
152                );
153                Some(DEFAULT_TCP_KEEPALIVE)
154            }
155            Ok(_) => Some(tcp_keepalive),
156        };
157
158        let connect_timeout = match u32::try_from(connect_timeout.as_millis()) {
159            Err(_) => {
160                error!(
161                    "connect_timeout is too large: {}. Maximum is {}.",
162                    connect_timeout.as_millis(),
163                    u32::MAX,
164                );
165                Some(DEFAULT_CONNECT_TIMEOUT)
166            }
167            Ok(_) => Some(connect_timeout),
168        };
169
170        Self {
171            snapshot_max_execution_time,
172            snapshot_lock_wait_timeout,
173            snapshot_wait_timeout,
174            tcp_keepalive,
175            connect_timeout,
176        }
177    }
178
179    /// Apply relevant timeout configurations to a `mysql_async::OptsBuilder`.
180    pub fn apply_to_opts(&self, mut opts_builder: OptsBuilder) -> Result<OptsBuilder, MySqlError> {
181        if let Some(tcp_keepalive) = self.tcp_keepalive {
182            opts_builder = opts_builder.tcp_keepalive(Some(
183                u32::try_from(tcp_keepalive.as_millis()).map_err(|e| {
184                    MySqlError::InvalidClientConfig(format!(
185                        "invalid tcp_keepalive duration: {}",
186                        e
187                    ))
188                })?,
189            ));
190        }
191        Ok(opts_builder)
192    }
193}
194
195/// A MySQL connection with an optional SSH tunnel handle.
196///
197/// This wrapper is intended to be used in place of `mysql_async::Conn` to
198/// keep the SSH tunnel alive for the lifecycle of the connection by holding
199/// a reference to the tunnel handle.
200#[derive(Debug)]
201pub struct MySqlConn {
202    conn: Conn,
203    _ssh_tunnel_handle: Option<ManagedSshTunnelHandle>,
204}
205
206impl Deref for MySqlConn {
207    type Target = Conn;
208
209    fn deref(&self) -> &Self::Target {
210        &self.conn
211    }
212}
213
214impl DerefMut for MySqlConn {
215    fn deref_mut(&mut self) -> &mut Self::Target {
216        &mut self.conn
217    }
218}
219
220impl MySqlConn {
221    pub async fn disconnect(mut self) -> Result<(), MySqlError> {
222        self.conn.disconnect().await?;
223        self._ssh_tunnel_handle.take();
224        Ok(())
225    }
226
227    pub fn take(self) -> (Conn, Option<ManagedSshTunnelHandle>) {
228        (self.conn, self._ssh_tunnel_handle)
229    }
230}
231
232/// Configuration for MySQL connections.
233///
234/// This wraps [`mysql_async::Opts`] to allow the configuration of a
235/// tunnel via a [`TunnelConfig`].
236#[derive(Clone, Debug)]
237pub struct Config {
238    inner: Opts,
239    tunnel: TunnelConfig,
240    // Whether to poll I/O for this connection in a tokio task
241    // TODO(roshan): Make this apply to queries on the returned connection, not just the initial
242    // connection.
243    in_task: InTask,
244    ssh_timeout_config: SshTimeoutConfig,
245    mysql_timeout_config: TimeoutConfig,
246    aws_config: Option<SdkConfig>,
247}
248
249impl Config {
250    pub fn new(
251        builder: OptsBuilder,
252        tunnel: TunnelConfig,
253        ssh_timeout_config: SshTimeoutConfig,
254        in_task: InTask,
255        mysql_timeout_config: TimeoutConfig,
256        aws_config: Option<SdkConfig>,
257    ) -> Result<Self, MySqlError> {
258        let opts = mysql_timeout_config.apply_to_opts(builder)?;
259        Ok(Self {
260            inner: opts.into(),
261            tunnel,
262            in_task,
263            ssh_timeout_config,
264            mysql_timeout_config,
265            aws_config,
266        })
267    }
268
269    pub async fn connect(
270        &self,
271        task_name: &str,
272        ssh_tunnel_manager: &SshTunnelManager,
273    ) -> Result<MySqlConn, MySqlError> {
274        let address = format!(
275            "mysql:://{}@{}:{}/{}",
276            self.inner.user().display_or("<unknown-user>"),
277            self.inner.ip_or_hostname(),
278            self.inner.tcp_port(),
279            self.inner.db_name().display_or("<unknown-dbname>"),
280        );
281        info!(%task_name, %address, "connecting");
282        match self.connect_internal(ssh_tunnel_manager).await {
283            Ok(t) => {
284                info!(%task_name, %address, "connected");
285                Ok(t)
286            }
287            Err(e) => {
288                warn!(%task_name, %address, "connection failed: {e:#}");
289                Err(e)
290            }
291        }
292    }
293
294    fn address(&self) -> (&str, u16) {
295        (self.inner.ip_or_hostname(), self.inner.tcp_port())
296    }
297
298    async fn connect_internal(
299        &self,
300        ssh_tunnel_manager: &SshTunnelManager,
301    ) -> Result<MySqlConn, MySqlError> {
302        let mut opts_builder = OptsBuilder::from_opts(self.inner.clone());
303
304        if let Some(aws_config) = &self.aws_config {
305            let (host, port) = self.address();
306            let username = self.inner.user().expect("MySQL: username required");
307
308            let token = rds_auth_token(host, port, username, aws_config).await?;
309            // Cleartext plugin must be enabled for IAM authentication, for security,
310            // the network traffic is SSL/TLS encrypted.  The cleartext plugin is built
311            // into the MySQL client library.
312            opts_builder = opts_builder
313                .pass(Some(token.to_string()))
314                .enable_cleartext_plugin(true);
315        }
316
317        match &self.tunnel {
318            TunnelConfig::Direct { resolved_ips } => {
319                opts_builder = opts_builder.resolved_ips(
320                    resolved_ips
321                        .clone()
322                        .map(|ips| ips.into_iter().collect::<Vec<_>>()),
323                );
324
325                Ok(MySqlConn {
326                    conn: self.connect_with_timeout(opts_builder).await?,
327                    _ssh_tunnel_handle: None,
328                })
329            }
330            TunnelConfig::Ssh { config } => {
331                let (host, port) = self.address();
332                let tunnel = ssh_tunnel_manager
333                    .connect(
334                        config.clone(),
335                        host,
336                        port,
337                        self.ssh_timeout_config,
338                        self.in_task,
339                    )
340                    .await
341                    .map_err(MySqlError::Ssh)?;
342
343                let tunnel_addr = tunnel.local_addr();
344                // Override the connection host and port for the actual TCP connection to point to
345                // the local tunnel instead.
346                opts_builder = opts_builder
347                    .ip_or_hostname(tunnel_addr.ip().to_string())
348                    .tcp_port(tunnel_addr.port());
349
350                if let Some(ssl_opts) = self.inner.ssl_opts() {
351                    if !ssl_opts.skip_domain_validation() {
352                        // If the TLS configuration will validate the hostname, we need to set
353                        // the TLS hostname back to the actual upstream host and not the hostname
354                        // of the local SSH tunnel
355                        opts_builder = opts_builder.ssl_opts(Some(
356                            ssl_opts.clone().with_danger_tls_hostname_override(Some(
357                                self.inner.ip_or_hostname().to_string(),
358                            )),
359                        ));
360                    }
361                }
362
363                Ok(MySqlConn {
364                    conn: self.connect_with_timeout(opts_builder).await?,
365                    _ssh_tunnel_handle: Some(tunnel),
366                })
367            }
368            TunnelConfig::AwsPrivatelink { connection_id } => {
369                let privatelink_host = mz_cloud_resources::vpc_endpoint_name(*connection_id);
370
371                // Override the connection host for the actual TCP connection to point to
372                // the privatelink hostname instead.
373                let mut opts_builder = opts_builder.ip_or_hostname(privatelink_host);
374
375                if let Some(ssl_opts) = self.inner.ssl_opts() {
376                    if !ssl_opts.skip_domain_validation() {
377                        // If the TLS configuration will validate the hostname, we need to set
378                        // the TLS hostname back to the actual upstream host and not the
379                        // privatelink hostname.
380                        opts_builder = opts_builder.ssl_opts(Some(
381                            ssl_opts.clone().with_danger_tls_hostname_override(Some(
382                                self.inner.ip_or_hostname().to_string(),
383                            )),
384                        ));
385                    }
386                }
387
388                Ok(MySqlConn {
389                    conn: self.connect_with_timeout(opts_builder).await?,
390                    _ssh_tunnel_handle: None,
391                })
392            }
393        }
394    }
395
396    async fn connect_with_timeout(
397        &self,
398        opts_builder: OptsBuilder,
399    ) -> Result<mysql_async::Conn, MySqlError> {
400        let connection_future = if let InTask::Yes = self.in_task {
401            Box::pin(spawn(|| "mysql_connect".to_string(), Conn::new(opts_builder)).abort_on_drop())
402        } else {
403            Conn::new(opts_builder)
404        };
405
406        if let Some(connect_timeout) = self.mysql_timeout_config.connect_timeout {
407            mz_ore::future::timeout(connect_timeout, connection_future)
408                .await
409                .map_err(|err| match err {
410                    // match instead of impl From<> for MySqlError so we can capture the timeout value
411                    TimeoutError::DeadlineElapsed => MySqlError::ConnectionTimeout(connect_timeout),
412                    TimeoutError::Inner(e) => MySqlError::from(e),
413                })
414        } else {
415            connection_future.await.map_err(MySqlError::from)
416        }
417    }
418}
mz_mysql_util/tunnel.rs

mz_mysql_util/
tunnel.rs
