github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 08:13:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(kafka): bind SSH tunnel to advertised loopback port

Browse Source 
Closes #5571
This commit is contained in:
t8y2
2026-08-09 14:47:14 +08:00
parent 71ccb3d97d
commit e7f496bf0b
3 changed files with 210 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/connection.rs
+78 -6
View File
@@ -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() {
crates/dbx-core/src/db/ssh_tunnel.rs
+107 -11
View File
@@ -1012,11 +1012,58 @@ impl TunnelManager {
remote_port: u16,
expose_to_lan: bool,
allow_exec_channel_proxy: bool,
) -> Result<u16, String> {
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<u16>,
) -> Result<u16, String> {
{
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<u16>,
) -> 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<u16>,
) -> 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<u8>, 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(),
crates/dbx-core/src/db/transport_layer_tunnel.rs
+25 -1
View File
@@ -47,6 +47,29 @@ pub async fn start_transport_layers(
ssh_tunnels: &TunnelManager,
proxy_tunnels: &ProxyTunnelManager,
http_tunnels: &HttpTunnelManager,
) -> Result<u16, String> {
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<u16>,
ssh_tunnels: &TunnelManager,
proxy_tunnels: &ProxyTunnelManager,
http_tunnels: &HttpTunnelManager,
) -> Result<u16, String> {
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))?,