diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index b662d6b94..c4bdfa2c9 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -10,6 +10,7 @@ use crate::agent_connection::{ agent_connect_params, mongo_legacy_error_with_auth_hint, oracle_alternate_connect_config, oracle_auth_fallback_profiles, should_retry_oracle_with_10g_driver, }; +use crate::agent_manager::{JavaRuntimeMode, DEFAULT_JRE_KEY}; use crate::database_capabilities; use crate::db; use crate::db::agent_driver::AgentMethod; @@ -19,7 +20,7 @@ use crate::external; use crate::models::connection::{ parse_jdbc_host_port, parse_mongo_first_host, rewrite_jdbc_url_host, ConnectionConfig, DatabaseType, }; -use crate::plugins::{PluginDriverSession, PluginRegistry}; +use crate::plugins::{PluginDriverSession, PluginRegistry, PluginRuntimeEnv}; use crate::query_cancel::RunningQueries; use crate::storage::Storage; @@ -137,17 +138,32 @@ impl AppState { pub async fn test_external_driver(&self, driver_id: &str, config: &ConnectionConfig) -> Result { let params = serde_json::json!({ "connection": config }); - self.plugins.invoke_driver::(driver_id, "testConnection", params).await?; + let env = self.external_driver_runtime_env(driver_id)?; + self.plugins.invoke_driver_with_env::(driver_id, "testConnection", params, env).await?; Ok("Connection successful".to_string()) } pub async fn external_driver_pool(&self, driver_id: &str, config: &ConnectionConfig) -> Result { - let session = self.plugins.start_driver_session(driver_id).await?; + let env = self.external_driver_runtime_env(driver_id)?; + let session = self.plugins.start_driver_session_with_env(driver_id, env).await?; let params = serde_json::json!({ "connection": config }); session.invoke::("connect", params).await?; Ok(PoolKind::ExternalDriver { driver_id: driver_id.to_string(), config: Arc::new(config.clone()), session }) } + fn external_driver_runtime_env(&self, driver_id: &str) -> Result { + if driver_id != "jdbc" { + return Ok(PluginRuntimeEnv::default()); + } + let state = self.agent_manager.load_state(); + if state.java_runtime.mode == JavaRuntimeMode::Managed && !self.agent_manager.is_jre_installed(DEFAULT_JRE_KEY) + { + return Ok(PluginRuntimeEnv::default()); + } + let java = self.agent_manager.resolve_java_runtime(&state, DEFAULT_JRE_KEY)?; + Ok(PluginRuntimeEnv::default().with_var("DBX_JAVA_BIN", java.to_string_lossy().to_string())) + } + pub async fn get_or_create_pool(&self, connection_id: &str, database: Option<&str>) -> Result { self.get_or_create_pool_for_session(connection_id, database, None).await } @@ -911,6 +927,7 @@ mod tests { agent_connect_params, mongo_legacy_error_with_auth_hint, oracle_alternate_connect_config, should_retry_oracle_with_10g_driver, }; + use crate::agent_manager::{AgentState, JavaRuntimeConfig, JavaRuntimeMode, DEFAULT_JRE_KEY}; use crate::db; use crate::models::connection::{default_connect_timeout_secs, ConnectionConfig, DatabaseType, ProxyType}; use crate::schema; @@ -1201,6 +1218,21 @@ mod tests { (AppState::new(storage), dir) } + fn touch_executable(path: &std::path::Path) { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).unwrap(); + } + std::fs::write(path, b"").unwrap(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + + let mut permissions = std::fs::metadata(path).unwrap().permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(path, permissions).unwrap(); + } + } + #[tokio::test] async fn app_state_uses_explicit_agent_dir() { let dir = std::env::temp_dir().join(format!("dbx-core-agent-dir-test-{}", uuid::Uuid::new_v4())); @@ -1219,6 +1251,74 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + #[tokio::test] + async fn jdbc_plugin_env_uses_managed_jre_when_installed() { + let dir = std::env::temp_dir().join(format!("dbx-core-jdbc-managed-jre-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new_with_plugin_and_agent_dir_and_app_version( + storage, + dir.join("plugins"), + dir.join("agents"), + "0.0.0-test", + ); + let java = state.agent_manager.jre_java_path(DEFAULT_JRE_KEY); + touch_executable(&java); + + let env = state.external_driver_runtime_env("jdbc").unwrap(); + + assert_eq!(env.get("DBX_JAVA_BIN"), Some(java.to_string_lossy().as_ref())); + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn jdbc_plugin_env_keeps_wrapper_fallback_when_managed_jre_is_missing() { + let dir = std::env::temp_dir().join(format!("dbx-core-jdbc-missing-jre-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new_with_plugin_and_agent_dir_and_app_version( + storage, + dir.join("plugins"), + dir.join("agents"), + "0.0.0-test", + ); + + let env = state.external_driver_runtime_env("jdbc").unwrap(); + + assert_eq!(env.get("DBX_JAVA_BIN"), None); + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn jdbc_plugin_env_uses_custom_java_runtime() { + let dir = std::env::temp_dir().join(format!("dbx-core-jdbc-custom-jre-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new_with_plugin_and_agent_dir_and_app_version( + storage, + dir.join("plugins"), + dir.join("agents"), + "0.0.0-test", + ); + let java = dir.join("custom").join("bin").join(if cfg!(windows) { "java.exe" } else { "java" }); + touch_executable(&java); + state + .agent_manager + .save_state(&AgentState { + java_runtime: JavaRuntimeConfig { + mode: JavaRuntimeMode::Custom, + custom_java_path: Some(java.to_string_lossy().to_string()), + }, + ..AgentState::default() + }) + .unwrap(); + + let env = state.external_driver_runtime_env("jdbc").unwrap(); + + assert_eq!(env.get("DBX_JAVA_BIN"), Some(java.to_string_lossy().as_ref())); + let _ = std::fs::remove_dir_all(dir); + } + fn live_postgres_like_config( db_type: DatabaseType, host: &str, diff --git a/crates/dbx-core/src/plugins.rs b/crates/dbx-core/src/plugins.rs index 73ccce6fd..95e4d1f71 100644 --- a/crates/dbx-core/src/plugins.rs +++ b/crates/dbx-core/src/plugins.rs @@ -51,6 +51,28 @@ pub struct InstalledPlugin { pub path: PathBuf, } +#[derive(Debug, Clone, Default)] +pub struct PluginRuntimeEnv { + vars: Vec<(String, String)>, +} + +impl PluginRuntimeEnv { + pub fn with_var(mut self, key: impl Into, value: impl Into) -> Self { + self.vars.push((key.into(), value.into())); + self + } + + pub fn get(&self, key: &str) -> Option<&str> { + self.vars.iter().find_map(|(name, value)| (name == key).then_some(value.as_str())) + } + + fn apply_to(&self, command: &mut Command) { + for (key, value) in &self.vars { + command.env(key, value); + } + } +} + #[derive(Debug, Clone)] pub struct PluginRegistry { root_dir: PathBuf, @@ -103,22 +125,43 @@ impl PluginRegistry { } pub async fn invoke_driver(&self, driver_id: &str, method: &str, params: serde_json::Value) -> Result + where + T: DeserializeOwned, + { + self.invoke_driver_with_env(driver_id, method, params, PluginRuntimeEnv::default()).await + } + + pub async fn invoke_driver_with_env( + &self, + driver_id: &str, + method: &str, + params: serde_json::Value, + env: PluginRuntimeEnv, + ) -> Result where T: DeserializeOwned, { let plugin = self.find_driver(driver_id)?.ok_or_else(|| format!("Plugin driver '{driver_id}' is not installed"))?; ensure_plugin_protocol_compatible(&plugin.manifest)?; - timeout(PLUGIN_REQUEST_TIMEOUT, invoke_plugin(&plugin, driver_id, method, params)).await.map_err(|_| { - format!("Plugin '{}' timed out after {} seconds", plugin.manifest.id, PLUGIN_REQUEST_TIMEOUT.as_secs()) - })? + timeout(PLUGIN_REQUEST_TIMEOUT, invoke_plugin(&plugin, driver_id, method, params, &env)).await.map_err( + |_| format!("Plugin '{}' timed out after {} seconds", plugin.manifest.id, PLUGIN_REQUEST_TIMEOUT.as_secs()), + )? } pub async fn start_driver_session(&self, driver_id: &str) -> Result, String> { + self.start_driver_session_with_env(driver_id, PluginRuntimeEnv::default()).await + } + + pub async fn start_driver_session_with_env( + &self, + driver_id: &str, + env: PluginRuntimeEnv, + ) -> Result, String> { let plugin = self.find_driver(driver_id)?.ok_or_else(|| format!("Plugin driver '{driver_id}' is not installed"))?; ensure_plugin_protocol_compatible(&plugin.manifest)?; - PluginDriverSession::start(plugin, driver_id.to_string()).await.map(Arc::new) + PluginDriverSession::start(plugin, driver_id.to_string(), env).await.map(Arc::new) } } @@ -169,8 +212,8 @@ struct PluginProcess { } impl PluginDriverSession { - async fn start(plugin: InstalledPlugin, driver_id: String) -> Result { - let mut child = spawn_plugin_child(&plugin)?; + async fn start(plugin: InstalledPlugin, driver_id: String, env: PluginRuntimeEnv) -> Result { + let mut child = spawn_plugin_child(&plugin, &env)?; let stdin = child.stdin.take().ok_or("Plugin stdin unavailable")?; let stdout = child.stdout.take().ok_or("Plugin stdout unavailable")?; if let Some(stderr) = child.stderr.take() { @@ -275,11 +318,12 @@ async fn invoke_plugin( driver_id: &str, method: &str, params: serde_json::Value, + env: &PluginRuntimeEnv, ) -> Result where T: DeserializeOwned, { - let mut child = spawn_plugin_child(plugin)?; + let mut child = spawn_plugin_child(plugin, env)?; let request = PluginRequest { jsonrpc: "2.0", id: 1, driver: driver_id.to_string(), method: method.to_string(), params }; @@ -337,7 +381,7 @@ fn encode_plugin_request_line(request: &PluginRequest) -> Result, String Ok(bytes) } -fn spawn_plugin_child(plugin: &InstalledPlugin) -> Result { +fn spawn_plugin_child(plugin: &InstalledPlugin, env: &PluginRuntimeEnv) -> Result { let executable = plugin .manifest .executable @@ -353,6 +397,7 @@ fn spawn_plugin_child(plugin: &InstalledPlugin) -> Result { .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(true); + env.apply_to(&mut command); #[cfg(windows)] { diff --git a/plugins/jdbc/bin/dbx-jdbc-plugin b/plugins/jdbc/bin/dbx-jdbc-plugin index e0b203710..a2e1f672c 100755 --- a/plugins/jdbc/bin/dbx-jdbc-plugin +++ b/plugins/jdbc/bin/dbx-jdbc-plugin @@ -2,7 +2,9 @@ set -eu DIR="$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)" -if [ -n "${JAVA_HOME:-}" ] && [ -x "$JAVA_HOME/bin/java" ]; then +if [ -n "${DBX_JAVA_BIN:-}" ] && [ -x "$DBX_JAVA_BIN" ]; then + JAVA_BIN="$DBX_JAVA_BIN" +elif [ -n "${JAVA_HOME:-}" ] && [ -x "$JAVA_HOME/bin/java" ]; then JAVA_BIN="$JAVA_HOME/bin/java" elif [ -x "/opt/homebrew/opt/openjdk/bin/java" ]; then JAVA_BIN="/opt/homebrew/opt/openjdk/bin/java" diff --git a/plugins/jdbc/bin/dbx-jdbc-plugin.bat b/plugins/jdbc/bin/dbx-jdbc-plugin.bat index 4d759ecef..c9672a466 100644 --- a/plugins/jdbc/bin/dbx-jdbc-plugin.bat +++ b/plugins/jdbc/bin/dbx-jdbc-plugin.bat @@ -1,6 +1,13 @@ @echo off setlocal +if defined DBX_JAVA_BIN ( + if exist "%DBX_JAVA_BIN%" ( + set "JAVA_BIN=%DBX_JAVA_BIN%" + goto :run + ) +) + if defined JAVA_HOME ( if exist "%JAVA_HOME%\bin\java.exe" ( set "JAVA_BIN=%JAVA_HOME%\bin\java.exe" diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index d76313586..c169fe222 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -7,6 +7,7 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import java.io.BufferedReader; import java.io.InputStreamReader; +import java.net.URLEncoder; import java.math.BigDecimal; import java.net.URL; import java.net.URLClassLoader; @@ -180,7 +181,7 @@ public final class DbxJdbcPlugin { } private static Connection openConnection(JsonNode connection) throws SQLException { - String url = optionalText(connection, "connection_string"); + String url = jdbcUrlWithPasswordKey(optionalText(connection, "connection_string"), optionalText(connection, "password")); if (url == null) { throw new IllegalArgumentException("JDBC URL is required."); } @@ -547,6 +548,46 @@ public final class DbxJdbcPlugin { return url != null && url.regionMatches(true, 0, "jdbc:oracle:", 0, 12); } + static String jdbcUrlWithPasswordKey(String url, String password) { + if (url == null || password == null || password.isBlank() || !isSqliteUrl(url)) { + return url; + } + if (!urlHasQueryParam(url, "cipher") || urlHasQueryParam(url, "key")) { + return url; + } + return appendJdbcUrlParam(url, "key", password); + } + + private static boolean isSqliteUrl(String url) { + return url.regionMatches(true, 0, "jdbc:sqlite:", 0, 12); + } + + private static boolean urlHasQueryParam(String url, String key) { + int queryStart = url.indexOf('?'); + if (queryStart < 0) { + return false; + } + int fragmentStart = url.indexOf('#', queryStart + 1); + String query = fragmentStart < 0 ? url.substring(queryStart + 1) : url.substring(queryStart + 1, fragmentStart); + for (String part : query.split("[&;]")) { + int equals = part.indexOf('='); + String name = equals < 0 ? part : part.substring(0, equals); + if (name.equalsIgnoreCase(key)) { + return true; + } + } + return false; + } + + private static String appendJdbcUrlParam(String url, String key, String value) { + int fragmentStart = url.indexOf('#'); + String base = fragmentStart < 0 ? url : url.substring(0, fragmentStart); + String fragment = fragmentStart < 0 ? "" : url.substring(fragmentStart); + String separator = base.contains("?") ? (base.endsWith("?") || base.endsWith("&") ? "" : "&") : "?"; + String encodedValue = URLEncoder.encode(value, StandardCharsets.UTF_8); + return base + separator + key + "=" + encodedValue + fragment; + } + private static String oracleEffectiveSchema(Connection conn, String schema) throws SQLException { if (schema != null && !schema.isBlank()) { return schema.toUpperCase(); diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java index 5056f9389..6d60fd258 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -92,6 +92,36 @@ final class DbxJdbcPluginTest { assertEquals(false, DbxJdbcPlugin.driverQuirks(h2).useOracleMetadata()); } + @Test + void sqliteCipherUrlUsesPasswordAsKeyWhenKeyIsMissing() { + String url = DbxJdbcPlugin.jdbcUrlWithPasswordKey( + "jdbc:sqlite:/tmp/library.db?cipher=chacha20", + "my password" + ); + + assertEquals("jdbc:sqlite:/tmp/library.db?cipher=chacha20&key=my+password", url); + } + + @Test + void sqliteCipherUrlKeepsExplicitKey() { + String url = DbxJdbcPlugin.jdbcUrlWithPasswordKey( + "jdbc:sqlite:/tmp/library.db?cipher=chacha20&key=from-url", + "from-password" + ); + + assertEquals("jdbc:sqlite:/tmp/library.db?cipher=chacha20&key=from-url", url); + } + + @Test + void nonSqliteUrlDoesNotUsePasswordAsKey() { + String url = DbxJdbcPlugin.jdbcUrlWithPasswordKey( + "jdbc:h2:mem:dbx_cipher?cipher=sqlcipher", + "secret" + ); + + assertEquals("jdbc:h2:mem:dbx_cipher?cipher=sqlcipher", url); + } + @Test void listTablesFallsBackWhenCatalogFiltersEverything() throws Exception { request("executeQuery", """