diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 45eeb7c6d1..4f9de17217 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -4,7 +4,7 @@ use tokio::sync::Mutex; use crate::db; use crate::db::ssh_tunnel::TunnelManager; -use crate::models::connection::{ConnectionConfig, DatabaseType}; +use crate::models::connection::{parse_mongo_first_host, ConnectionConfig, DatabaseType}; use crate::query_cancel::RunningQueries; use crate::storage::Storage; @@ -209,6 +209,17 @@ impl AppState { 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 .tunnels .start_tunnel( @@ -219,8 +230,8 @@ impl AppState { &config.ssh_password, &config.ssh_key_path, &config.ssh_key_passphrase, - &config.host, - config.port, + &remote_host, + remote_port, config.ssh_expose_lan, ) .await?; diff --git a/crates/dbx-core/src/connection_secrets.rs b/crates/dbx-core/src/connection_secrets.rs index 09b2c2f395..b6a65c4b85 100644 --- a/crates/dbx-core/src/connection_secrets.rs +++ b/crates/dbx-core/src/connection_secrets.rs @@ -286,6 +286,7 @@ mod tests { ssh_key_passphrase: String::new(), ssh_expose_lan: false, ssl: false, + sysdba: false, connection_string: None, } } diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index 4d965791e0..9155167bd5 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -122,10 +122,21 @@ impl ConnectionConfig { format!("server=tcp:{host},{port};database={}", self.database.as_deref().unwrap_or("master")) } DatabaseType::MongoDb => { + let is_tunneled = host != self.host.as_str() || port != self.port; if let Some(cs) = self.connection_string.as_deref().filter(|s| !s.is_empty()) { + if is_tunneled { + return rewrite_mongo_uri_host(cs, &host, port); + } return cs.to_string(); } - let suffix = if params.is_empty() { String::new() } else { format!("?{params}") }; + let mut suffix = if params.is_empty() { String::new() } else { format!("?{params}") }; + if is_tunneled && !suffix.contains("directConnection=") { + if suffix.is_empty() { + suffix = "?directConnection=true".to_string(); + } else { + suffix.push_str("&directConnection=true"); + } + } format!("mongodb://{host}:{port}{db_part}{suffix}") } DatabaseType::Oracle => format!("oracle://{host}:{port}{db_part}"), @@ -176,10 +187,21 @@ impl ConnectionConfig { self.database.as_deref().unwrap_or("master") ), DatabaseType::MongoDb => { + let is_tunneled = host != self.host.as_str() || port != self.port; if let Some(cs) = self.connection_string.as_deref().filter(|s| !s.is_empty()) { + if is_tunneled { + return rewrite_mongo_uri_host(cs, &host, port); + } return cs.to_string(); } - let suffix = if params.is_empty() { String::new() } else { format!("?{params}") }; + let mut suffix = if params.is_empty() { String::new() } else { format!("?{params}") }; + if is_tunneled && !suffix.contains("directConnection=") { + if suffix.is_empty() { + suffix = "?directConnection=true".to_string(); + } else { + suffix.push_str("&directConnection=true"); + } + } if self.username.is_empty() { format!("mongodb://{host}:{port}{db_part}{suffix}") } else { @@ -254,6 +276,52 @@ impl ConnectionConfig { } } +pub fn parse_mongo_first_host(uri: &str) -> Option<(String, u16)> { + let rest = uri.strip_prefix("mongodb://").or_else(|| uri.strip_prefix("mongodb+srv://"))?; + let authority = rest.split('/').next()?; + let host_section = match authority.rfind('@') { + Some(idx) => &authority[idx + 1..], + None => authority, + }; + let first = host_section.split(',').next()?; + match first.rsplit_once(':') { + Some((h, p)) => Some((h.to_string(), p.parse().unwrap_or(27017))), + None => Some((first.to_string(), 27017)), + } +} + +fn rewrite_mongo_uri_host(uri: &str, new_host: &str, new_port: u16) -> String { + let (_scheme, rest) = if let Some(r) = uri.strip_prefix("mongodb+srv://") { + ("mongodb://", r) + } else if let Some(r) = uri.strip_prefix("mongodb://") { + ("mongodb://", r) + } else { + return uri.to_string(); + }; + + let (creds_prefix, after_creds) = match rest.find('@') { + Some(idx) => (&rest[..=idx], &rest[idx + 1..]), + None => ("", rest), + }; + + let after_hosts = match after_creds.find('/') { + Some(idx) => &after_creds[idx..], + None => "", + }; + + let mut result = format!("mongodb://{creds_prefix}{new_host}:{new_port}{after_hosts}"); + + if !result.contains("directConnection=") { + if result.contains('?') { + result.push_str("&directConnection=true"); + } else { + result.push_str("?directConnection=true"); + } + } + + result +} + fn encode_url_part(value: &str) -> String { utf8_percent_encode(value, NON_ALPHANUMERIC).to_string() } @@ -411,4 +479,83 @@ mod tests { assert!(!url.contains("root")); assert!(!url.contains("secret")); } + + #[test] + fn parse_mongo_first_host_replica_set() { + let uri = "mongodb://user:pass@host1:27017,host2:27017,host3:27017/admin?replicaSet=rs0"; + let (host, port) = super::parse_mongo_first_host(uri).unwrap(); + assert_eq!(host, "host1"); + assert_eq!(port, 27017); + } + + #[test] + fn parse_mongo_first_host_single() { + let uri = "mongodb://user:pass@myhost:30000/db"; + let (host, port) = super::parse_mongo_first_host(uri).unwrap(); + assert_eq!(host, "myhost"); + assert_eq!(port, 30000); + } + + #[test] + fn parse_mongo_first_host_no_creds() { + let uri = "mongodb://host1:27017,host2:27017/admin"; + let (host, port) = super::parse_mongo_first_host(uri).unwrap(); + assert_eq!(host, "host1"); + assert_eq!(port, 27017); + } + + #[test] + fn parse_mongo_first_host_srv() { + let uri = "mongodb+srv://user:pass@cluster0.example.net/db"; + let (host, port) = super::parse_mongo_first_host(uri).unwrap(); + assert_eq!(host, "cluster0.example.net"); + assert_eq!(port, 27017); + } + + #[test] + fn mongodb_connection_string_rewritten_when_tunneled() { + let mut config = mongodb_config("root", "secret", Some("admin")); + config.connection_string = + Some("mongodb://read:pass@host1:27017,host2:27017/admin?replicaSet=rs0&authSource=admin".to_string()); + + let url = config.connection_url_with_host("127.0.0.1", 54321); + + assert_eq!( + url, + "mongodb://read:pass@127.0.0.1:54321/admin?replicaSet=rs0&authSource=admin&directConnection=true" + ); + } + + #[test] + fn mongodb_connection_string_unchanged_when_not_tunneled() { + let mut config = mongodb_config("root", "secret", Some("admin")); + config.connection_string = Some("mongodb://read:pass@host1:27017,host2:27017/admin?replicaSet=rs0".to_string()); + + let url = config.connection_url(); + + assert_eq!(url, "mongodb://read:pass@host1:27017,host2:27017/admin?replicaSet=rs0"); + } + + #[test] + fn mongodb_form_url_adds_direct_connection_when_tunneled() { + let mut config = mongodb_config("root", "secret", Some("admin")); + config.url_params = Some("replicaSet=rs0&authSource=admin".to_string()); + + let url = config.connection_url_with_host("127.0.0.1", 54321); + + assert_eq!( + url, + "mongodb://root:secret@127.0.0.1:54321/admin?replicaSet=rs0&authSource=admin&directConnection=true" + ); + } + + #[test] + fn mongodb_form_url_no_duplicate_direct_connection() { + let mut config = mongodb_config("root", "secret", Some("admin")); + config.url_params = Some("directConnection=true&authSource=admin".to_string()); + + let url = config.connection_url_with_host("127.0.0.1", 54321); + + assert!(url.matches("directConnection").count() == 1); + } }