diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index bc29f24407..e85395c738 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -2722,6 +2722,26 @@ impl AppState { return Ok(mqc); } + if mqc.system_kind == crate::mq::types::MqSystemKind::Kafka { + if let Some((bootstrap_host, bootstrap_port)) = kafka_single_loopback_bootstrap_endpoint(&mqc.extra) { + let transport_layers = self.resolved_transport_layers(config).await?; + if matches!(transport_layers.last(), Some(TransportLayerConfig::Ssh(_))) { + let local_port = db::transport_layer_tunnel::start_transport_layers_with_final_ssh_local_port( + connection_id, + &transport_layers, + &bootstrap_host, + bootstrap_port, + Some(bootstrap_port), + &self.tunnels, + &self.proxy_tunnels, + &self.http_tunnels, + ) + .await?; + return Ok(mqc.with_connect_override("127.0.0.1", local_port)); + } + } + } + let (host, port) = self.connection_host_port(connection_id, config).await?; Ok(mqc.with_connect_override(&host, port)) } @@ -4396,6 +4416,30 @@ fn parse_mq_admin_host_port(config: &ConnectionConfig) -> Option<(String, u16)> Some((host, port)) } +#[cfg(any(feature = "mq-admin", test))] +fn kafka_single_loopback_bootstrap_endpoint(extra: &serde_json::Value) -> Option<(String, u16)> { + let value = extra.get("bootstrapServers").or_else(|| extra.get("bootstrap_servers"))?.as_str()?.trim(); + let mut endpoints = value + .split(|ch: char| ch.is_whitespace() || matches!(ch, ',' | ';' | '，' | '；')) + .filter(|endpoint| !endpoint.is_empty()); + let endpoint = endpoints.next()?; + if endpoints.next().is_some() { + return None; + } + let address = endpoint.rsplit_once("://").map_or(endpoint, |(_, address)| address); + let url = reqwest::Url::parse(&format!("kafka://{address}")).ok()?; + if url.host_str()? != "127.0.0.1" + || url.username() != "" + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || !matches!(url.path(), "" | "/") + { + return None; + } + Some(("127.0.0.1".to_string(), url.port()?)) +} + fn parse_nacos_server_host_port(config: &ConnectionConfig) -> Option<(String, u16)> { let value = config .external_config @@ -4961,12 +5005,12 @@ mod tests { agent_connect_timeout, connection_probe_endpoints, connection_remote_endpoint, connection_url_for_endpoint, database_connection_config, database_connection_config_with_catalog, gaussdb_identifier_quote_from_query_result, gaussdb_m_jdbc_config_for_endpoint, gaussdb_uses_m_jdbc_driver, - metadata_connection_config, mysql_metadata_fallback_url, mysql_pool_setup_queries, - oceanbase_mysql_query_timeout_sql, oceanbase_mysql_setup_queries, prestosql_jdbc_config_for_endpoint, - redacted_connection_url_for_endpoint, redis_sentinel_transport_id, redis_sentinel_transport_prefix, - sqlserver_legacy_agent_config, sqlserver_legacy_driver_error, sqlserver_uses_legacy_driver, - task_client_session_id, upsert_connection_url_param, uses_bare_mysql_pool, uses_tcp_probe, - validate_connection_url_params, validate_h2_database_path, AppState, MysqlMode, PoolKind, + kafka_single_loopback_bootstrap_endpoint, metadata_connection_config, mysql_metadata_fallback_url, + mysql_pool_setup_queries, oceanbase_mysql_query_timeout_sql, oceanbase_mysql_setup_queries, + prestosql_jdbc_config_for_endpoint, redacted_connection_url_for_endpoint, redis_sentinel_transport_id, + redis_sentinel_transport_prefix, sqlserver_legacy_agent_config, sqlserver_legacy_driver_error, + sqlserver_uses_legacy_driver, task_client_session_id, upsert_connection_url_param, uses_bare_mysql_pool, + uses_tcp_probe, validate_connection_url_params, validate_h2_database_path, AppState, MysqlMode, PoolKind, GAUSSDB_M_JDBC_DRIVER_CLASS, GAUSSDB_M_JDBC_DRIVER_PROFILE, PRESTOSQL_JDBC_DRIVER_CLASS, }; use crate::agent_connection::{ @@ -7633,6 +7677,34 @@ for line in sys.stdin: assert_eq!(connection_remote_endpoint(&config), ("broker.internal".to_string(), 8443)); } + #[test] + fn kafka_loopback_bootstrap_endpoint_requires_one_ipv4_loopback_server() { + assert_eq!( + kafka_single_loopback_bootstrap_endpoint(&serde_json::json!({ + "bootstrapServers": "127.0.0.1:9093" + })), + Some(("127.0.0.1".to_string(), 9093)) + ); + assert_eq!( + kafka_single_loopback_bootstrap_endpoint(&serde_json::json!({ + "bootstrap_servers": "PLAINTEXT://127.0.0.1:19093" + })), + Some(("127.0.0.1".to_string(), 19093)) + ); + assert_eq!( + kafka_single_loopback_bootstrap_endpoint(&serde_json::json!({ + "bootstrapServers": "127.0.0.1:9093,127.0.0.1:9094" + })), + None + ); + assert_eq!( + kafka_single_loopback_bootstrap_endpoint(&serde_json::json!({ + "bootstrapServers": "broker.internal:9093" + })), + None + ); + } + #[cfg(feature = "mq-admin")] #[tokio::test] async fn mq_admin_config_preserves_admin_url_and_uses_forwarded_connect_override() { diff --git a/crates/dbx-core/src/db/ssh_tunnel.rs b/crates/dbx-core/src/db/ssh_tunnel.rs index f409f67776..7fe05840c6 100644 --- a/crates/dbx-core/src/db/ssh_tunnel.rs +++ b/crates/dbx-core/src/db/ssh_tunnel.rs @@ -1012,11 +1012,58 @@ impl TunnelManager { remote_port: u16, expose_to_lan: bool, allow_exec_channel_proxy: bool, + ) -> Result { + self.start_tunnel_on_local_port( + connection_id, + connect_host, + connect_port, + host_key_host, + host_key_port, + ssh_user, + ssh_password, + ssh_key_path, + ssh_key_passphrase, + use_ssh_agent, + ssh_agent_sock_path, + auth_method, + connect_timeout_secs, + remote_host, + remote_port, + expose_to_lan, + allow_exec_channel_proxy, + None, + ) + .await + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn start_tunnel_on_local_port( + &self, + connection_id: &str, + connect_host: &str, + connect_port: u16, + host_key_host: &str, + host_key_port: u16, + ssh_user: &str, + ssh_password: &str, + ssh_key_path: &str, + ssh_key_passphrase: &str, + use_ssh_agent: bool, + ssh_agent_sock_path: &str, + auth_method: &str, + connect_timeout_secs: u64, + remote_host: &str, + remote_port: u16, + expose_to_lan: bool, + allow_exec_channel_proxy: bool, + requested_local_port: Option, ) -> Result { { let mut tunnels = self.tunnels.lock().await; if let Some(port) = Self::get_active_port(&mut tunnels, connection_id) { - return Ok(port); + if requested_local_port.is_none_or(|requested| requested == port) { + return Ok(port); + } } } @@ -1027,7 +1074,14 @@ impl TunnelManager { { let mut tunnels = self.tunnels.lock().await; if let Some(port) = Self::get_active_port(&mut tunnels, connection_id) { - return Ok(port); + if requested_local_port.is_none_or(|requested| requested == port) { + return Ok(port); + } + if let Some(entry) = tunnels.remove(connection_id) { + for handle in entry.handles { + handle.abort(); + } + } } } let (handle, local_port) = spawn_tunnel( @@ -1048,6 +1102,7 @@ impl TunnelManager { remote_port, expose_to_lan, allow_exec_channel_proxy, + requested_local_port, ) .await?; @@ -1123,6 +1178,7 @@ impl TunnelManager { target_port, is_last && hop.expose_lan, hop.allow_exec_channel_proxy, + None, ) .await .map_err(|err| format!("SSH hop {} failed: {err}", index + 1))?; @@ -1199,12 +1255,9 @@ async fn spawn_tunnel( remote_port: u16, expose_to_lan: bool, allow_exec_channel_proxy: bool, + requested_local_port: Option, ) -> Result<(JoinHandle<()>, u16), String> { - let local_port = portpicker::pick_unused_port().ok_or("No available port")?; - - let bind_addr = if expose_to_lan { "0.0.0.0" } else { "127.0.0.1" }; - let listener = - TcpListener::bind((bind_addr, local_port)).await.map_err(|e| format!("Failed to bind local port: {e}"))?; + let (listener, local_port) = bind_tunnel_listener(expose_to_lan, requested_local_port).await?; // Initial connection: fail fast on bad credentials let session = connect_and_authenticate( @@ -1248,6 +1301,26 @@ async fn spawn_tunnel( Ok((handle, local_port)) } +async fn bind_tunnel_listener( + expose_to_lan: bool, + requested_local_port: Option, +) -> Result<(TcpListener, u16), String> { + let local_port = match requested_local_port { + Some(port) => port, + None => portpicker::pick_unused_port().ok_or("No available port")?, + }; + + let bind_addr = if expose_to_lan { "0.0.0.0" } else { "127.0.0.1" }; + let listener = TcpListener::bind((bind_addr, local_port)).await.map_err(|error| { + if requested_local_port.is_some() { + format!("Failed to bind requested local port {local_port}: {error}") + } else { + format!("Failed to bind local port: {error}") + } + })?; + Ok((listener, local_port)) +} + fn effective_hop_timeout(hop: &SshTunnelConfig) -> u64 { if hop.connect_timeout_secs == 0 { crate::models::connection::default_ssh_connect_timeout_secs() @@ -1300,9 +1373,10 @@ mod tests { use super::SshClient; use super::PROMPT_TEST_LOCK; use super::{ - connect_and_authenticate, effective_hop_timeout, netcat_proxy_command, openssh_padding_len, plan_chain, - read_ssh_string, sanitize_unencrypted_openssh_comment_bytes, server_offers_keyboard_interactive, - server_offers_password, ssh_client_config, HostKeyState, HostKeyVerifier, PlannedTunnel, TunnelManager, + bind_tunnel_listener, connect_and_authenticate, effective_hop_timeout, netcat_proxy_command, + openssh_padding_len, plan_chain, read_ssh_string, sanitize_unencrypted_openssh_comment_bytes, + server_offers_keyboard_interactive, server_offers_password, ssh_client_config, HostKeyState, HostKeyVerifier, + PlannedTunnel, TunnelManager, }; use crate::db::ssh_prompt; use crate::models::connection::{default_ssh_connect_timeout_secs, SshTunnelConfig}; @@ -1319,7 +1393,7 @@ mod tests { use std::sync::Arc; use tempfile::tempdir; use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::TcpStream; + use tokio::net::{TcpListener, TcpStream}; use tokio::sync::mpsc; fn push_u32(bytes: &mut Vec, value: u32) { @@ -1354,6 +1428,28 @@ mod tests { bytes } + #[tokio::test] + async fn requested_local_port_is_bound_exactly() { + let probe = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let port = probe.local_addr().unwrap().port(); + drop(probe); + + let (listener, local_port) = bind_tunnel_listener(false, Some(port)).await.unwrap(); + + assert_eq!(local_port, port); + assert_eq!(listener.local_addr().unwrap().port(), port); + } + + #[tokio::test] + async fn requested_local_port_conflict_is_reported() { + let occupied = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let port = occupied.local_addr().unwrap().port(); + + let error = bind_tunnel_listener(false, Some(port)).await.unwrap_err(); + + assert!(error.contains(&format!("requested local port {port}"))); + } + fn hop(id: &str, host: &str, port: u16) -> SshTunnelConfig { SshTunnelConfig { profile_id: String::new(), diff --git a/crates/dbx-core/src/db/transport_layer_tunnel.rs b/crates/dbx-core/src/db/transport_layer_tunnel.rs index 7e9a547e62..8166f07794 100644 --- a/crates/dbx-core/src/db/transport_layer_tunnel.rs +++ b/crates/dbx-core/src/db/transport_layer_tunnel.rs @@ -47,6 +47,29 @@ pub async fn start_transport_layers( ssh_tunnels: &TunnelManager, proxy_tunnels: &ProxyTunnelManager, http_tunnels: &HttpTunnelManager, +) -> Result { + start_transport_layers_with_final_ssh_local_port( + connection_id, + layers, + remote_host, + remote_port, + None, + ssh_tunnels, + proxy_tunnels, + http_tunnels, + ) + .await +} + +pub(crate) async fn start_transport_layers_with_final_ssh_local_port( + connection_id: &str, + layers: &[TransportLayerConfig], + remote_host: &str, + remote_port: u16, + final_ssh_local_port: Option, + ssh_tunnels: &TunnelManager, + proxy_tunnels: &ProxyTunnelManager, + http_tunnels: &HttpTunnelManager, ) -> Result { if layers.is_empty() { return Err("No transport layers configured".to_string()); @@ -75,7 +98,7 @@ pub async fn start_transport_layers( let local_port = match layer { TransportLayerConfig::Ssh(resolved) => ssh_tunnels - .start_tunnel( + .start_tunnel_on_local_port( &layer_id, &connect_endpoint.host, connect_endpoint.port, @@ -93,6 +116,7 @@ pub async fn start_transport_layers( target_endpoint.port, is_last && resolved.expose_lan, resolved.allow_exec_channel_proxy, + if is_last { final_ssh_local_port } else { None }, ) .await .map_err(|err| format!("SSH layer {} failed: {err}", index + 1))?,