fix(jdbc): support encrypted sqlite with managed runtime

This commit is contained in:
t8y2 2026-06-01 17:52:11 +08:00
parent 84a7ff6b92
commit 6da367a47b
6 changed files with 238 additions and 13 deletions

View File

@ -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<String, String> {
let params = serde_json::json!({ "connection": config });
self.plugins.invoke_driver::<serde_json::Value>(driver_id, "testConnection", params).await?;
let env = self.external_driver_runtime_env(driver_id)?;
self.plugins.invoke_driver_with_env::<serde_json::Value>(driver_id, "testConnection", params, env).await?;
Ok("Connection successful".to_string())
}
pub async fn external_driver_pool(&self, driver_id: &str, config: &ConnectionConfig) -> Result<PoolKind, String> {
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::<serde_json::Value>("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<PluginRuntimeEnv, String> {
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<String, String> {
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,

View File

@ -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<String>, value: impl Into<String>) -> 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<T>(&self, driver_id: &str, method: &str, params: serde_json::Value) -> Result<T, String>
where
T: DeserializeOwned,
{
self.invoke_driver_with_env(driver_id, method, params, PluginRuntimeEnv::default()).await
}
pub async fn invoke_driver_with_env<T>(
&self,
driver_id: &str,
method: &str,
params: serde_json::Value,
env: PluginRuntimeEnv,
) -> Result<T, String>
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<Arc<PluginDriverSession>, 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<Arc<PluginDriverSession>, 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<Self, String> {
let mut child = spawn_plugin_child(&plugin)?;
async fn start(plugin: InstalledPlugin, driver_id: String, env: PluginRuntimeEnv) -> Result<Self, String> {
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<T>(
driver_id: &str,
method: &str,
params: serde_json::Value,
env: &PluginRuntimeEnv,
) -> Result<T, String>
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<Vec<u8>, String
Ok(bytes)
}
fn spawn_plugin_child(plugin: &InstalledPlugin) -> Result<Child, String> {
fn spawn_plugin_child(plugin: &InstalledPlugin, env: &PluginRuntimeEnv) -> Result<Child, String> {
let executable = plugin
.manifest
.executable
@ -353,6 +397,7 @@ fn spawn_plugin_child(plugin: &InstalledPlugin) -> Result<Child, String> {
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
env.apply_to(&mut command);
#[cfg(windows)]
{

View File

@ -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"

View File

@ -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"

View File

@ -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();

View File

@ -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", """