From 5351fc50664059cd11dc6a2eb7f922d09295392b Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 11 Aug 2026 03:16:08 +0800 Subject: [PATCH] fix(rocketmq): route tunneled traffic through SOCKS5 Closes #5742 --- .../com/dbx/agent/rocketmq/RocketMqAgent.java | 35 ++ .../dbx/agent/rocketmq/RocketMqAgentTest.java | 35 ++ crates/dbx-core/src/connection.rs | 121 ++++- crates/dbx-core/src/db/proxy_tunnel.rs | 21 + crates/dbx-core/src/db/ssh_tunnel.rs | 457 +++++++++++++++++- .../dbx-core/src/db/transport_layer_tunnel.rs | 59 +++ crates/dbx-core/src/mq/adapters/kafka.rs | 1 + crates/dbx-core/src/mq/adapters/pulsar.rs | 1 + crates/dbx-core/src/mq/adapters/rabbitmq.rs | 1 + crates/dbx-core/src/mq/adapters/rocketmq.rs | 155 +++++- crates/dbx-core/src/mq/config.rs | 23 + 11 files changed, 883 insertions(+), 26 deletions(-) diff --git a/agents/drivers/rocketmq/src/main/java/com/dbx/agent/rocketmq/RocketMqAgent.java b/agents/drivers/rocketmq/src/main/java/com/dbx/agent/rocketmq/RocketMqAgent.java index a428885d9f..06c445712f 100644 --- a/agents/drivers/rocketmq/src/main/java/com/dbx/agent/rocketmq/RocketMqAgent.java +++ b/agents/drivers/rocketmq/src/main/java/com/dbx/agent/rocketmq/RocketMqAgent.java @@ -3,6 +3,7 @@ package com.dbx.agent.rocketmq; import com.google.gson.*; import org.apache.rocketmq.acl.common.AclClientRPCHook; import org.apache.rocketmq.acl.common.SessionCredentials; +import org.apache.rocketmq.client.ClientConfig; import org.apache.rocketmq.client.QueryResult; import org.apache.rocketmq.client.exception.MQBrokerException; import org.apache.rocketmq.client.exception.MQClientException; @@ -2623,6 +2624,7 @@ public final class RocketMqAgent { ? new DefaultMQAdminExt(rpcHook, timeoutMs) : new DefaultMQAdminExt(timeoutMs); admin.setNamesrvAddr(namesrvAddr(conn)); + applySocksProxy(admin, conn); admin.setAdminExtGroup("_DBX_ROCKETMQ_ADMIN_" + UUID.randomUUID()); admin.setInstanceName("DBX_" + UUID.randomUUID()); admin.start(); @@ -2635,6 +2637,7 @@ public final class RocketMqAgent { ? new DefaultMQProducer("_DBX_ROCKETMQ_PRODUCER", rpcHook) : new DefaultMQProducer("_DBX_ROCKETMQ_PRODUCER"); nextProducer.setNamesrvAddr(namesrvAddr(conn)); + applySocksProxy(nextProducer, conn); nextProducer.setInstanceName("DBX_" + UUID.randomUUID()); // Follow Advanced query timeout — SDK default sendMsgTimeout is only 3s. nextProducer.setSendMsgTimeout(operationBudgetMsAsInt(conn)); @@ -2648,6 +2651,7 @@ public final class RocketMqAgent { ? new DefaultLitePullConsumer(rpcHook) : new DefaultLitePullConsumer(); consumer.setNamesrvAddr(namesrvAddr(conn)); + applySocksProxy(consumer, conn); consumer.setConsumerGroup("_DBX_PEEK_" + UUID.randomUUID()); consumer.setInstanceName("DBX_" + UUID.randomUUID()); consumer.setAutoCommit(false); @@ -2660,6 +2664,7 @@ public final class RocketMqAgent { ? new DefaultMQPullConsumer(MixAll.TOOLS_CONSUMER_GROUP, rpcHook) : new DefaultMQPullConsumer(MixAll.TOOLS_CONSUMER_GROUP); consumer.setNamesrvAddr(namesrvAddr(conn)); + applySocksProxy(consumer, conn); consumer.setInstanceName("DBX_" + UUID.randomUUID()); return consumer; } @@ -2684,6 +2689,33 @@ public final class RocketMqAgent { return addr; } + static boolean applySocksProxy(ClientConfig client, JsonObject conn) { + if (!conn.has("socks_proxy") || !conn.get("socks_proxy").isJsonObject()) { + return false; + } + JsonObject proxy = conn.getAsJsonObject("socks_proxy"); + String host = stringOrEmpty(proxy, "host"); + int port = intOrDefault(proxy, "port", 0); + if (host.isBlank() || port <= 0 || port > 65_535) { + throw new IllegalArgumentException("socks_proxy host and port are required"); + } + + JsonObject route = new JsonObject(); + route.addProperty("addr", formatSocketAddress(host, Integer.toString(port))); + String username = stringOrEmpty(proxy, "username"); + String password = stringOrEmpty(proxy, "password"); + if (!username.isBlank()) { + route.addProperty("username", username); + } + if (!password.isBlank()) { + route.addProperty("password", password); + } + JsonObject routes = new JsonObject(); + routes.add("0.0.0.0/0", route); + client.setSocksProxyConfig(GSON.toJson(routes)); + return true; + } + private static TopicList fetchTopicList(DefaultMQAdminExt admin, String cluster) throws Exception { if (cluster != null && !cluster.isBlank()) { return admin.fetchTopicsByCLuster(cluster); @@ -3146,6 +3178,9 @@ public final class RocketMqAgent { if (!explicit.isBlank()) { return explicit; } + if (conn.has("socks_proxy") && conn.get("socks_proxy").isJsonObject()) { + return brokerAddr; + } String brokerHost = parseHostFromSocketAddress(brokerAddr); String brokerPort = parsePortFromSocketAddress(brokerAddr); if (brokerPort.isBlank()) { diff --git a/agents/drivers/rocketmq/src/test/java/com/dbx/agent/rocketmq/RocketMqAgentTest.java b/agents/drivers/rocketmq/src/test/java/com/dbx/agent/rocketmq/RocketMqAgentTest.java index 83af091c84..5f55b8ea40 100644 --- a/agents/drivers/rocketmq/src/test/java/com/dbx/agent/rocketmq/RocketMqAgentTest.java +++ b/agents/drivers/rocketmq/src/test/java/com/dbx/agent/rocketmq/RocketMqAgentTest.java @@ -23,6 +23,7 @@ import org.apache.rocketmq.remoting.protocol.body.ConsumerConnection; import org.apache.rocketmq.remoting.protocol.route.QueueData; import org.apache.rocketmq.remoting.protocol.route.TopicRouteData; import org.apache.rocketmq.common.message.MessageQueue; +import org.apache.rocketmq.tools.admin.DefaultMQAdminExt; import com.google.gson.JsonObject; import com.google.gson.JsonParser; import java.util.ArrayList; @@ -71,6 +72,28 @@ class RocketMqAgentTest { ); } + @Test + void applySocksProxyConfiguresAllRocketMqDestinations() { + JsonObject conn = JsonParser.parseString(""" + { + "socks_proxy": { + "host": "127.0.0.1", + "port": 41080, + "username": "proxy-user", + "password": "proxy-secret" + } + } + """).getAsJsonObject(); + DefaultMQAdminExt admin = new DefaultMQAdminExt(); + + assertTrue(RocketMqAgent.applySocksProxy(admin, conn)); + + JsonObject routes = JsonParser.parseString(admin.getSocksProxyConfig()).getAsJsonObject(); + assertEquals("127.0.0.1:41080", routes.getAsJsonObject("0.0.0.0/0").get("addr").getAsString()); + assertEquals("proxy-user", routes.getAsJsonObject("0.0.0.0/0").get("username").getAsString()); + assertEquals("proxy-secret", routes.getAsJsonObject("0.0.0.0/0").get("password").getAsString()); + } + @Test void isUserTopicFiltersRocketMqSystemTopics() { assertEquals(false, RocketMqAgent.isUserTopic("RMQ_SYS_TRANS_HALF_TOPIC")); @@ -139,6 +162,18 @@ class RocketMqAgentTest { assertEquals("172.32.0.1:10911", RocketMqAgent.remapBrokerAddrForClient("172.32.0.1:10911", conn)); } + @Test + void remapBrokerAddrPreservesAdvertisedAddressWhenSocksProxyIsActive() { + JsonObject conn = JsonParser.parseString(""" + { + "namesrv_addr": "172.19.191.166:9876", + "socks_proxy": {"host":"127.0.0.1","port":41080} + } + """).getAsJsonObject(); + + assertEquals("172.18.0.3:10911", RocketMqAgent.remapBrokerAddrForClient("172.18.0.3:10911", conn)); + } + @Test void isRfc1918PrivateIpv4MatchesOnlyPrivateRanges() { assertTrue(RocketMqAgent.isRfc1918PrivateIpv4("10.0.0.1")); diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 6ac48970db..4d761881ef 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -24,7 +24,7 @@ use crate::db::proxy_tunnel::ProxyTunnelManager; use crate::db::ssh_tunnel::TunnelManager; use crate::models::connection::{ database_info_from_protocol_value, parse_jdbc_host_port, parse_mongo_first_host, rewrite_jdbc_url_host, - ConnectionConfig, ConnectionTestResult, DatabaseConnectionInfo, DatabaseType, TransportLayerConfig, + ConnectionConfig, ConnectionTestResult, DatabaseConnectionInfo, DatabaseType, ProxyType, TransportLayerConfig, }; use crate::path_utils::expand_tilde; use crate::plugins::{PluginDriverSession, PluginRegistry, PluginRuntimeEnv}; @@ -2465,6 +2465,15 @@ impl AppState { return Err("Oracle TNS connections cannot be combined with SSH, proxy, or HTTP tunnel layers. Remove the transport layer or use Service Name/SID mode.".to_string()); } + #[cfg(feature = "mq-admin")] + if config.db_type == DatabaseType::MessageQueue + && crate::mq::config::MqAdminConfig::from_connection(config)?.system_kind + == crate::mq::types::MqSystemKind::RocketMq + { + self.rocketmq_socks_proxy_for_transport_layers(connection_id, &transport_layers).await?; + return Ok((config.host.clone(), config.port)); + } + let (remote_host, remote_port) = connection_remote_endpoint(config); let local_port = db::transport_layer_tunnel::start_transport_layers( connection_id, @@ -2669,6 +2678,66 @@ impl AppState { Ok(routes) } + #[cfg(feature = "mq-admin")] + async fn rocketmq_socks_proxy_for_transport_layers( + &self, + connection_id: &str, + transport_layers: &[TransportLayerConfig], + ) -> Result { + let final_layer = transport_layers.last().ok_or("No transport layers configured")?; + match final_layer { + TransportLayerConfig::Ssh(_) => { + let local_port = db::transport_layer_tunnel::start_transport_layers_with_final_ssh_socks5( + connection_id, + transport_layers, + &self.tunnels, + &self.proxy_tunnels, + &self.http_tunnels, + ) + .await?; + Ok(crate::mq::config::MqSocksProxy { + host: "127.0.0.1".to_string(), + port: local_port, + username: String::new(), + password: String::new(), + }) + } + TransportLayerConfig::Proxy(proxy) if proxy.proxy_type == ProxyType::Socks5 => { + if transport_layers.len() == 1 { + Ok(crate::mq::config::MqSocksProxy { + host: proxy.host.clone(), + port: proxy.port, + username: proxy.username.clone(), + password: proxy.password.clone(), + }) + } else { + let local_port = db::transport_layer_tunnel::start_transport_layers( + connection_id, + &transport_layers[..transport_layers.len() - 1], + &proxy.host, + proxy.port, + &self.tunnels, + &self.proxy_tunnels, + &self.http_tunnels, + ) + .await?; + Ok(crate::mq::config::MqSocksProxy { + host: "127.0.0.1".to_string(), + port: local_port, + username: proxy.username.clone(), + password: proxy.password.clone(), + }) + } + } + TransportLayerConfig::Proxy(_) => { + Err("RocketMQ requires a SOCKS5 proxy as the final proxy layer".to_string()) + } + TransportLayerConfig::HttpTunnel(_) => { + Err("RocketMQ does not support an HTTP tunnel as the final transport layer".to_string()) + } + } + } + #[cfg(feature = "mq-admin")] pub async fn mq_admin_config_for_connection( &self, @@ -2680,6 +2749,12 @@ impl AppState { return Ok(mqc); } + if mqc.system_kind == crate::mq::types::MqSystemKind::RocketMq { + let transport_layers = self.resolved_transport_layers(config).await?; + let proxy = self.rocketmq_socks_proxy_for_transport_layers(connection_id, &transport_layers).await?; + return Ok(mqc.with_socks_proxy(&proxy.host, proxy.port, &proxy.username, &proxy.password)); + } + if mqc.system_kind == crate::mq::types::MqSystemKind::RabbitMq { let transport_layers = self.resolved_transport_layers(config).await?; let (amqp_host, amqp_port) = crate::mq::adapters::rabbitmq::primary_amqp_endpoint(&mqc)?; @@ -7829,6 +7904,50 @@ for line in sys.stdin: let _ = std::fs::remove_dir_all(dir); } + #[cfg(feature = "mq-admin")] + #[tokio::test] + async fn rocketmq_transport_passes_socks_proxy_to_multi_endpoint_client() { + let (state, dir) = test_app_state().await; + let mut config = mysql_config(None); + config.id = "proxied-rocketmq".to_string(); + config.db_type = DatabaseType::MessageQueue; + config.host = "172.19.191.166".to_string(); + config.port = 9876; + config.external_config = Some(serde_json::json!({ + "systemKind": "rocketmq", + "adminUrl": "", + "auth": { "kind": "none" }, + "extra": { "namesrvAddr": "172.19.191.166:9876" } + })); + config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig { + profile_id: String::new(), + id: "proxy".to_string(), + name: String::new(), + enabled: true, + proxy_type: ProxyType::Socks5, + host: "proxy.internal".to_string(), + port: 1080, + username: "proxy-user".to_string(), + password: "proxy-secret".to_string(), + test_target: None, + })]; + + let endpoint = state.connection_host_port("proxied-rocketmq", &config).await.unwrap(); + assert_eq!(endpoint, ("172.19.191.166".to_string(), 9876)); + assert!(state.proxy_tunnels.local_port("proxied-rocketmq:transport:0").await.is_none()); + + let mqc = state.mq_admin_config_for_connection("proxied-rocketmq", &config).await.unwrap(); + + let socks_proxy = mqc.socks_proxy.expect("RocketMQ transport should configure SOCKS5 routing"); + assert_eq!(socks_proxy.host, "proxy.internal"); + assert_eq!(socks_proxy.port, 1080); + assert_eq!(socks_proxy.username, "proxy-user"); + assert_eq!(socks_proxy.password, "proxy-secret"); + assert!(mqc.connect_override.is_none()); + assert!(state.proxy_tunnels.local_port("proxied-rocketmq:transport:0").await.is_none()); + let _ = std::fs::remove_dir_all(dir); + } + #[cfg(feature = "mq-admin")] #[tokio::test] async fn rabbitmq_transport_uses_separate_amqp_and_management_tunnels() { diff --git a/crates/dbx-core/src/db/proxy_tunnel.rs b/crates/dbx-core/src/db/proxy_tunnel.rs index 5c2dc550ba..a4437a348f 100644 --- a/crates/dbx-core/src/db/proxy_tunnel.rs +++ b/crates/dbx-core/src/db/proxy_tunnel.rs @@ -136,6 +136,27 @@ async fn connect_via_proxy(proxy: &ProxyEndpoint, remote: &RemoteEndpoint) -> Re } } +pub(crate) async fn connect_via_socks5_proxy( + proxy_host: &str, + proxy_port: u16, + proxy_username: &str, + proxy_password: &str, + remote_host: &str, + remote_port: u16, +) -> Result { + connect_via_proxy( + &ProxyEndpoint { + proxy_type: ProxyType::Socks5, + host: proxy_host.to_string(), + port: proxy_port, + username: proxy_username.to_string(), + password: proxy_password.to_string(), + }, + &RemoteEndpoint { host: remote_host.to_string(), port: remote_port }, + ) + .await +} + async fn http_connect( mut stream: TcpStream, proxy: &ProxyEndpoint, diff --git a/crates/dbx-core/src/db/ssh_tunnel.rs b/crates/dbx-core/src/db/ssh_tunnel.rs index 7fe05840c6..92fda9debf 100644 --- a/crates/dbx-core/src/db/ssh_tunnel.rs +++ b/crates/dbx-core/src/db/ssh_tunnel.rs @@ -15,6 +15,7 @@ use russh::keys::{decode_secret_key, key::PrivateKeyWithHashAlg, PrivateKey}; use russh::MethodKind; use russh::MethodSet; use russh::{kex, mac, ChannelOpenFailure, Preferred}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use tokio::sync::Mutex; use tokio::task::JoinHandle; @@ -754,13 +755,90 @@ fn netcat_proxy_command(remote_host: &str, remote_port: u16) -> Result String { + match self { + Self::Fixed { host, port } => format!("{host}:{port}"), + Self::Socks5 => "dynamic SOCKS5 targets".to_string(), + } + } +} + +async fn read_socks5_target(stream: &mut tokio::net::TcpStream) -> Result<(String, u16), String> { + let mut greeting = [0_u8; 2]; + stream.read_exact(&mut greeting).await.map_err(|e| format!("SOCKS5 greeting failed: {e}"))?; + if greeting[0] != 0x05 || greeting[1] == 0 { + return Err("Invalid SOCKS5 greeting".to_string()); + } + let mut methods = vec![0_u8; greeting[1] as usize]; + stream.read_exact(&mut methods).await.map_err(|e| format!("SOCKS5 methods failed: {e}"))?; + if !methods.contains(&0x00) { + let _ = stream.write_all(&[0x05, 0xff]).await; + return Err("SOCKS5 client does not support no-auth mode".to_string()); + } + stream.write_all(&[0x05, 0x00]).await.map_err(|e| format!("SOCKS5 method reply failed: {e}"))?; + + let mut request = [0_u8; 4]; + stream.read_exact(&mut request).await.map_err(|e| format!("SOCKS5 request failed: {e}"))?; + if request[0] != 0x05 { + return Err("Invalid SOCKS5 request version".to_string()); + } + if request[1] != 0x01 { + let _ = write_socks5_reply(stream, 0x07).await; + return Err("SOCKS5 command is not CONNECT".to_string()); + } + + let host = match request[3] { + 0x01 => { + let mut addr = [0_u8; 4]; + stream.read_exact(&mut addr).await.map_err(|e| format!("SOCKS5 IPv4 target failed: {e}"))?; + std::net::Ipv4Addr::from(addr).to_string() + } + 0x03 => { + let mut len = [0_u8; 1]; + stream.read_exact(&mut len).await.map_err(|e| format!("SOCKS5 domain length failed: {e}"))?; + if len[0] == 0 { + let _ = write_socks5_reply(stream, 0x08).await; + return Err("SOCKS5 target host is empty".to_string()); + } + let mut host = vec![0_u8; len[0] as usize]; + stream.read_exact(&mut host).await.map_err(|e| format!("SOCKS5 domain target failed: {e}"))?; + String::from_utf8(host).map_err(|_| "SOCKS5 target host is not UTF-8".to_string())? + } + 0x04 => { + let mut addr = [0_u8; 16]; + stream.read_exact(&mut addr).await.map_err(|e| format!("SOCKS5 IPv6 target failed: {e}"))?; + std::net::Ipv6Addr::from(addr).to_string() + } + _ => { + let _ = write_socks5_reply(stream, 0x08).await; + return Err("Unsupported SOCKS5 address type".to_string()); + } + }; + let mut port = [0_u8; 2]; + stream.read_exact(&mut port).await.map_err(|e| format!("SOCKS5 target port failed: {e}"))?; + Ok((host, u16::from_be_bytes(port))) +} + +async fn write_socks5_reply(stream: &mut tokio::net::TcpStream, status: u8) -> Result<(), String> { + stream + .write_all(&[0x05, status, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]) + .await + .map_err(|e| format!("SOCKS5 reply failed: {e}")) +} + /// Accept connections on the local listener and forward them through the SSH session. /// Returns when the SSH session dies (listener error or session.is_closed()). async fn forward_loop( session: &Handle, listener: &TcpListener, - remote_host: &str, - remote_port: u16, + target: &TunnelTarget, allow_exec_channel_proxy: bool, ) { let mut idle_check = tokio::time::interval(IDLE_SESSION_CHECK_INTERVAL); @@ -802,9 +880,20 @@ async fn forward_loop( break; } + let (remote_host, remote_port) = match target { + TunnelTarget::Fixed { host, port } => (host.clone(), *port), + TunnelTarget::Socks5 => match read_socks5_target(&mut stream).await { + Ok(target) => target, + Err(e) => { + log::debug!("SSH SOCKS5 request rejected: {e}"); + continue; + } + }, + }; + let channel = match session .channel_open_direct_tcpip( - remote_host, + &remote_host, remote_port.into(), peer_addr.ip().to_string(), peer_addr.port().into(), @@ -818,37 +907,63 @@ async fn forward_loop( // JumpServer/Koko deliberately disables SSH direct-tcpip even // for a directly selected asset. An exec channel is still // proxied to that asset, so use netcat there as a byte stream. - let command = match netcat_proxy_command(remote_host, remote_port) { + let command = match netcat_proxy_command(&remote_host, remote_port) { Ok(command) => command, Err(e) => { log::error!("SSH netcat fallback rejected the target: {e}"); - continue; + if matches!(target, TunnelTarget::Socks5) { + let _ = write_socks5_reply(&mut stream, 0x01).await; + continue; + } + break; } }; let channel = match session.channel_open_session().await { Ok(channel) => channel, Err(e) => { log::error!("SSH netcat fallback could not open a session channel: {e}"); + if matches!(target, TunnelTarget::Socks5) { + let _ = write_socks5_reply(&mut stream, 0x01).await; + } continue; } }; if let Err(e) = channel.exec(true, command).await { log::error!("SSH netcat fallback could not start nc: {e}"); - continue; + if matches!(target, TunnelTarget::Socks5) { + let _ = write_socks5_reply(&mut stream, 0x01).await; + continue; + } + break; } log::info!("SSH direct-tcpip is disabled; forwarding through a remote nc session"); channel } Err(russh::Error::ChannelOpenFailure(ChannelOpenFailure::AdministrativelyProhibited)) => { log::warn!("SSH direct-tcpip was administratively prohibited; exec-channel proxy fallback is disabled"); + if matches!(target, TunnelTarget::Socks5) { + let _ = write_socks5_reply(&mut stream, 0x02).await; + continue; + } break; } Err(e) => { log::error!("SSH direct-tcpip failed: {e}"); + if matches!(target, TunnelTarget::Socks5) { + let _ = write_socks5_reply(&mut stream, 0x01).await; + continue; + } break; } }; + if matches!(target, TunnelTarget::Socks5) { + if let Err(e) = write_socks5_reply(&mut stream, 0x00).await { + log::debug!("SSH SOCKS5 success reply failed: {e}"); + continue; + } + } + tokio::spawn(async move { let mut channel_stream = channel.into_stream(); if let Err(e) = tokio::io::copy_bidirectional(&mut stream, &mut channel_stream).await { @@ -880,14 +995,13 @@ async fn tunnel_reconnect_loop( connect_timeout_secs: u64, known_hosts_path: PathBuf, listener: TcpListener, - remote_host: String, - remote_port: u16, + target: TunnelTarget, allow_exec_channel_proxy: bool, ) { loop { - log::info!("SSH tunnel active: {}:{} -> {}:{}", connect_host, connect_port, remote_host, remote_port); + log::info!("SSH tunnel active: {}:{} -> {}", connect_host, connect_port, target.description()); - forward_loop(&session, &listener, &remote_host, remote_port, allow_exec_channel_proxy).await; + forward_loop(&session, &listener, &target, allow_exec_channel_proxy).await; log::warn!("SSH tunnel connection lost ({}:{}), reconnecting...", connect_host, connect_port); @@ -947,9 +1061,16 @@ async fn tunnel_reconnect_loop( } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TunnelKind { + Fixed, + Socks5, +} + struct TunnelEntry { handles: Vec>, local_port: u16, + kind: TunnelKind, } #[cfg(test)] @@ -1060,7 +1181,7 @@ impl TunnelManager { ) -> Result { { let mut tunnels = self.tunnels.lock().await; - if let Some(port) = Self::get_active_port(&mut tunnels, connection_id) { + if let Some(port) = Self::get_active_port(&mut tunnels, connection_id, TunnelKind::Fixed) { if requested_local_port.is_none_or(|requested| requested == port) { return Ok(port); } @@ -1073,7 +1194,7 @@ impl TunnelManager { // A concurrent caller may have completed while this task waited. { let mut tunnels = self.tunnels.lock().await; - if let Some(port) = Self::get_active_port(&mut tunnels, connection_id) { + if let Some(port) = Self::get_active_port(&mut tunnels, connection_id, TunnelKind::Fixed) { if requested_local_port.is_none_or(|requested| requested == port) { return Ok(port); } @@ -1106,16 +1227,91 @@ impl TunnelManager { ) .await?; - self.tunnels.lock().await.insert(connection_id.to_string(), TunnelEntry { handles: vec![handle], local_port }); + self.tunnels.lock().await.insert( + connection_id.to_string(), + TunnelEntry { handles: vec![handle], local_port, kind: TunnelKind::Fixed }, + ); + Ok(local_port) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn start_socks5_proxy( + &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, + allow_exec_channel_proxy: bool, + ) -> Result { + { + let mut tunnels = self.tunnels.lock().await; + if let Some(port) = Self::get_active_port(&mut tunnels, connection_id, TunnelKind::Socks5) { + return Ok(port); + } + } + + let start_lock = self.start_lock(connection_id).await; + let _start_guard = start_lock.lock().await; + { + let mut tunnels = self.tunnels.lock().await; + if let Some(port) = Self::get_active_port(&mut tunnels, connection_id, TunnelKind::Socks5) { + return Ok(port); + } + if let Some(entry) = tunnels.remove(connection_id) { + for handle in entry.handles { + handle.abort(); + } + } + } + + let (handle, local_port) = spawn_socks5_proxy( + 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, + &self.known_hosts_path, + allow_exec_channel_proxy, + ) + .await?; + + self.tunnels.lock().await.insert( + connection_id.to_string(), + TunnelEntry { handles: vec![handle], local_port, kind: TunnelKind::Socks5 }, + ); Ok(local_port) } /// Returns the local port for a cached tunnel entry, or `None` if the entry /// is stale (all background handles have exited). - fn get_active_port(tunnels: &mut HashMap, connection_id: &str) -> Option { + fn get_active_port( + tunnels: &mut HashMap, + connection_id: &str, + expected_kind: TunnelKind, + ) -> Option { let entry = tunnels.get(connection_id)?; - if entry.handles.iter().all(|h| h.is_finished()) { - tunnels.remove(connection_id); + if entry.kind != expected_kind || entry.handles.iter().all(|h| h.is_finished()) { + if let Some(entry) = tunnels.remove(connection_id) { + for handle in entry.handles { + handle.abort(); + } + } return None; } Some(entry.local_port) @@ -1133,7 +1329,7 @@ impl TunnelManager { } { let mut tunnels = self.tunnels.lock().await; - if let Some(port) = Self::get_active_port(&mut tunnels, connection_id) { + if let Some(port) = Self::get_active_port(&mut tunnels, connection_id, TunnelKind::Fixed) { return Ok(port); } } @@ -1142,7 +1338,7 @@ impl TunnelManager { let _start_guard = start_lock.lock().await; { let mut tunnels = self.tunnels.lock().await; - if let Some(port) = Self::get_active_port(&mut tunnels, connection_id) { + if let Some(port) = Self::get_active_port(&mut tunnels, connection_id, TunnelKind::Fixed) { return Ok(port); } } @@ -1188,10 +1384,10 @@ impl TunnelManager { next_connect_endpoint = Some(("127.0.0.1".to_string(), local_port)); } - self.tunnels - .lock() - .await - .insert(connection_id.to_string(), TunnelEntry { handles, local_port: final_local_port }); + self.tunnels.lock().await.insert( + connection_id.to_string(), + TunnelEntry { handles, local_port: final_local_port, kind: TunnelKind::Fixed }, + ); Ok(final_local_port) } @@ -1256,6 +1452,87 @@ async fn spawn_tunnel( expose_to_lan: bool, allow_exec_channel_proxy: bool, requested_local_port: Option, +) -> Result<(JoinHandle<()>, u16), String> { + spawn_tunnel_target( + 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, + known_hosts_path, + TunnelTarget::Fixed { host: remote_host.to_string(), port: remote_port }, + expose_to_lan, + allow_exec_channel_proxy, + requested_local_port, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn spawn_socks5_proxy( + 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, + known_hosts_path: &Path, + allow_exec_channel_proxy: bool, +) -> Result<(JoinHandle<()>, u16), String> { + spawn_tunnel_target( + 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, + known_hosts_path, + TunnelTarget::Socks5, + false, + allow_exec_channel_proxy, + None, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn spawn_tunnel_target( + 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, + known_hosts_path: &Path, + target: TunnelTarget, + expose_to_lan: bool, + allow_exec_channel_proxy: bool, + requested_local_port: Option, ) -> Result<(JoinHandle<()>, u16), String> { let (listener, local_port) = bind_tunnel_listener(expose_to_lan, requested_local_port).await?; @@ -1293,8 +1570,7 @@ async fn spawn_tunnel( connect_timeout_secs, known_hosts_path.to_path_buf(), listener, - remote_host.to_string(), - remote_port, + target, allow_exec_channel_proxy, )); @@ -1982,6 +2258,139 @@ uveF/dLmnVN1IriEyEvHAAAACGRieC10ZXN0AQIDBAU= (port, task) } + struct DirectEchoServer { + target: Arc>>, + } + + impl server::Server for DirectEchoServer { + type Handler = DirectEchoHandler; + + fn new_client(&mut self, _peer: Option) -> DirectEchoHandler { + DirectEchoHandler { target: self.target.clone() } + } + } + + struct DirectEchoHandler { + target: Arc>>, + } + + impl server::Handler for DirectEchoHandler { + type Error = russh::Error; + + async fn auth_none(&mut self, _user: &str) -> Result { + Ok(Auth::Accept) + } + + async fn channel_open_direct_tcpip( + &mut self, + _channel: Channel, + host_to_connect: &str, + port_to_connect: u32, + _originator_address: &str, + _originator_port: u32, + _session: &mut server::Session, + ) -> Result { + *self.target.lock().unwrap() = Some((host_to_connect.to_string(), port_to_connect)); + Ok(true) + } + + async fn data( + &mut self, + channel: ChannelId, + data: &[u8], + session: &mut server::Session, + ) -> Result<(), Self::Error> { + session.data(channel, data.to_vec())?; + Ok(()) + } + } + + #[tokio::test] + async fn socks5_proxy_forwards_requested_target_over_ssh() { + let _guard = PROMPT_TEST_LOCK.lock().await; + ssh_prompt::clear_ssh_prompt_gateway(); + let server_key = decode_secret_key(TEST_SERVER_KEY_PEM, None).expect("decode test server key"); + let server_config = server::Config { keys: vec![server_key], ..Default::default() }; + let ssh_port = portpicker::pick_unused_port().expect("no free port"); + let target = Arc::new(std::sync::Mutex::new(None)); + let mut server = DirectEchoServer { target: target.clone() }; + let server_task = tokio::spawn(async move { + let _ = server.run_on_address(Arc::new(server_config), ("127.0.0.1", ssh_port)).await; + }); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + + let dir = tempdir().unwrap(); + let known_hosts_path = dir.path().join("known_hosts"); + HostKeyVerifier::new(known_hosts_path).learn("127.0.0.1", ssh_port, &test_server_public_key()).unwrap(); + let manager = TunnelManager::new(dir.path().to_path_buf()); + let local_port = manager + .start_socks5_proxy( + "rocketmq-socks", + "127.0.0.1", + ssh_port, + "127.0.0.1", + ssh_port, + "user", + "", + "", + "", + false, + "", + "none", + 5, + false, + ) + .await + .unwrap(); + let reused_port = manager + .start_socks5_proxy( + "rocketmq-socks", + "127.0.0.1", + ssh_port, + "127.0.0.1", + ssh_port, + "user", + "", + "", + "", + false, + "", + "none", + 5, + false, + ) + .await + .unwrap(); + assert_eq!(reused_port, local_port); + + let mut client = TcpStream::connect(("127.0.0.1", local_port)).await.unwrap(); + client.write_all(&[0x05, 0x01, 0x00]).await.unwrap(); + let mut greeting = [0_u8; 2]; + client.read_exact(&mut greeting).await.unwrap(); + assert_eq!(greeting, [0x05, 0x00]); + + let host = b"broker.internal"; + let mut request = vec![0x05, 0x01, 0x00, 0x03, host.len() as u8]; + request.extend_from_slice(host); + request.extend_from_slice(&10911_u16.to_be_bytes()); + client.write_all(&request).await.unwrap(); + let mut reply = [0_u8; 10]; + client.read_exact(&mut reply).await.unwrap(); + assert_eq!(reply[1], 0x00); + + client.write_all(b"ping").await.unwrap(); + let mut echoed = [0_u8; 4]; + tokio::time::timeout(std::time::Duration::from_secs(5), client.read_exact(&mut echoed)) + .await + .expect("SOCKS5 echo timeout") + .expect("SOCKS5 echo read"); + assert_eq!(&echoed, b"ping"); + assert_eq!(*target.lock().unwrap(), Some(("broker.internal".to_string(), 10911))); + + manager.stop_tunnel("rocketmq-socks").await; + server_task.abort(); + } + struct NetcatFallbackServer { command: Arc>>, } diff --git a/crates/dbx-core/src/db/transport_layer_tunnel.rs b/crates/dbx-core/src/db/transport_layer_tunnel.rs index 8166f07794..ad33eca79b 100644 --- a/crates/dbx-core/src/db/transport_layer_tunnel.rs +++ b/crates/dbx-core/src/db/transport_layer_tunnel.rs @@ -70,6 +70,46 @@ pub(crate) async fn start_transport_layers_with_final_ssh_local_port( ssh_tunnels: &TunnelManager, proxy_tunnels: &ProxyTunnelManager, http_tunnels: &HttpTunnelManager, +) -> Result { + start_transport_layers_internal( + connection_id, + layers, + remote_host, + remote_port, + final_ssh_local_port, + false, + ssh_tunnels, + proxy_tunnels, + http_tunnels, + ) + .await +} + +pub(crate) async fn start_transport_layers_with_final_ssh_socks5( + connection_id: &str, + layers: &[TransportLayerConfig], + ssh_tunnels: &TunnelManager, + proxy_tunnels: &ProxyTunnelManager, + http_tunnels: &HttpTunnelManager, +) -> Result { + if !matches!(layers.last(), Some(TransportLayerConfig::Ssh(_))) { + return Err("Dynamic SOCKS5 routing requires SSH as the final transport layer".to_string()); + } + start_transport_layers_internal(connection_id, layers, "", 0, None, true, ssh_tunnels, proxy_tunnels, http_tunnels) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn start_transport_layers_internal( + connection_id: &str, + layers: &[TransportLayerConfig], + remote_host: &str, + remote_port: u16, + final_ssh_local_port: Option, + final_ssh_socks5: bool, + ssh_tunnels: &TunnelManager, + proxy_tunnels: &ProxyTunnelManager, + http_tunnels: &HttpTunnelManager, ) -> Result { if layers.is_empty() { return Err("No transport layers configured".to_string()); @@ -97,6 +137,25 @@ pub(crate) async fn start_transport_layers_with_final_ssh_local_port( }; let local_port = match layer { + TransportLayerConfig::Ssh(resolved) if is_last && final_ssh_socks5 => ssh_tunnels + .start_socks5_proxy( + &layer_id, + &connect_endpoint.host, + connect_endpoint.port, + &resolved.host, + resolved.port, + &resolved.user, + &resolved.password, + &resolved.key_path, + &resolved.key_passphrase, + resolved.use_ssh_agent, + &resolved.ssh_agent_sock_path, + &resolved.auth_method, + effective_ssh_connect_timeout_secs(resolved.connect_timeout_secs), + resolved.allow_exec_channel_proxy, + ) + .await + .map_err(|err| format!("SSH layer {} failed: {err}", index + 1))?, TransportLayerConfig::Ssh(resolved) => ssh_tunnels .start_tunnel_on_local_port( &layer_id, diff --git a/crates/dbx-core/src/mq/adapters/kafka.rs b/crates/dbx-core/src/mq/adapters/kafka.rs index 98c4caf777..c798a24852 100644 --- a/crates/dbx-core/src/mq/adapters/kafka.rs +++ b/crates/dbx-core/src/mq/adapters/kafka.rs @@ -803,6 +803,7 @@ mod tests { token_signing: None, connect_override: None, management_connect_override: None, + socks_proxy: None, query_timeout_secs: crate::mq::config::DEFAULT_MQ_QUERY_TIMEOUT_SECS, connect_timeout_secs: crate::mq::config::DEFAULT_MQ_CONNECT_TIMEOUT_SECS, extra, diff --git a/crates/dbx-core/src/mq/adapters/pulsar.rs b/crates/dbx-core/src/mq/adapters/pulsar.rs index dc4b18d75b..436ce0193f 100644 --- a/crates/dbx-core/src/mq/adapters/pulsar.rs +++ b/crates/dbx-core/src/mq/adapters/pulsar.rs @@ -1154,6 +1154,7 @@ mod tests { token_signing: None, connect_override: None, management_connect_override: None, + socks_proxy: None, query_timeout_secs: crate::mq::config::DEFAULT_MQ_QUERY_TIMEOUT_SECS, connect_timeout_secs: crate::mq::config::DEFAULT_MQ_CONNECT_TIMEOUT_SECS, extra: serde_json::Value::Null, diff --git a/crates/dbx-core/src/mq/adapters/rabbitmq.rs b/crates/dbx-core/src/mq/adapters/rabbitmq.rs index b4f9956299..8c7d783c87 100644 --- a/crates/dbx-core/src/mq/adapters/rabbitmq.rs +++ b/crates/dbx-core/src/mq/adapters/rabbitmq.rs @@ -1029,6 +1029,7 @@ mod tests { token_signing: None, connect_override: None, management_connect_override: None, + socks_proxy: None, query_timeout_secs: crate::mq::config::DEFAULT_MQ_QUERY_TIMEOUT_SECS, connect_timeout_secs: crate::mq::config::DEFAULT_MQ_CONNECT_TIMEOUT_SECS, extra, diff --git a/crates/dbx-core/src/mq/adapters/rocketmq.rs b/crates/dbx-core/src/mq/adapters/rocketmq.rs index d035ddd327..46a1bf4c8f 100644 --- a/crates/dbx-core/src/mq/adapters/rocketmq.rs +++ b/crates/dbx-core/src/mq/adapters/rocketmq.rs @@ -842,6 +842,12 @@ fn build_connection_params(cfg: &MqAdminConfig) -> serde_json::Value { "tls_skip_verify": cfg.tls_skip_verify, "request_timeout_ms": cfg.request_timeout_ms(), "connect_timeout_ms": cfg.connect_timeout_ms(), + "socks_proxy": cfg.socks_proxy.as_ref().map(|proxy| serde_json::json!({ + "host": proxy.host, + "port": proxy.port, + "username": proxy.username, + "password": proxy.password, + })), }) } @@ -952,7 +958,90 @@ fn rocketmq_subscription_from_group(group: &serde_json::Value) -> SubscriptionIn /// Fail fast on unreachable NameServer before paying for JVM agent startup. /// Called outside the connect-timeout wall so the probe does not steal JVM budget. pub(crate) async fn probe_namesrv_before_connect(cfg: &MqAdminConfig, budget: Duration) -> Result<(), String> { - probe_namesrv_tcp(&namesrv_addr(cfg), budget).await + if let Some(proxy) = &cfg.socks_proxy { + probe_namesrv_via_socks5(&namesrv_addr(cfg), proxy, budget).await + } else { + probe_namesrv_tcp(&namesrv_addr(cfg), budget).await + } +} + +async fn probe_namesrv_via_socks5( + namesrv: &str, + proxy: &crate::mq::config::MqSocksProxy, + budget: Duration, +) -> Result<(), String> { + let targets: Vec = + namesrv.split(';').map(str::trim).filter(|part| !part.is_empty()).map(str::to_string).collect(); + if targets.is_empty() { + return Err("RocketMQ namesrv_addr is empty".to_string()); + } + + let deadline = tokio::time::Instant::now() + budget; + let per_attempt = + if targets.len() <= 1 { budget } else { (budget / targets.len() as u32).max(Duration::from_millis(500)) }; + let mut last_error = String::new(); + for target in targets { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + break; + } + let (host, port) = match parse_namesrv_target(&target) { + Ok(endpoint) => endpoint, + Err(error) => { + last_error = error; + continue; + } + }; + let attempt = remaining.min(per_attempt); + match timeout( + attempt, + crate::db::proxy_tunnel::connect_via_socks5_proxy( + &proxy.host, + proxy.port, + &proxy.username, + &proxy.password, + &host, + port, + ), + ) + .await + { + Ok(Ok(_stream)) => return Ok(()), + Ok(Err(error)) => { + last_error = format!("Cannot reach RocketMQ NameServer {target} through SOCKS5 proxy: {error}"); + } + Err(_) => { + last_error = format!( + "RocketMQ NameServer {target} connect through SOCKS5 proxy timed out after {}ms", + attempt.as_millis() + ); + } + } + } + Err(if last_error.is_empty() { + format!("RocketMQ NameServer connect through SOCKS5 proxy timed out after {}s", budget.as_secs()) + } else { + last_error + }) +} + +fn parse_namesrv_target(target: &str) -> Result<(String, u16), String> { + let target = target.trim(); + let (host, port) = if let Some(rest) = target.strip_prefix('[') { + let (host, port) = + rest.split_once("]:").ok_or_else(|| format!("RocketMQ NameServer address '{target}' is invalid"))?; + (host, port) + } else { + target.rsplit_once(':').ok_or_else(|| format!("RocketMQ NameServer address '{target}' is invalid"))? + }; + if host.trim().is_empty() { + return Err(format!("RocketMQ NameServer address '{target}' is invalid")); + } + let port = port + .trim() + .parse::() + .map_err(|_| format!("RocketMQ NameServer address '{target}' has an invalid port"))?; + Ok((host.trim().to_string(), port)) } /// Probe NameServer addresses so connect fails before spawning the JVM agent. @@ -1060,7 +1149,10 @@ fn peeked_message_from_agent_json(idx: usize, message: &serde_json::Value) -> Pe mod tests { use super::*; use crate::mq::auth::MqAuth; + use crate::mq::config::MqSocksProxy; use crate::mq::types::MqSystemKind; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; fn rocketmq_config(extra: serde_json::Value, auth: MqAuth) -> MqAdminConfig { MqAdminConfig { @@ -1072,6 +1164,7 @@ mod tests { token_signing: None, connect_override: None, management_connect_override: None, + socks_proxy: None, query_timeout_secs: crate::mq::config::DEFAULT_MQ_QUERY_TIMEOUT_SECS, connect_timeout_secs: crate::mq::config::DEFAULT_MQ_CONNECT_TIMEOUT_SECS, extra, @@ -1084,6 +1177,41 @@ mod tests { assert!(err.contains("empty"), "{err}"); } + #[tokio::test] + async fn probe_namesrv_uses_socks_proxy_for_logical_target() { + let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let proxy_port = listener.local_addr().unwrap().port(); + let (target_tx, target_rx) = tokio::sync::oneshot::channel(); + let proxy_task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut greeting = [0_u8; 3]; + stream.read_exact(&mut greeting).await.unwrap(); + assert_eq!(greeting, [0x05, 0x01, 0x00]); + stream.write_all(&[0x05, 0x00]).await.unwrap(); + + let mut request = [0_u8; 10]; + stream.read_exact(&mut request).await.unwrap(); + assert_eq!(&request[..4], &[0x05, 0x01, 0x00, 0x01]); + let host = std::net::Ipv4Addr::new(request[4], request[5], request[6], request[7]).to_string(); + let port = u16::from_be_bytes([request[8], request[9]]); + let _ = target_tx.send((host, port)); + stream.write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await.unwrap(); + }); + + let mut cfg = rocketmq_config(serde_json::json!({ "namesrvAddr": "172.19.191.166:9876" }), MqAuth::None); + cfg.socks_proxy = Some(MqSocksProxy { + host: "127.0.0.1".to_string(), + port: proxy_port, + username: String::new(), + password: String::new(), + }); + + probe_namesrv_before_connect(&cfg, Duration::from_secs(2)).await.unwrap(); + + assert_eq!(target_rx.await.unwrap(), ("172.19.191.166".to_string(), 9876)); + proxy_task.await.unwrap(); + } + #[test] fn connection_params_map_namesrv_and_acl_credentials() { let cfg = rocketmq_config( @@ -1118,6 +1246,31 @@ mod tests { assert_eq!(params.get("connect_timeout_ms").and_then(|v| v.as_u64()), Some(15_000)); } + #[test] + fn connection_params_include_runtime_socks_proxy() { + let mut cfg = rocketmq_config(serde_json::json!({ "namesrvAddr": "mq.internal:9876" }), MqAuth::None); + cfg.socks_proxy = Some(MqSocksProxy { + host: "127.0.0.1".to_string(), + port: 41080, + username: "proxy-user".to_string(), + password: "proxy-secret".to_string(), + }); + + let params = build_connection_params(&cfg); + + assert_eq!(params.pointer("/socks_proxy/host").and_then(|v| v.as_str()), Some("127.0.0.1")); + assert_eq!(params.pointer("/socks_proxy/port").and_then(|v| v.as_u64()), Some(41080)); + assert_eq!(params.pointer("/socks_proxy/username").and_then(|v| v.as_str()), Some("proxy-user")); + assert_eq!(params.pointer("/socks_proxy/password").and_then(|v| v.as_str()), Some("proxy-secret")); + } + + #[test] + fn parse_namesrv_target_accepts_ipv4_domain_and_ipv6() { + assert_eq!(parse_namesrv_target("172.19.191.166:9876").unwrap(), ("172.19.191.166".to_string(), 9876)); + assert_eq!(parse_namesrv_target("mq.internal:9876").unwrap(), ("mq.internal".to_string(), 9876)); + assert_eq!(parse_namesrv_target("[2001:db8::1]:9876").unwrap(), ("2001:db8::1".to_string(), 9876)); + } + #[test] fn reset_cursor_params_preserve_timestamp_position() { let topic = TopicRef { diff --git a/crates/dbx-core/src/mq/config.rs b/crates/dbx-core/src/mq/config.rs index 0105c8dc6b..ae3fcf643c 100644 --- a/crates/dbx-core/src/mq/config.rs +++ b/crates/dbx-core/src/mq/config.rs @@ -30,6 +30,17 @@ pub struct MqConnectOverride { pub port: u16, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct MqSocksProxy { + pub host: String, + pub port: u16, + #[serde(default)] + pub username: String, + #[serde(default)] + pub password: String, +} + /// Configuration for an MQ admin connection, decoded from /// `ConnectionConfig.external_config`. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -60,6 +71,8 @@ pub struct MqAdminConfig { /// its Management HTTP API listens on an independently configured port. #[serde(skip)] pub management_connect_override: Option, + #[serde(skip)] + pub socks_proxy: Option, /// Runtime-only: from `ConnectionConfig.query_timeout_secs` (`0` = unlimited). #[serde(skip)] pub query_timeout_secs: u64, @@ -141,6 +154,16 @@ impl MqAdminConfig { self.management_connect_override = Some(MqConnectOverride { host: host.to_string(), port }); self } + + pub fn with_socks_proxy(mut self, host: &str, port: u16, username: &str, password: &str) -> Self { + self.socks_proxy = Some(MqSocksProxy { + host: host.to_string(), + port, + username: username.to_string(), + password: password.to_string(), + }); + self + } } pub fn admin_url_with_endpoint(admin_url: &str, host: &str, port: u16) -> Result {