diff --git a/Cargo.lock b/Cargo.lock index ed1dcc6e94..6c97d3d145 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1880,7 +1880,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2214,7 +2214,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3212,7 +3212,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.61.2", + "windows-core 0.62.2", ] [[package]] @@ -4196,7 +4196,7 @@ dependencies = [ "png 0.18.1", "serde", "thiserror 2.0.18", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -4287,7 +4287,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6235,7 +6235,7 @@ dependencies = [ [[package]] name = "rust-oracle" version = "0.1.6" -source = "git+https://github.com/t8y2/rust-oracle?branch=main#aa5a936c7bfe566123302e885a6965f279c041d9" +source = "git+https://github.com/t8y2/rust-oracle?branch=main#dddc2565d3e7f8c8275f6afc3ba3bc531ed5a2d2" dependencies = [ "aes 0.8.4", "async-trait", @@ -6351,7 +6351,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6452,7 +6452,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -7114,7 +7114,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -8054,7 +8054,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -8553,7 +8553,7 @@ dependencies = [ "png 0.18.1", "serde", "thiserror 2.0.18", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] diff --git a/crates/dbx-core/Cargo.toml b/crates/dbx-core/Cargo.toml index 139d3183e5..9c55515cae 100644 --- a/crates/dbx-core/Cargo.toml +++ b/crates/dbx-core/Cargo.toml @@ -23,7 +23,7 @@ redis = { version = "0.32.2", features = ["tokio-comp", "tls-rustls", "tokio-rus rustls = { version = "0.23", features = ["aws-lc-rs"] } duckdb = "1.3.2" tiberius = { version = "0.12.3", default-features = false, features = ["tds73", "chrono", "rust_decimal", "rustls"] } -reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls", "socks"] } futures = "0.3" mongodb = "3.2.5" rust-oracle = { git = "https://github.com/t8y2/rust-oracle", branch = "main" } diff --git a/crates/dbx-core/src/ai.rs b/crates/dbx-core/src/ai.rs index 23c0a88779..a8923eb570 100644 --- a/crates/dbx-core/src/ai.rs +++ b/crates/dbx-core/src/ai.rs @@ -70,6 +70,10 @@ pub struct AiConfig { pub model: String, #[serde(default)] pub api_style: AiApiStyle, + #[serde(default)] + pub proxy_enabled: bool, + #[serde(default)] + pub proxy_url: String, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -218,6 +222,15 @@ fn validate_config(config: &AiConfig) -> Result<(), String> { Ok(()) } +pub fn build_ai_http_client(config: &AiConfig, timeout_secs: u64) -> Result { + let mut builder = reqwest::Client::builder().timeout(std::time::Duration::from_secs(timeout_secs)); + if config.proxy_enabled && !config.proxy_url.trim().is_empty() { + let proxy = reqwest::Proxy::all(config.proxy_url.trim()).map_err(|e| format!("Invalid AI proxy URL: {e}"))?; + builder = builder.proxy(proxy); + } + builder.build().map_err(|e| e.to_string()) +} + // --------------------------------------------------------------------------- // Non-streaming calls // --------------------------------------------------------------------------- @@ -339,8 +352,7 @@ pub async fn call_responses_api(client: &reqwest::Client, request: AiCompletionR pub async fn test_connection_core(config: &AiConfig) -> Result { validate_config(config)?; - let client = - reqwest::Client::builder().timeout(std::time::Duration::from_secs(15)).build().map_err(|e| e.to_string())?; + let client = build_ai_http_client(config, 15)?; let request = AiCompletionRequest { config: config.clone(), @@ -366,8 +378,7 @@ pub async fn test_connection_core(config: &AiConfig) -> Result { pub async fn complete(request: &AiCompletionRequest) -> Result { validate_config(&request.config)?; - let client = - reqwest::Client::builder().timeout(std::time::Duration::from_secs(60)).build().map_err(|e| e.to_string())?; + let client = build_ai_http_client(&request.config, 60)?; match request.config.provider { AiProvider::Claude => call_claude(&client, request.clone()).await, @@ -393,8 +404,7 @@ pub async fn stream( ) -> Result<(), String> { validate_config(&request.config)?; - let client = - reqwest::Client::builder().timeout(std::time::Duration::from_secs(120)).build().map_err(|e| e.to_string())?; + let client = build_ai_http_client(&request.config, 120)?; match request.config.provider { AiProvider::Claude => stream_claude(&client, session_id, request, cancelled, &on_chunk).await, @@ -724,3 +734,40 @@ pub fn load_config(path: &Path) -> Result, String> { let json = std::fs::read_to_string(path).map_err(|e| e.to_string())?; serde_json::from_str(&json).map(Some).map_err(|e| e.to_string()) } + +#[cfg(test)] +mod tests { + use super::{build_ai_http_client, AiApiStyle, AiConfig, AiProvider}; + + #[test] + fn ai_config_proxy_fields_default_for_legacy_config() { + let config: AiConfig = serde_json::from_value(serde_json::json!({ + "provider": "openai", + "apiKey": "key", + "endpoint": "https://api.openai.com/v1/chat/completions", + "model": "gpt-4o", + "apiStyle": "completions" + })) + .unwrap(); + + assert_eq!(config.proxy_enabled, false); + assert_eq!(config.proxy_url, ""); + } + + #[test] + fn ai_http_client_rejects_invalid_proxy_url() { + let config = AiConfig { + provider: AiProvider::Openai, + api_key: "key".to_string(), + endpoint: "https://api.openai.com/v1/chat/completions".to_string(), + model: "gpt-4o".to_string(), + api_style: AiApiStyle::Completions, + proxy_enabled: true, + proxy_url: "not a proxy url".to_string(), + }; + + let err = build_ai_http_client(&config, 1).unwrap_err(); + + assert!(err.contains("Invalid AI proxy URL")); + } +} diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 19ed4ea095..56280d6d95 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use tokio::sync::RwLock; use crate::db; +use crate::db::proxy_tunnel::ProxyTunnelManager; use crate::db::ssh_tunnel::TunnelManager; use crate::external; use crate::models::connection::{parse_mongo_first_host, ConnectionConfig, DatabaseType}; @@ -93,6 +94,7 @@ pub struct AppState { pub configs: RwLock>, pub running_queries: RunningQueries, pub tunnels: TunnelManager, + pub proxy_tunnels: ProxyTunnelManager, pub storage: Storage, pub plugins: PluginRegistry, } @@ -126,6 +128,7 @@ impl AppState { configs: RwLock::new(HashMap::new()), running_queries: RunningQueries::default(), tunnels: TunnelManager::new(), + proxy_tunnels: ProxyTunnelManager::new(), storage, plugins: PluginRegistry::new(plugin_dir), } @@ -302,6 +305,37 @@ impl AppState { config: &ConnectionConfig, ) -> Result<(String, u16), String> { if !config.ssh_enabled || config.ssh_host.is_empty() { + if config.proxy_enabled && !config.proxy_host.is_empty() { + if let Some(local_port) = self.proxy_tunnels.local_port(connection_id).await { + return Ok(("127.0.0.1".to_string(), local_port)); + } + + let (remote_host, remote_port) = if config.db_type == DatabaseType::MongoDb { + config + .connection_string + .as_deref() + .filter(|s| !s.is_empty()) + .and_then(parse_mongo_first_host) + .unwrap_or_else(|| (config.host.clone(), config.port)) + } else { + (config.host.clone(), config.port) + }; + + let local_port = self + .proxy_tunnels + .start_tunnel( + connection_id, + config.proxy_type, + &config.proxy_host, + config.proxy_port, + &config.proxy_username, + &config.proxy_password, + &remote_host, + remote_port, + ) + .await?; + return Ok(("127.0.0.1".to_string(), local_port)); + } return Ok((config.host.clone(), config.port)); } @@ -412,7 +446,7 @@ async fn detect_ob_oracle_mode(config: &ConnectionConfig, pool: &sqlx::mysql::My #[cfg(test)] mod tests { use super::{database_connection_config, metadata_connection_config, AppState}; - use crate::models::connection::{ConnectionConfig, DatabaseType}; + use crate::models::connection::{ConnectionConfig, DatabaseType, ProxyType}; use crate::schema; use crate::storage::Storage; @@ -439,6 +473,12 @@ mod tests { ssh_key_passphrase: String::new(), ssh_expose_lan: false, ssh_connect_timeout_secs: crate::models::connection::default_ssh_connect_timeout_secs(), + proxy_enabled: false, + proxy_type: ProxyType::Socks5, + proxy_host: String::new(), + proxy_port: 1080, + proxy_username: String::new(), + proxy_password: String::new(), ssl: false, sysdba: false, connection_string: None, @@ -517,4 +557,20 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + + #[tokio::test] + async fn proxy_connection_uses_local_forward_endpoint() { + let (state, dir) = test_app_state().await; + let mut config = mysql_config(Some("app")); + config.proxy_enabled = true; + config.proxy_host = "127.0.0.1".to_string(); + config.proxy_port = 65000; + + let (host, port) = state.connection_host_port("proxied", &config).await.unwrap(); + + assert_eq!(host, "127.0.0.1"); + assert_ne!(port, config.port); + state.proxy_tunnels.stop_tunnel("proxied").await; + let _ = std::fs::remove_dir_all(dir); + } } diff --git a/crates/dbx-core/src/connection_secrets.rs b/crates/dbx-core/src/connection_secrets.rs index 2bf3c1a68d..0ef74783d3 100644 --- a/crates/dbx-core/src/connection_secrets.rs +++ b/crates/dbx-core/src/connection_secrets.rs @@ -6,6 +6,7 @@ use std::path::{Path, PathBuf}; pub const MAIN_PASSWORD_KEY: &str = "password"; pub const SSH_PASSWORD_KEY: &str = "ssh_password"; pub const SSH_KEY_PASSPHRASE_KEY: &str = "ssh_key_passphrase"; +pub const PROXY_PASSWORD_KEY: &str = "proxy_password"; pub const CONNECTION_STRING_KEY: &str = "connection_string"; pub trait ConnectionSecretStore { @@ -61,6 +62,7 @@ pub fn save_connections_to_file( persist_secret(store, &config.id, MAIN_PASSWORD_KEY, &config.password)?; persist_secret(store, &config.id, SSH_PASSWORD_KEY, &config.ssh_password)?; persist_secret(store, &config.id, SSH_KEY_PASSPHRASE_KEY, &config.ssh_key_passphrase)?; + persist_secret(store, &config.id, PROXY_PASSWORD_KEY, &config.proxy_password)?; persist_optional_secret(store, &config.id, CONNECTION_STRING_KEY, config.connection_string.as_deref())?; } @@ -105,6 +107,15 @@ pub fn load_connections_from_file( needs_rewrite = true; } + if config.proxy_password.is_empty() { + if let Some(secret) = store.get_secret(&config.id, PROXY_PASSWORD_KEY)? { + config.proxy_password = secret; + } + } else { + store.set_secret(&config.id, PROXY_PASSWORD_KEY, &config.proxy_password)?; + needs_rewrite = true; + } + match config.connection_string.as_deref().filter(|secret| !secret.is_empty()) { Some(secret) => { store.set_secret(&config.id, CONNECTION_STRING_KEY, secret)?; @@ -195,6 +206,7 @@ fn sanitize_connections(configs: &[ConnectionConfig]) -> Vec { config.password.clear(); config.ssh_password.clear(); config.ssh_key_passphrase.clear(); + config.proxy_password.clear(); config.connection_string = None; config }) @@ -211,7 +223,7 @@ mod tests { load_connections_from_file, save_connections_to_file, ConnectionSecretStore, CONNECTION_STRING_KEY, MAIN_PASSWORD_KEY, SSH_PASSWORD_KEY, }; - use crate::models::connection::{ConnectionConfig, DatabaseType}; + use crate::models::connection::{ConnectionConfig, DatabaseType, ProxyType}; use std::cell::RefCell; use std::collections::HashMap; use std::path::Path; @@ -286,6 +298,12 @@ mod tests { ssh_key_passphrase: String::new(), ssh_expose_lan: false, ssh_connect_timeout_secs: crate::models::connection::default_ssh_connect_timeout_secs(), + proxy_enabled: false, + proxy_type: ProxyType::Socks5, + proxy_host: String::new(), + proxy_port: 1080, + proxy_username: String::new(), + proxy_password: String::new(), ssl: false, sysdba: false, connection_string: None, diff --git a/crates/dbx-core/src/db/mod.rs b/crates/dbx-core/src/db/mod.rs index c7dbdc5746..2ef320001b 100644 --- a/crates/dbx-core/src/db/mod.rs +++ b/crates/dbx-core/src/db/mod.rs @@ -9,6 +9,7 @@ pub mod mysql; pub mod ob_oracle; pub mod oracle_driver; pub mod postgres; +pub mod proxy_tunnel; pub mod redis_driver; pub mod sqlite; pub mod sqlserver; diff --git a/crates/dbx-core/src/db/oracle_driver.rs b/crates/dbx-core/src/db/oracle_driver.rs index d4006f8988..363843a5af 100644 --- a/crates/dbx-core/src/db/oracle_driver.rs +++ b/crates/dbx-core/src/db/oracle_driver.rs @@ -332,10 +332,12 @@ pub async fn list_triggers(conn: &OracleClient, schema: &str, table: &str) -> Re pub async fn execute_query_with_schema(conn: &OracleClient, schema: &str, sql: &str) -> Result { let set_schema = format!("ALTER SESSION SET CURRENT_SCHEMA = \"{}\"", schema); + log::info!("[oracle][set-schema:start] schema={schema}"); conn.execute(&set_schema, &[]).await.map_err(|e| { log::error!("[oracle] set current_schema failed: {e}"); e.to_string() })?; + log::info!("[oracle][set-schema:done] schema={schema}"); execute_query(conn, sql).await } @@ -343,17 +345,31 @@ pub async fn execute_query(conn: &OracleClient, sql: &str) -> Result = result.columns.iter().map(|c| c.name.clone()).collect(); let mut rows: Vec> = result .rows @@ -369,11 +385,24 @@ pub async fn execute_query(conn: &OracleClient, sql: &str) -> Result { let _ = conn.commit().await; + log::info!( + "[oracle][execute-non-select:done] affected_rows={} elapsed_ms={}", + result.rows_affected, + start.elapsed().as_millis() + ); Ok(QueryResult { columns: vec![], rows: vec![], @@ -384,8 +413,9 @@ pub async fn execute_query(conn: &OracleClient, sql: &str) -> Result { let msg = e.to_string(); + log::error!("[oracle][execute-non-select:error] {msg}"); if msg.contains("Server rejected") || msg.contains("closed the connection") { - Err("Operation failed (connection closed) — possibly a constraint violation (foreign key, unique, or check constraint).".to_string()) + Err(format!("Operation failed (connection closed). Original driver error: {msg}")) } else { Err(msg) } diff --git a/crates/dbx-core/src/db/proxy_tunnel.rs b/crates/dbx-core/src/db/proxy_tunnel.rs new file mode 100644 index 0000000000..2328f7305c --- /dev/null +++ b/crates/dbx-core/src/db/proxy_tunnel.rs @@ -0,0 +1,212 @@ +use crate::models::connection::ProxyType; +use base64::engine::general_purpose::STANDARD as BASE64; +use base64::Engine; +use std::collections::HashMap; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::task::JoinHandle; +use tokio::time::{timeout, Duration}; + +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Default)] +pub struct ProxyTunnelManager { + tunnels: tokio::sync::Mutex, u16)>>, +} + +impl ProxyTunnelManager { + pub fn new() -> Self { + Self { tunnels: tokio::sync::Mutex::new(HashMap::new()) } + } + + pub async fn start_tunnel( + &self, + connection_id: &str, + proxy_type: ProxyType, + proxy_host: &str, + proxy_port: u16, + proxy_username: &str, + proxy_password: &str, + remote_host: &str, + remote_port: u16, + ) -> Result { + let local_port = portpicker::pick_unused_port().ok_or("No available port")?; + let listener = TcpListener::bind(("127.0.0.1", local_port)) + .await + .map_err(|e| format!("Failed to bind proxy tunnel local port: {e}"))?; + + let proxy = ProxyEndpoint { + proxy_type, + host: proxy_host.to_string(), + port: proxy_port, + username: proxy_username.to_string(), + password: proxy_password.to_string(), + }; + let remote = RemoteEndpoint { host: remote_host.to_string(), port: remote_port }; + let handle = tokio::spawn(proxy_forward_loop(listener, proxy, remote)); + self.tunnels.lock().await.insert(connection_id.to_string(), (handle, local_port)); + Ok(local_port) + } + + pub async fn local_port(&self, connection_id: &str) -> Option { + self.tunnels.lock().await.get(connection_id).map(|(_, port)| *port) + } + + pub async fn stop_tunnel(&self, connection_id: &str) { + if let Some((handle, _)) = self.tunnels.lock().await.remove(connection_id) { + handle.abort(); + } + } +} + +#[derive(Clone)] +struct ProxyEndpoint { + proxy_type: ProxyType, + host: String, + port: u16, + username: String, + password: String, +} + +#[derive(Clone)] +struct RemoteEndpoint { + host: String, + port: u16, +} + +async fn proxy_forward_loop(listener: TcpListener, proxy: ProxyEndpoint, remote: RemoteEndpoint) { + loop { + let (mut inbound, _) = match listener.accept().await { + Ok(pair) => pair, + Err(_) => break, + }; + let proxy = proxy.clone(); + let remote = remote.clone(); + tokio::spawn(async move { + let Ok(mut outbound) = connect_via_proxy(&proxy, &remote).await else { + return; + }; + let _ = tokio::io::copy_bidirectional(&mut inbound, &mut outbound).await; + }); + } +} + +async fn connect_via_proxy(proxy: &ProxyEndpoint, remote: &RemoteEndpoint) -> Result { + let stream = timeout(CONNECT_TIMEOUT, TcpStream::connect((proxy.host.as_str(), proxy.port))) + .await + .map_err(|_| "Proxy connection timed out".to_string())? + .map_err(|e| format!("Failed to connect proxy: {e}"))?; + + match proxy.proxy_type { + ProxyType::Http => http_connect(stream, proxy, remote).await, + ProxyType::Socks5 => socks5_connect(stream, proxy, remote).await, + } +} + +async fn http_connect( + mut stream: TcpStream, + proxy: &ProxyEndpoint, + remote: &RemoteEndpoint, +) -> Result { + let target = format!("{}:{}", remote.host, remote.port); + let mut request = format!("CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n"); + if !proxy.username.is_empty() || !proxy.password.is_empty() { + let token = BASE64.encode(format!("{}:{}", proxy.username, proxy.password)); + request.push_str(&format!("Proxy-Authorization: Basic {token}\r\n")); + } + request.push_str("\r\n"); + stream.write_all(request.as_bytes()).await.map_err(|e| format!("Failed to send CONNECT request: {e}"))?; + + let mut response = Vec::new(); + let mut buf = [0_u8; 1]; + while !response.ends_with(b"\r\n\r\n") && response.len() < 8192 { + let n = stream.read(&mut buf).await.map_err(|e| format!("Failed to read CONNECT response: {e}"))?; + if n == 0 { + break; + } + response.push(buf[0]); + } + let text = String::from_utf8_lossy(&response); + if text.starts_with("HTTP/1.1 200") || text.starts_with("HTTP/1.0 200") { + Ok(stream) + } else { + let status = text.lines().next().unwrap_or("invalid proxy response"); + Err(format!("HTTP proxy CONNECT failed: {status}")) + } +} + +async fn socks5_connect( + mut stream: TcpStream, + proxy: &ProxyEndpoint, + remote: &RemoteEndpoint, +) -> Result { + let wants_auth = !proxy.username.is_empty() || !proxy.password.is_empty(); + let methods: &[u8] = if wants_auth { &[0x00, 0x02] } else { &[0x00] }; + let mut hello = vec![0x05, methods.len() as u8]; + hello.extend_from_slice(methods); + stream.write_all(&hello).await.map_err(|e| format!("Failed to send SOCKS greeting: {e}"))?; + + let mut method = [0_u8; 2]; + stream.read_exact(&mut method).await.map_err(|e| format!("Failed to read SOCKS greeting: {e}"))?; + if method[0] != 0x05 { + return Err("Invalid SOCKS proxy version".to_string()); + } + match method[1] { + 0x00 => {} + 0x02 => socks5_authenticate(&mut stream, proxy).await?, + 0xff => return Err("SOCKS proxy rejected supported authentication methods".to_string()), + other => return Err(format!("SOCKS proxy selected unsupported auth method: {other}")), + } + + let host = remote.host.as_bytes(); + if host.len() > u8::MAX as usize { + return Err("Remote host is too long for SOCKS5 domain address".to_string()); + } + let mut req = vec![0x05, 0x01, 0x00, 0x03, host.len() as u8]; + req.extend_from_slice(host); + req.extend_from_slice(&remote.port.to_be_bytes()); + stream.write_all(&req).await.map_err(|e| format!("Failed to send SOCKS connect request: {e}"))?; + + let mut head = [0_u8; 4]; + stream.read_exact(&mut head).await.map_err(|e| format!("Failed to read SOCKS connect response: {e}"))?; + if head[0] != 0x05 { + return Err("Invalid SOCKS connect response version".to_string()); + } + if head[1] != 0x00 { + return Err(format!("SOCKS proxy connect failed with code {}", head[1])); + } + let addr_len = match head[3] { + 0x01 => 4, + 0x03 => { + let mut len = [0_u8; 1]; + stream.read_exact(&mut len).await.map_err(|e| format!("Failed to read SOCKS bound address length: {e}"))?; + len[0] as usize + } + 0x04 => 16, + other => return Err(format!("Unsupported SOCKS bound address type: {other}")), + }; + let mut discard = vec![0_u8; addr_len + 2]; + stream.read_exact(&mut discard).await.map_err(|e| format!("Failed to read SOCKS bound address: {e}"))?; + Ok(stream) +} + +async fn socks5_authenticate(stream: &mut TcpStream, proxy: &ProxyEndpoint) -> Result<(), String> { + let username = proxy.username.as_bytes(); + let password = proxy.password.as_bytes(); + if username.len() > u8::MAX as usize || password.len() > u8::MAX as usize { + return Err("SOCKS username or password is too long".to_string()); + } + let mut req = vec![0x01, username.len() as u8]; + req.extend_from_slice(username); + req.push(password.len() as u8); + req.extend_from_slice(password); + stream.write_all(&req).await.map_err(|e| format!("Failed to send SOCKS authentication: {e}"))?; + + let mut res = [0_u8; 2]; + stream.read_exact(&mut res).await.map_err(|e| format!("Failed to read SOCKS authentication response: {e}"))?; + if res == [0x01, 0x00] { + Ok(()) + } else { + Err("SOCKS proxy authentication failed".to_string()) + } +} diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index 5315673246..3c7f5d8ea3 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -38,6 +38,18 @@ pub struct ConnectionConfig { #[serde(default = "default_ssh_connect_timeout_secs")] pub ssh_connect_timeout_secs: u64, #[serde(default)] + pub proxy_enabled: bool, + #[serde(default)] + pub proxy_type: ProxyType, + #[serde(default)] + pub proxy_host: String, + #[serde(default = "default_proxy_port")] + pub proxy_port: u16, + #[serde(default)] + pub proxy_username: String, + #[serde(default)] + pub proxy_password: String, + #[serde(default)] pub ssl: bool, #[serde(default)] pub sysdba: bool, @@ -60,6 +72,23 @@ pub fn default_ssh_connect_timeout_secs() -> u64 { 5 } +fn default_proxy_port() -> u16 { + 1080 +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum ProxyType { + Socks5, + Http, +} + +impl Default for ProxyType { + fn default() -> Self { + Self::Socks5 + } +} + #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "lowercase")] pub enum DatabaseType { @@ -372,7 +401,7 @@ fn bracket_ipv6(host: &str) -> String { #[cfg(test)] mod tests { - use super::{default_ssh_connect_timeout_secs, ConnectionConfig, DatabaseType}; + use super::{default_ssh_connect_timeout_secs, ConnectionConfig, DatabaseType, ProxyType}; fn mysql_config(username: &str, password: &str, database: Option<&str>) -> ConnectionConfig { ConnectionConfig { @@ -397,6 +426,12 @@ mod tests { ssh_key_passphrase: String::new(), ssh_expose_lan: false, ssh_connect_timeout_secs: default_ssh_connect_timeout_secs(), + proxy_enabled: false, + proxy_type: ProxyType::Socks5, + proxy_host: String::new(), + proxy_port: 1080, + proxy_username: String::new(), + proxy_password: String::new(), ssl: false, sysdba: false, connection_string: None, @@ -431,6 +466,28 @@ mod tests { assert_eq!(config.effective_ssh_connect_timeout_secs(), default_ssh_connect_timeout_secs()); } + #[test] + fn proxy_fields_default_for_legacy_config() { + let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ + "id": "id", + "name": "name", + "db_type": "mysql", + "host": "10.1.2.3", + "port": 3306, + "username": "root", + "password": "", + "database": null + })) + .unwrap(); + + assert_eq!(config.proxy_enabled, false); + assert_eq!(config.proxy_type, ProxyType::Socks5); + assert_eq!(config.proxy_host, ""); + assert_eq!(config.proxy_port, 1080); + assert_eq!(config.proxy_username, ""); + assert_eq!(config.proxy_password, ""); + } + #[test] fn ssh_connect_timeout_zero_uses_default() { let mut config = mysql_config("root", "", None); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index b500aab333..cc819cb271 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -205,6 +205,7 @@ pub async fn do_execute( let client = pool.client(); let schema = schema.map(|s| s.to_string()); drop(connections); + log::info!("[query][oracle:lock:start] schema={:?} sql={}", schema, sql); let client = match cancel_token.as_ref() { Some(token) => tokio::select! { biased; @@ -213,6 +214,7 @@ pub async fn do_execute( }, None => client.lock().await, }; + log::info!("[query][oracle:lock:done] schema={:?}", schema); if let Some(schema) = schema { wait_for_query(cancel_token, db::oracle_driver::execute_query_with_schema(&*client, &schema, sql)) .await @@ -635,9 +637,11 @@ async fn exec_tx_none_inner( ) -> Result { let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { + log::info!("[query][tx-none:statement:start] index={} sql={}", i + 1, sql); match do_execute(state, pool_key, sql, schema, None).await { Ok(result) => { total_affected += result.affected_rows; + log::info!("[query][tx-none:statement:done] index={} affected_rows={}", i + 1, result.affected_rows); } Err(e) => { log::warn!("Statement {} failed (no transaction support): {}", i + 1, e); diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index f567f46595..27d5887371 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -397,6 +397,7 @@ impl Storage { sanitized.password = String::new(); sanitized.ssh_password = String::new(); sanitized.ssh_key_passphrase = String::new(); + sanitized.proxy_password = String::new(); sanitized.connection_string = None; let json = serde_json::to_string(&sanitized).map_err(|e| e.to_string())?; @@ -411,6 +412,7 @@ impl Storage { persist_secret_in_tx(&mut tx, &config.id, "password", &config.password).await?; persist_secret_in_tx(&mut tx, &config.id, "ssh_password", &config.ssh_password).await?; persist_secret_in_tx(&mut tx, &config.id, "ssh_key_passphrase", &config.ssh_key_passphrase).await?; + persist_secret_in_tx(&mut tx, &config.id, "proxy_password", &config.proxy_password).await?; if let Some(cs) = &config.connection_string { persist_secret_in_tx(&mut tx, &config.id, "connection_string", cs).await?; } else { @@ -452,6 +454,7 @@ impl Storage { config.password = self.get_secret(&id, "password").await?.unwrap_or_default(); config.ssh_password = self.get_secret(&id, "ssh_password").await?.unwrap_or_default(); config.ssh_key_passphrase = self.get_secret(&id, "ssh_key_passphrase").await?.unwrap_or_default(); + config.proxy_password = self.get_secret(&id, "proxy_password").await?.unwrap_or_default(); config.connection_string = self.get_secret(&id, "connection_string").await?; configs.push(config); } diff --git a/mcp/src/connections.ts b/mcp/src/connections.ts index c9640aa538..969de48efa 100644 --- a/mcp/src/connections.ts +++ b/mcp/src/connections.ts @@ -15,6 +15,12 @@ export interface ConnectionConfig { database?: string; url_params?: string; ssh_enabled: boolean; + proxy_enabled?: boolean; + proxy_type?: "socks5" | "http"; + proxy_host?: string; + proxy_port?: number; + proxy_username?: string; + proxy_password?: string; ssl: boolean; } @@ -50,6 +56,7 @@ export async function loadConnections(): Promise { const config: ConnectionConfig = JSON.parse(row.config_json); config.id = row.id; if (!config.password) config.password = getSecret(db, row.id, "password"); + if (!config.proxy_password) config.proxy_password = getSecret(db, row.id, "proxy_password"); configs.push(config); } @@ -90,6 +97,12 @@ export async function addConnection(config: Omit): Promi ssh_key_path: "", ssh_key_passphrase: "", ssh_expose_lan: false, + proxy_enabled: config.proxy_enabled ?? false, + proxy_type: config.proxy_type ?? "socks5", + proxy_host: config.proxy_host ?? "", + proxy_port: config.proxy_port ?? 1080, + proxy_username: config.proxy_username ?? "", + proxy_password: "", ssl: config.ssl ?? false, sysdba: false, connection_string: null, @@ -101,6 +114,9 @@ export async function addConnection(config: Omit): Promi if (config.password) { db.prepare("INSERT INTO connection_secrets (connection_id, key, secret) VALUES (?, ?, ?)").run(id, "password", config.password); } + if (config.proxy_password) { + db.prepare("INSERT INTO connection_secrets (connection_id, key, secret) VALUES (?, ?, ?)").run(id, "proxy_password", config.proxy_password); + } }); insert(); db.close(); diff --git a/mcp/src/database.ts b/mcp/src/database.ts index 967154da39..39ed5fa436 100644 --- a/mcp/src/database.ts +++ b/mcp/src/database.ts @@ -1,4 +1,5 @@ import type { ConnectionConfig } from "./connections.js"; +import { createServer, connect as netConnect, type Server, type Socket } from "node:net"; export interface TableInfo { name: string; @@ -31,6 +32,7 @@ interface PoolEntry { } const pools = new Map(); +const proxyTunnels = new Map(); function poolKey(config: ConnectionConfig): string { return `${config.id}:${config.database || ""}`; @@ -59,8 +61,9 @@ async function getPgPool(config: ConnectionConfig): Promise { } const pg = await import("pg"); + const endpoint = await connectionEndpoint(config); const pool = new pg.default.Pool({ - connectionString: buildConnectionUrl(config), + connectionString: buildConnectionUrl(config, endpoint), max: 3, idleTimeoutMillis: 30_000, connectionTimeoutMillis: 10_000, @@ -81,8 +84,9 @@ async function getMysqlPool(config: ConnectionConfig): Promise { + if (!config.proxy_enabled || !config.proxy_host) return { host: config.host, port: config.port }; + const existing = proxyTunnels.get(config.id); + if (existing) return { host: "127.0.0.1", port: existing.port }; + + const server = createServer((inbound) => { + connectViaProxy(config) + .then((outbound) => { + inbound.pipe(outbound); + outbound.pipe(inbound); + }) + .catch(() => inbound.destroy()); + }); + const port = await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + if (address && typeof address === "object") resolve(address.port); + else reject(new Error("Failed to bind proxy tunnel")); + }); + }); + proxyTunnels.set(config.id, { server, port }); + return { host: "127.0.0.1", port }; +} + +function buildConnectionUrl(config: ConnectionConfig, endpoint: { host: string; port: number }): string { const db = config.database || ""; const params = config.url_params || ""; const suffix = params ? `?${params}` : ""; if (isMysqlType(config.db_type)) { - return `mysql://${encodeURIComponent(config.username)}:${encodeURIComponent(config.password)}@${config.host}:${config.port}/${db}${suffix}`; + return `mysql://${encodeURIComponent(config.username)}:${encodeURIComponent(config.password)}@${endpoint.host}:${endpoint.port}/${db}${suffix}`; } - return `postgres://${encodeURIComponent(config.username)}:${encodeURIComponent(config.password)}@${config.host}:${config.port}/${db}${suffix}`; + return `postgres://${encodeURIComponent(config.username)}:${encodeURIComponent(config.password)}@${endpoint.host}:${endpoint.port}/${db}${suffix}`; +} + +function connectViaProxy(config: ConnectionConfig): Promise { + return new Promise((resolve, reject) => { + const socket = netConnect(config.proxy_port || 1080, config.proxy_host || "127.0.0.1"); + socket.once("error", reject); + socket.once("connect", () => { + if ((config.proxy_type || "socks5") === "http") { + httpConnect(socket, config, resolve, reject); + } else { + socks5Connect(socket, config, resolve, reject); + } + }); + }); +} + +function httpConnect(socket: Socket, config: ConnectionConfig, resolve: (socket: Socket) => void, reject: (err: Error) => void) { + const target = `${config.host}:${config.port}`; + const lines = [`CONNECT ${target} HTTP/1.1`, `Host: ${target}`]; + if (config.proxy_username || config.proxy_password) { + const token = Buffer.from(`${config.proxy_username || ""}:${config.proxy_password || ""}`).toString("base64"); + lines.push(`Proxy-Authorization: Basic ${token}`); + } + socket.write(`${lines.join("\r\n")}\r\n\r\n`); + let buffer = Buffer.alloc(0); + socket.on("data", function onData(chunk: Buffer) { + buffer = Buffer.concat([buffer, chunk]); + const end = buffer.indexOf("\r\n\r\n"); + if (end < 0) return; + socket.off("data", onData); + const head = buffer.subarray(0, end).toString("utf8"); + if (!/^HTTP\/1\.[01] 200\b/.test(head)) { + reject(new Error(`HTTP proxy CONNECT failed: ${head.split("\r\n")[0] || "invalid response"}`)); + socket.destroy(); + return; + } + const rest = buffer.subarray(end + 4); + if (rest.length) socket.unshift(rest); + resolve(socket); + }); +} + +function socks5Connect(socket: Socket, config: ConnectionConfig, resolve: (socket: Socket) => void, reject: (err: Error) => void) { + const wantsAuth = !!(config.proxy_username || config.proxy_password); + socket.write(Buffer.from(wantsAuth ? [0x05, 0x02, 0x00, 0x02] : [0x05, 0x01, 0x00])); + socket.once("data", (method) => { + if (method.length < 2 || method[0] !== 0x05) { + reject(new Error("Invalid SOCKS greeting")); + socket.destroy(); + return; + } + if (method[1] === 0x02) { + const user = Buffer.from(config.proxy_username || ""); + const pass = Buffer.from(config.proxy_password || ""); + socket.write(Buffer.concat([Buffer.from([0x01, user.length]), user, Buffer.from([pass.length]), pass])); + socket.once("data", (auth) => { + if (auth.length < 2 || auth[1] !== 0x00) { + reject(new Error("SOCKS proxy authentication failed")); + socket.destroy(); + return; + } + sendSocksConnect(socket, config, resolve, reject); + }); + } else if (method[1] === 0x00) { + sendSocksConnect(socket, config, resolve, reject); + } else { + reject(new Error("SOCKS proxy rejected authentication methods")); + socket.destroy(); + } + }); +} + +function sendSocksConnect(socket: Socket, config: ConnectionConfig, resolve: (socket: Socket) => void, reject: (err: Error) => void) { + const host = Buffer.from(config.host); + socket.write(Buffer.concat([Buffer.from([0x05, 0x01, 0x00, 0x03, host.length]), host, portBytes(config.port)])); + socket.once("data", (res) => { + if (res.length < 4 || res[0] !== 0x05 || res[1] !== 0x00) { + reject(new Error(`SOCKS proxy connect failed with code ${res[1] ?? "unknown"}`)); + socket.destroy(); + return; + } + resolve(socket); + }); +} + +function portBytes(port: number): Buffer { + const buf = Buffer.alloc(2); + buf.writeUInt16BE(port); + return buf; } function isMysqlType(dbType: string): boolean { diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index a33cad3816..4df5915119 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -137,6 +137,9 @@ pub async fn test_connection(state: State<'_, Arc>, config: Connection if config.ssh_enabled && !config.ssh_host.is_empty() { state.tunnels.stop_tunnel(&tunnel_id).await; } + if config.proxy_enabled && !config.proxy_host.is_empty() { + state.proxy_tunnels.stop_tunnel(&tunnel_id).await; + } result } @@ -268,5 +271,6 @@ pub async fn disconnect_db(state: State<'_, Arc>, connection_id: Strin drop(conns); state.configs.write().await.remove(&connection_id); state.tunnels.stop_tunnel(&connection_id).await; + state.proxy_tunnels.stop_tunnel(&connection_id).await; Ok(()) } diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index 387d6a5873..b45ca02fd6 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -37,8 +37,29 @@ pub async fn execute_multi( let registered_query = execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); let cancel_token = registered_query.as_ref().map(|query| query.token()); + let trace_id = execution_id.as_deref().unwrap_or("no-execution-id"); + log::info!( + "[query][execute_multi:start] trace_id={} connection_id={} database={} schema={:?} sql={}", + trace_id, + connection_id, + database, + schema, + sql + ); - dbx_core::query::execute_multi_core(&state, &connection_id, &database, &sql, schema.as_deref(), cancel_token).await + let result = + dbx_core::query::execute_multi_core(&state, &connection_id, &database, &sql, schema.as_deref(), cancel_token) + .await; + match &result { + Ok(results) => log::info!( + "[query][execute_multi:done] trace_id={} result_count={} row_counts={:?}", + trace_id, + results.len(), + results.iter().map(|result| result.rows.len()).collect::>() + ), + Err(error) => log::error!("[query][execute_multi:error] trace_id={} error={}", trace_id, error), + } + result } #[tauri::command] diff --git a/src-web/src/routes/connection.rs b/src-web/src/routes/connection.rs index c1bf0bbcec..ac8a4001b5 100644 --- a/src-web/src/routes/connection.rs +++ b/src-web/src/routes/connection.rs @@ -88,6 +88,7 @@ pub async fn disconnect_db( app.configs.write().await.remove(&body.connection_id); app.tunnels.stop_tunnel(&body.connection_id).await; + app.proxy_tunnels.stop_tunnel(&body.connection_id).await; Ok(Json(())) } @@ -121,7 +122,7 @@ mod tests { use axum::extract::State; use axum::Json; use dbx_core::connection::{AppState, PoolKind}; - use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; + use dbx_core::models::connection::{ConnectionConfig, DatabaseType, ProxyType}; use dbx_core::storage::Storage; use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -150,6 +151,12 @@ mod tests { ssh_key_passphrase: String::new(), ssh_expose_lan: false, ssh_connect_timeout_secs: dbx_core::models::connection::default_ssh_connect_timeout_secs(), + proxy_enabled: false, + proxy_type: ProxyType::Socks5, + proxy_host: String::new(), + proxy_port: 1080, + proxy_username: String::new(), + proxy_password: String::new(), ssl: false, sysdba: false, connection_string: None, diff --git a/src/components/connection/ConnectionDialog.vue b/src/components/connection/ConnectionDialog.vue index cd9d7af39c..dcad24c26d 100644 --- a/src/components/connection/ConnectionDialog.vue +++ b/src/components/connection/ConnectionDialog.vue @@ -22,7 +22,7 @@ type DbOption = { value: string; label: string }; type DbCategory = { key: string; title: string; options: DbOption[] }; type DialogStep = "select" | "config"; type DbPickerView = "icon" | "list"; -type ConfigTab = "connection" | "ssh"; +type ConfigTab = "connection" | "ssh" | "proxy"; const { t } = useI18n(); const { toast } = useToast(); @@ -67,6 +67,12 @@ const defaultForm = (): Omit => ({ ssh_key_passphrase: "", ssh_expose_lan: false, ssh_connect_timeout_secs: 5, + proxy_enabled: false, + proxy_type: "socks5", + proxy_host: "", + proxy_port: 1080, + proxy_username: "", + proxy_password: "", ssl: false, connection_string: undefined, jdbc_driver_class: undefined, @@ -264,6 +270,12 @@ watch( ssh_key_passphrase: config.ssh_key_passphrase || "", ssh_expose_lan: config.ssh_expose_lan || false, ssh_connect_timeout_secs: config.ssh_connect_timeout_secs || 5, + proxy_enabled: config.proxy_enabled || false, + proxy_type: config.proxy_type || "socks5", + proxy_host: config.proxy_host || "", + proxy_port: config.proxy_port || 1080, + proxy_username: config.proxy_username || "", + proxy_password: config.proxy_password || "", ssl: config.ssl || false, connection_string: config.connection_string, jdbc_driver_class: config.jdbc_driver_class, @@ -408,6 +420,9 @@ const hasDbPickerResults = computed(() => filteredDbCategories.value.some((categ const selectedDbIcon = computed(() => iconTypeMap[selectedType.value] || selectedProfile().icon || selectedType.value); const isJdbcConnection = computed(() => form.value.db_type === "jdbc"); const canUseSsh = computed(() => form.value.db_type !== "sqlite" && form.value.db_type !== "jdbc"); +const canUseProxy = computed( + () => form.value.db_type !== "sqlite" && form.value.db_type !== "duckdb" && form.value.db_type !== "jdbc", +); const testResultMessage = computed(() => { if (!testResult.value) return ""; return testResult.value.ok ? t("connection.testSuccess") : testResult.value.message; @@ -456,6 +471,8 @@ function connectionConfigForSubmit(id: string): ConnectionConfig { const config: ConnectionConfig = { ...form.value, id }; const sshTimeout = Number(config.ssh_connect_timeout_secs); config.ssh_connect_timeout_secs = Number.isFinite(sshTimeout) && sshTimeout > 0 ? sshTimeout : 5; + const proxyPort = Number(config.proxy_port); + config.proxy_port = Number.isFinite(proxyPort) && proxyPort > 0 ? proxyPort : 1080; if (config.db_type === "mongodb" && !mongoUseUrl.value) { config.connection_string = undefined; } @@ -534,6 +551,12 @@ watch(canUseSsh, (value) => { } }); +watch(canUseProxy, (value) => { + if (!value && configTab.value === "proxy") { + configTab.value = "connection"; + } +}); + async function save() { if (isSaving.value) return; isSaving.value = true; @@ -778,10 +801,11 @@ function openExternalUrl(url: string) {