feat(agents): add versioned offline driver packages

This commit is contained in:
t8y2 committed 2026-07-21 13:42:39 +08:00
1 parent 8d68a527a0
commit 79cd12d13a
29 files changed
+1372 -162

No files matched your search

+396 -48
View File
@@ -1,5 +1,5 @@
use std::hash::{Hash, Hasher};
use std::io::Read;
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::time::Duration;
@@ -306,11 +306,13 @@ pub fn install_local_agent(am: &AgentManager, db_type: &str, source: PathBuf) ->
let jar_path = am.driver_jar_path(db_type);
let parent = jar_path.parent().ok_or_else(|| format!("Invalid driver path: {}", jar_path.display()))?;
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
std::fs::copy(&source, &jar_path).map_err(|e| format!("Failed to copy local agent jar: {e}"))?;
if !am.is_driver_jar_valid(db_type) {
std::fs::remove_file(&jar_path).ok();
let staging_path = parent.join(format!(".agent-jar-import-{}", uuid::Uuid::new_v4()));
std::fs::copy(&source, &staging_path).map_err(|e| format!("Failed to copy local agent jar: {e}"))?;
if !is_valid_agent_jar(&staging_path) {
std::fs::remove_file(&staging_path).ok();
return Err(format!("Local agent jar is invalid or corrupt: {}", source.display()));
}
replace_imported_agent_file(&staging_path, &jar_path)?;
let mut local_state = am.load_state();
local_state.installed_drivers.insert(
@@ -324,6 +326,20 @@ pub fn install_local_agent(am: &AgentManager, db_type: &str, source: PathBuf) ->
am.save_state(&local_state)
}
fn is_valid_agent_jar(path: &Path) -> bool {
let Ok(file) = std::fs::File::open(path) else {
return false;
};
let Ok(mut archive) = zip::ZipArchive::new(file) else {
return false;
};
let Ok(mut manifest) = archive.by_name("META-INF/MANIFEST.MF") else {
return false;
};
let mut manifest_text = String::new();
manifest.read_to_string(&mut manifest_text).is_ok() && manifest_text.contains("Main-Class:")
}
pub async fn fetch_registry() -> Result<AgentRegistry, String> {
fetch_registry_from(DownloadSource::Official).await
}
@@ -1228,6 +1244,25 @@ pub struct OfflineImportResult {
pub drivers_skipped: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct OfflineImportPlan {
pub driver_keys: Vec<String>,
pub includes_jre: bool,
}
type OfflineDriverEntry = (String, String, bool);
pub fn inspect_offline_zip(zip_path: &Path) -> Result<OfflineImportPlan, String> {
let file = std::fs::File::open(zip_path).map_err(|e| format!("Failed to open ZIP file: {e}"))?;
let mut archive = zip::ZipArchive::new(file).map_err(|e| format!("Invalid ZIP file: {e}"))?;
let registry = read_registry_from_zip(&mut archive)?;
let (jre_entries, driver_entries) = collect_offline_entries(&mut archive, &registry)?;
Ok(OfflineImportPlan {
driver_keys: driver_entries.into_iter().map(|(db_type, _, _)| db_type).collect(),
includes_jre: !jre_entries.is_empty(),
})
}
pub fn import_offline_zip(
am: &AgentManager,
zip_path: &Path,
@@ -1244,36 +1279,13 @@ pub fn import_offline_zip(
let mut result =
OfflineImportResult { jre_installed: Vec::new(), drivers_installed: Vec::new(), drivers_skipped: Vec::new() };
let jre_entries: Vec<(String, String)> = (0..archive.len())
.filter_map(|i| {
let entry = archive.by_index(i).ok()?;
let name = entry.name().to_string();
if name.starts_with("jre/") && name.ends_with(".tar.gz") && name.contains(platform) {
let jre_key = extract_jre_key_from_filename(&name)?;
Some((jre_key, name))
} else {
None
}
})
.collect();
let driver_entries: Vec<(String, String, bool)> = (0..archive.len())
.filter_map(|i| {
let entry = archive.by_index(i).ok()?;
let name = entry.name().to_string();
if name.starts_with("drivers/") && name.ends_with(".jar") {
let db_type = extract_db_type_from_filename(&name)?;
Some((db_type, name, false))
} else if name.starts_with("drivers/") {
let db_type = db_type_for_native_offline_entry(&registry, platform, &name)?;
Some((db_type, name, true))
} else {
None
}
})
.collect();
let (jre_entries, driver_entries) = collect_offline_entries(&mut archive, &registry)?;
let total = (jre_entries.len() + driver_entries.len()) as u32;
if total == 0 {
return Err(format!("Offline package contains no drivers compatible with platform: {platform}"));
}
validate_offline_driver_entries(am, &mut archive, &driver_entries)?;
let mut current: u32 = 0;
for (jre_key, entry_name) in &jre_entries {
@@ -1295,13 +1307,22 @@ pub fn import_offline_zip(
}
let jre_dir = am.jre_dir(jre_key);
// Daemons cannot be stopped from a sync function safely; the retry +
// Windows rename fallback in replace_old_jre_dir still handles a
// locked directory. Daemon shutdown for the foreground install paths
// happens in `reinstall_agent_jre` and `install_agent_driver_*`.
replace_old_jre_dir(am, &jre_dir)?;
extract_tar_gz(&tmp_archive, &jre_dir)?;
let staging_dir = am.base_dir().join(format!(".jre-offline-import-{}", uuid::Uuid::new_v4()));
if let Err(error) = extract_tar_gz(&tmp_archive, &staging_dir) {
std::fs::remove_dir_all(&staging_dir).ok();
std::fs::remove_file(&tmp_archive).ok();
return Err(error);
}
if !jre_dir_contains_java(&staging_dir) {
std::fs::remove_dir_all(&staging_dir).ok();
std::fs::remove_file(&tmp_archive).ok();
return Err(format!("Offline JRE archive does not contain a Java executable: {entry_name}"));
}
let pending_cleanup = replace_imported_jre_dir(&staging_dir, &jre_dir)?;
std::fs::remove_file(&tmp_archive).ok();
if let Some(path) = pending_cleanup {
local_state.pending_jre_cleanup.push(path);
}
if let Some(ver) = jre_version {
local_state.jre_versions.insert(jre_key.clone(), ver);
@@ -1336,17 +1357,29 @@ pub fn import_offline_zip(
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
}
let mut entry = archive.by_name(entry_name).map_err(|e| format!("Failed to read {entry_name}: {e}"))?;
let mut out = std::fs::File::create(&driver_path).map_err(|e| format!("Failed to write driver: {e}"))?;
let parent = driver_path.parent().ok_or_else(|| format!("Invalid driver path: {}", driver_path.display()))?;
let staging_path = parent.join(format!(".offline-agent-import-{}", uuid::Uuid::new_v4()));
let mut out = std::fs::File::create(&staging_path).map_err(|e| format!("Failed to write driver: {e}"))?;
std::io::copy(&mut entry, &mut out).map_err(|e| format!("Failed to copy driver: {e}"))?;
drop(out);
if *is_native {
mark_executable(&driver_path)?;
std::fs::remove_file(am.driver_jar_path(db_type)).ok();
if let Err(error) = validate_native_agent_binary(&staging_path) {
std::fs::remove_file(&staging_path).ok();
return Err(error);
}
mark_executable(&staging_path)?;
} else {
// Offline bundles are user-supplied; reject corrupt JARs before state says the driver is installed.
if !am.is_driver_jar_valid(db_type) {
std::fs::remove_file(&driver_path).ok();
// Validate the staged JAR before replacing a working driver so a
// corrupt offline package cannot destroy the previous installation.
if !is_valid_agent_jar(&staging_path) {
std::fs::remove_file(&staging_path).ok();
return Err(format!("Offline agent jar is invalid or corrupt: {entry_name}"));
}
}
replace_imported_agent_file(&staging_path, &driver_path)?;
if *is_native {
std::fs::remove_file(am.driver_jar_path(db_type)).ok();
} else {
std::fs::remove_file(am.driver_native_path(db_type)).ok();
}
@@ -1365,6 +1398,88 @@ pub fn import_offline_zip(
Ok(result)
}
fn collect_offline_entries(
archive: &mut zip::ZipArchive<std::fs::File>,
registry: &AgentRegistry,
) -> Result<(Vec<(String, String)>, Vec<OfflineDriverEntry>), String> {
let platform = AgentManager::current_platform();
let mut jre_entries = Vec::new();
let mut drivers = std::collections::BTreeMap::<String, (String, bool)>::new();
for index in 0..archive.len() {
let entry = archive.by_index(index).map_err(|e| format!("Failed to inspect ZIP entry: {e}"))?;
let Some(path) = entry.enclosed_name() else {
return Err(format!("Offline package contains an unsafe path: {}", entry.name()));
};
let name = path.to_string_lossy().replace('\\', "/");
if name.starts_with("jre/") && name.ends_with(".tar.gz") && name.contains(platform) {
let jre_key = extract_jre_key_from_filename(&name)
.ok_or_else(|| format!("Invalid JRE filename in offline package: {name}"))?;
validate_offline_identifier(&jre_key, "JRE")?;
jre_entries.push((jre_key, name));
} else if name.starts_with("drivers/") && name.ends_with(".jar") {
let db_type = db_type_for_jar_offline_entry(registry, &name)
.or_else(|| extract_db_type_from_filename(&name))
.ok_or_else(|| format!("Unable to identify offline driver: {name}"))?;
validate_offline_driver_key(&db_type)?;
drivers.entry(db_type).or_insert((name, false));
} else if name.starts_with("drivers/") {
if let Some(db_type) = db_type_for_native_offline_entry(registry, platform, &name) {
validate_offline_driver_key(&db_type)?;
// Prefer the native artifact when a package contains both the
// platform executable and a Java fallback for the same driver.
drivers.insert(db_type, (name, true));
}
}
}
Ok((jre_entries, drivers.into_iter().map(|(db_type, (name, is_native))| (db_type, name, is_native)).collect()))
}
fn validate_offline_driver_entries(
am: &AgentManager,
archive: &mut zip::ZipArchive<std::fs::File>,
driver_entries: &[OfflineDriverEntry],
) -> Result<(), String> {
for (_, entry_name, is_native) in driver_entries {
let staging_path = am.base_dir().join(format!(".offline-agent-validation-{}", uuid::Uuid::new_v4()));
let result = (|| {
let mut entry = archive.by_name(entry_name).map_err(|e| format!("Failed to read {entry_name}: {e}"))?;
let mut out = std::fs::File::create(&staging_path).map_err(|e| format!("Failed to write driver: {e}"))?;
std::io::copy(&mut entry, &mut out).map_err(|e| format!("Failed to copy driver: {e}"))?;
drop(out);
if *is_native {
validate_native_agent_binary(&staging_path)
} else if is_valid_agent_jar(&staging_path) {
Ok(())
} else {
Err(format!("Offline agent jar is invalid or corrupt: {entry_name}"))
}
})();
std::fs::remove_file(&staging_path).ok();
result?;
}
Ok(())
}
fn validate_offline_driver_key(db_type: &str) -> Result<(), String> {
validate_offline_identifier(db_type, "driver")?;
if agent_catalog::label_for_key(db_type).is_none() {
return Err(format!("Offline package contains an unknown driver type: {db_type}"));
}
Ok(())
}
fn validate_offline_identifier(value: &str, kind: &str) -> Result<(), String> {
if value.is_empty()
|| matches!(value, "." | "..")
|| !value.chars().all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '.' | '-' | '_'))
{
return Err(format!("Offline package contains an invalid {kind} identifier: {value}"));
}
Ok(())
}
fn read_registry_from_zip(archive: &mut zip::ZipArchive<std::fs::File>) -> Result<AgentRegistry, String> {
let mut entry = archive
.by_name("agent-registry.json")
@@ -1403,6 +1518,15 @@ fn db_type_for_native_offline_entry(registry: &AgentRegistry, platform: &str, na
})
}
fn db_type_for_jar_offline_entry(registry: &AgentRegistry, name: &str) -> Option<String> {
let filename = name.rsplit('/').next()?;
registry.drivers.iter().find_map(|(db_type, driver)| {
let artifact = driver.jar.as_ref()?;
let artifact_filename = artifact.url.rsplit('/').next()?;
(artifact_filename == filename).then(|| db_type.clone())
})
}
fn extract_tar_gz(archive: &Path, dest: &Path) -> Result<(), String> {
std::fs::create_dir_all(dest).map_err(|e| e.to_string())?;
let status = crate::process::new_std_command("tar")
@@ -1415,11 +1539,210 @@ fn extract_tar_gz(archive: &Path, dest: &Path) -> Result<(), String> {
Ok(())
}
pub fn import_agent_jar(am: &AgentManager, db_type: &str, jar_path: &Path) -> Result<(), String> {
if !jar_path.exists() {
return Err(format!("File not found: {}", jar_path.display()));
pub fn import_agent_driver(am: &AgentManager, db_type: &str, source_path: &Path) -> Result<(), String> {
if !source_path.is_file() {
return Err(format!("File not found: {}", source_path.display()));
}
install_local_agent(am, db_type, jar_path.to_path_buf())
if source_path.extension().is_some_and(|extension| extension.eq_ignore_ascii_case("jar")) {
install_local_agent(am, db_type, source_path.to_path_buf())?;
std::fs::remove_file(am.driver_native_path(db_type)).ok();
return Ok(());
}
validate_native_agent_binary(source_path)?;
let native_path = am.driver_native_path(db_type);
let parent = native_path.parent().ok_or_else(|| format!("Invalid driver path: {}", native_path.display()))?;
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
let staging_path = parent.join(format!(".agent-import-{}", uuid::Uuid::new_v4()));
std::fs::copy(source_path, &staging_path).map_err(|e| format!("Failed to copy native agent: {e}"))?;
mark_executable(&staging_path)?;
replace_imported_agent_file(&staging_path, &native_path)?;
std::fs::remove_file(am.driver_jar_path(db_type)).ok();
let mut local_state = am.load_state();
local_state.installed_drivers.insert(
db_type.to_string(),
InstalledDriver {
version: "0.1.0-local".to_string(),
installed_at: chrono::Utc::now().to_rfc3339(),
jre: DEFAULT_JRE_KEY.to_string(),
},
);
am.save_state(&local_state)
}
pub fn import_agent_jar(am: &AgentManager, db_type: &str, jar_path: &Path) -> Result<(), String> {
import_agent_driver(am, db_type, jar_path)
}
fn replace_imported_agent_file(staging_path: &Path, target_path: &Path) -> Result<(), String> {
let backup_path = target_path.with_file_name(format!(
".{}-backup-{}",
target_path.file_name().and_then(|name| name.to_str()).unwrap_or("agent"),
uuid::Uuid::new_v4()
));
let had_existing = target_path.exists();
if had_existing {
std::fs::rename(target_path, &backup_path).map_err(|e| format!("Failed to replace existing agent: {e}"))?;
}
if let Err(error) = std::fs::rename(staging_path, target_path) {
if had_existing {
let _ = std::fs::rename(&backup_path, target_path);
}
let _ = std::fs::remove_file(staging_path);
return Err(format!("Failed to install agent: {error}"));
}
if had_existing {
std::fs::remove_file(backup_path).ok();
}
Ok(())
}
fn replace_imported_jre_dir(staging_dir: &Path, target_dir: &Path) -> Result<Option<PathBuf>, String> {
let backup_dir = target_dir.with_file_name(format!(
".{}-backup-{}",
target_dir.file_name().and_then(|name| name.to_str()).unwrap_or("jre"),
uuid::Uuid::new_v4()
));
let had_existing = target_dir.exists();
if had_existing {
std::fs::rename(target_dir, &backup_dir).map_err(|error| {
let _ = std::fs::remove_dir_all(staging_dir);
format!("Failed to replace existing JRE: {error}")
})?;
}
if let Err(error) = std::fs::rename(staging_dir, target_dir) {
if had_existing {
let _ = std::fs::rename(&backup_dir, target_dir);
}
let _ = std::fs::remove_dir_all(staging_dir);
return Err(format!("Failed to install JRE: {error}"));
}
if had_existing && remove_jre_dir_with_retry(&backup_dir).is_err() {
// The new runtime is already installed. Keep the old directory for
// startup cleanup rather than turning a successful import into an error.
return Ok(Some(backup_dir));
}
Ok(None)
}
fn jre_dir_contains_java(path: &Path) -> bool {
let java_name = if cfg!(windows) { "java.exe" } else { "java" };
path.join("bin").join(java_name).is_file()
|| path.join("Contents").join("Home").join("bin").join(java_name).is_file()
}
fn validate_native_agent_binary(path: &Path) -> Result<(), String> {
let mut file = std::fs::File::open(path).map_err(|e| format!("Failed to read native agent: {e}"))?;
let mut magic = [0_u8; 4];
file.read_exact(&mut magic).map_err(|e| format!("Failed to read native agent header: {e}"))?;
let valid = if cfg!(target_os = "windows") {
is_windows_binary_for_current_arch(&mut file, &magic)
} else if cfg!(target_os = "linux") {
is_elf_binary_for_current_arch(&mut file, &magic)
} else if cfg!(target_os = "macos") {
is_macho_binary_for_current_arch(&mut file, &magic)
} else {
false
};
if valid {
Ok(())
} else {
Err(format!("The selected file is not a {} native agent for this platform", AgentManager::current_platform()))
}
}
fn is_elf_binary_for_current_arch(file: &mut std::fs::File, magic: &[u8; 4]) -> bool {
if magic != b"\x7fELF" || file.seek(SeekFrom::Start(4)).is_err() {
return false;
}
let mut header = [0_u8; 16];
if file.read_exact(&mut header).is_err() || header[0] != 2 {
return false;
}
let machine = match header[1] {
1 => u16::from_le_bytes([header[14], header[15]]),
2 => u16::from_be_bytes([header[14], header[15]]),
_ => return false,
};
(cfg!(target_arch = "x86_64") && machine == 62) || (cfg!(target_arch = "aarch64") && machine == 183)
}
fn is_macho_binary_for_current_arch(file: &mut std::fs::File, magic: &[u8; 4]) -> bool {
const CPU_TYPE_X86_64: u32 = 0x0100_0007;
const CPU_TYPE_ARM64: u32 = 0x0100_000c;
let expected = if cfg!(target_arch = "aarch64") { CPU_TYPE_ARM64 } else { CPU_TYPE_X86_64 };
let thin_endian = match magic {
[0xce, 0xfa, 0xed, 0xfe] | [0xcf, 0xfa, 0xed, 0xfe] => Some(true),
[0xfe, 0xed, 0xfa, 0xce] | [0xfe, 0xed, 0xfa, 0xcf] => Some(false),
_ => None,
};
if let Some(little_endian) = thin_endian {
if file.seek(SeekFrom::Start(4)).is_err() {
return false;
}
let mut cpu_type = [0_u8; 4];
if file.read_exact(&mut cpu_type).is_err() {
return false;
}
let cpu_type = if little_endian { u32::from_le_bytes(cpu_type) } else { u32::from_be_bytes(cpu_type) };
return cpu_type == expected;
}
let (little_endian, arch_size) = match magic {
[0xca, 0xfe, 0xba, 0xbe] => (false, 20_u64),
[0xbe, 0xba, 0xfe, 0xca] => (true, 20_u64),
[0xca, 0xfe, 0xba, 0xbf] => (false, 32_u64),
[0xbf, 0xba, 0xfe, 0xca] => (true, 32_u64),
_ => return false,
};
if file.seek(SeekFrom::Start(4)).is_err() {
return false;
}
let mut count = [0_u8; 4];
if file.read_exact(&mut count).is_err() {
return false;
}
let count = if little_endian { u32::from_le_bytes(count) } else { u32::from_be_bytes(count) };
// A real universal binary has only a handful of slices; cap the count so
// a malformed header cannot trigger unbounded seeks during import.
if count == 0 || count > 64 {
return false;
}
for index in 0..count {
if file.seek(SeekFrom::Start(8 + u64::from(index) * arch_size)).is_err() {
return false;
}
let mut cpu_type = [0_u8; 4];
if file.read_exact(&mut cpu_type).is_err() {
return false;
}
let cpu_type = if little_endian { u32::from_le_bytes(cpu_type) } else { u32::from_be_bytes(cpu_type) };
if cpu_type == expected {
return true;
}
}
false
}
fn is_windows_binary_for_current_arch(file: &mut std::fs::File, magic: &[u8; 4]) -> bool {
if &magic[..2] != b"MZ" || file.seek(SeekFrom::Start(0x3c)).is_err() {
return false;
}
let mut pe_offset = [0_u8; 4];
if file.read_exact(&mut pe_offset).is_err()
|| file.seek(SeekFrom::Start(u32::from_le_bytes(pe_offset) as u64)).is_err()
{
return false;
}
let mut pe_header = [0_u8; 6];
if file.read_exact(&mut pe_header).is_err() || &pe_header[..4] != b"PE\0\0" {
return false;
}
let machine = u16::from_le_bytes([pe_header[4], pe_header[5]]);
(cfg!(target_arch = "x86_64") && machine == 0x8664) || (cfg!(target_arch = "aarch64") && machine == 0xaa64)
}
// ──────────── Tests ────────────
@@ -1449,6 +1772,31 @@ mod agent_download_url_tests {
assert_eq!(extract_jre_key_from_filename("jre/dbx-jre-21-macos-aarch64.tar.gz").as_deref(), Some("21"));
assert_eq!(extract_jre_key_from_filename("jre/jre-21-macos-aarch64.tar.gz").as_deref(), Some("21"));
}
#[test]
fn windows_native_header_validator_checks_cpu_architecture() {
let path = std::env::temp_dir().join(format!("dbx-agent-pe-test-{}", uuid::Uuid::new_v4()));
let expected_machine = if cfg!(target_arch = "aarch64") { 0xaa64_u16 } else { 0x8664_u16 };
let wrong_machine = if expected_machine == 0xaa64 { 0x8664_u16 } else { 0xaa64_u16 };
std::fs::write(&path, test_pe_binary(expected_machine)).unwrap();
let mut file = std::fs::File::open(&path).unwrap();
assert!(is_windows_binary_for_current_arch(&mut file, b"MZ\0\0"));
std::fs::write(&path, test_pe_binary(wrong_machine)).unwrap();
let mut file = std::fs::File::open(&path).unwrap();
assert!(!is_windows_binary_for_current_arch(&mut file, b"MZ\0\0"));
std::fs::remove_file(path).ok();
}
fn test_pe_binary(machine: u16) -> Vec<u8> {
let mut bytes = vec![0_u8; 0x48];
bytes[..2].copy_from_slice(b"MZ");
bytes[0x3c..0x40].copy_from_slice(&(0x40_u32).to_le_bytes());
bytes[0x40..0x44].copy_from_slice(b"PE\0\0");
bytes[0x44..0x46].copy_from_slice(&machine.to_le_bytes());
bytes
}
}
#[cfg(test)]
+400 -9
View File
@@ -3,9 +3,9 @@ use dbx_core::agent_manager::{
JreInfo, DEFAULT_JRE_KEY,
};
use dbx_core::agent_service::{
build_agent_list, clear_agent_download_cache, github_url_to_r2_path, import_agent_jar, import_agents_from_zip,
is_app_version_compatible, jre_needs_install, local_agent_jar_candidates, replace_download, uninstall_agent_driver,
AgentProgressEvent,
build_agent_list, clear_agent_download_cache, github_url_to_r2_path, import_agent_driver, import_agent_jar,
import_agents_from_zip, inspect_offline_zip, is_app_version_compatible, jre_needs_install,
local_agent_jar_candidates, replace_download, uninstall_agent_driver, AgentProgressEvent,
};
fn test_manager(name: &str) -> AgentManager {
@@ -406,6 +406,51 @@ fn local_jar_import_rejects_corrupt_jar() {
assert!(!manager.load_state().installed_drivers.contains_key("h2"));
}
#[test]
fn local_native_import_installs_current_platform_executable() {
let manager = test_manager("local-native-import");
let source = test_path("local-native-import-source").join(if cfg!(windows) {
"dbx-agent-kingbase-windows.exe"
} else {
"dbx-agent-kingbase"
});
std::fs::create_dir_all(source.parent().unwrap()).unwrap();
std::fs::write(&source, current_platform_native_binary()).unwrap();
import_agent_driver(&manager, "kingbase", &source).unwrap();
assert_eq!(std::fs::read(manager.driver_native_path("kingbase")).unwrap(), std::fs::read(&source).unwrap());
assert!(!manager.driver_jar_path("kingbase").exists());
assert_eq!(manager.load_state().installed_drivers["kingbase"].version, "0.1.0-local");
}
#[test]
fn local_native_import_rejects_wrong_platform_binary() {
let manager = test_manager("local-native-import-invalid");
let source = test_path("local-native-import-invalid-source").join("dbx-agent-kingbase");
std::fs::create_dir_all(source.parent().unwrap()).unwrap();
std::fs::write(&source, b"not-an-executable").unwrap();
let err = import_agent_driver(&manager, "kingbase", &source).unwrap_err();
assert!(err.contains(AgentManager::current_platform()));
assert!(!manager.driver_native_path("kingbase").exists());
assert!(!manager.load_state().installed_drivers.contains_key("kingbase"));
}
#[test]
fn local_native_import_rejects_wrong_arch_binary() {
let manager = test_manager("local-native-wrong-arch");
let native_path = test_path("local-native-wrong-arch-file").join("agent");
std::fs::create_dir_all(native_path.parent().unwrap()).unwrap();
std::fs::write(&native_path, native_binary_for_arch(!cfg!(target_arch = "aarch64"))).unwrap();
let err = import_agent_driver(&manager, "kingbase", &native_path).unwrap_err();
assert!(err.contains("not a"));
assert!(!manager.driver_native_path("kingbase").exists());
}
#[tokio::test]
async fn uninstall_driver_removes_artifact_and_state() {
let manager = test_manager("uninstall");
@@ -511,6 +556,53 @@ fn offline_zip_import_installs_release_named_jre() {
assert!(events.lock().unwrap().iter().any(|event| event.step == "jre-extract"));
}
#[test]
fn offline_zip_import_preserves_existing_jre_when_archive_is_corrupt() {
let manager = test_manager("offline-corrupt-jre-preserves-existing");
let root = test_path("offline-corrupt-jre-preserves-existing-zip");
let valid_zip = root.join("valid.zip");
let corrupt_zip = root.join("corrupt.zip");
std::fs::create_dir_all(&root).unwrap();
write_offline_driver_zip_with_jre(&valid_zip, "h2", "0.2.0", "21.0.12");
import_agents_from_zip(&manager, &valid_zip, |_| {}).unwrap();
let java_path = manager.jre_java_path(DEFAULT_JRE_KEY);
let original_java = std::fs::read(&java_path).unwrap();
write_offline_driver_zip_with_jre_bytes(&corrupt_zip, "h2", "0.3.0", "21.0.13", b"not-a-tar-gz".to_vec());
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).unwrap_err();
assert!(err.contains("Failed to extract JRE archive"));
assert_eq!(std::fs::read(java_path).unwrap(), original_java);
assert_eq!(manager.load_state().jre_versions.get(DEFAULT_JRE_KEY).map(String::as_str), Some("21.0.12"));
}
#[test]
fn offline_zip_import_preserves_existing_jre_when_driver_is_corrupt() {
let manager = test_manager("offline-corrupt-driver-preserves-jre");
let root = test_path("offline-corrupt-driver-preserves-jre-zip");
let valid_zip = root.join("valid.zip");
let corrupt_zip = root.join("corrupt.zip");
std::fs::create_dir_all(&root).unwrap();
write_offline_driver_zip_with_jre(&valid_zip, "h2", "0.2.0", "21.0.12");
import_agents_from_zip(&manager, &valid_zip, |_| {}).unwrap();
let java_path = manager.jre_java_path(DEFAULT_JRE_KEY);
let original_java = std::fs::read(&java_path).unwrap();
write_offline_driver_zip_with_jre_and_jar_bytes(
&corrupt_zip,
"h2",
"0.3.0",
"21.0.13",
test_jre_archive_bytes(),
b"jar".to_vec(),
);
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).unwrap_err();
assert!(err.contains("invalid or corrupt"));
assert_eq!(std::fs::read(java_path).unwrap(), original_java);
assert_eq!(manager.load_state().jre_versions.get(DEFAULT_JRE_KEY).map(String::as_str), Some("21.0.12"));
}
#[test]
fn offline_zip_import_rejects_corrupt_jar() {
let manager = test_manager("offline-corrupt-driver");
@@ -525,16 +617,284 @@ fn offline_zip_import_rejects_corrupt_jar() {
assert!(!manager.load_state().installed_drivers.contains_key("h2"));
}
#[test]
fn offline_zip_import_preserves_existing_driver_when_jar_is_corrupt() {
let manager = test_manager("offline-corrupt-driver-preserves-existing");
let root = test_path("offline-corrupt-driver-preserves-existing-zip");
let valid_zip = root.join("valid.zip");
let corrupt_zip = root.join("corrupt.zip");
std::fs::create_dir_all(&root).unwrap();
write_offline_driver_zip(&valid_zip, "h2", "0.2.0");
import_agents_from_zip(&manager, &valid_zip, |_| {}).unwrap();
let original = std::fs::read(manager.driver_jar_path("h2")).unwrap();
write_offline_driver_zip_with_jar(&corrupt_zip, "h2", "0.3.0", b"jar".to_vec());
let err = import_agents_from_zip(&manager, &corrupt_zip, |_| {}).unwrap_err();
assert!(err.contains("invalid or corrupt"));
assert_eq!(std::fs::read(manager.driver_jar_path("h2")).unwrap(), original);
assert_eq!(manager.load_state().installed_drivers["h2"].version, "0.2.0");
}
#[test]
fn offline_zip_import_keeps_legacy_unversioned_jar_compatibility() {
let manager = test_manager("offline-legacy-jar-zip");
let zip_path = test_path("offline-legacy-jar-zip").join("h2.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
let file = std::fs::File::create(&zip_path).unwrap();
let mut zip = zip::ZipWriter::new(file);
let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
let jar = test_agent_jar_bytes();
let registry = serde_json::json!({
"drivers": {
"h2": {
"version": "0.1.9",
"label": "H2",
"min_app_version": "0.1.0",
"jre": DEFAULT_JRE_KEY,
"jar": { "url": "https://example.com/dbx-agent-h2.jar", "size": jar.len() }
}
}
});
zip.start_file("agent-registry.json", options).unwrap();
std::io::Write::write_all(&mut zip, registry.to_string().as_bytes()).unwrap();
zip.start_file("drivers/dbx-agent-h2.jar", options).unwrap();
std::io::Write::write_all(&mut zip, &jar).unwrap();
zip.finish().unwrap();
import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap();
assert_eq!(manager.load_state().installed_drivers["h2"].version, "0.1.9");
}
#[test]
fn offline_zip_import_installs_versioned_native_driver_package() {
let manager = test_manager("offline-native-driver-zip");
let zip_path = test_path("offline-native-driver-zip").join("kingbase.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
write_offline_native_driver_zip(&zip_path, "kingbase", "0.1.34");
let result = import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap();
assert_eq!(result.drivers_installed, vec!["kingbase"]);
assert_eq!(manager.load_state().installed_drivers["kingbase"].version, "0.1.34");
assert_eq!(std::fs::read(manager.driver_native_path("kingbase")).unwrap(), current_platform_native_binary());
assert!(!manager.driver_jar_path("kingbase").exists());
}
#[test]
fn offline_zip_import_rejects_native_driver_for_another_platform() {
let manager = test_manager("offline-wrong-platform-native-driver-zip");
let zip_path = test_path("offline-wrong-platform-native-driver-zip").join("kingbase.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
let other_platform = if AgentManager::current_platform() == "windows-x64" { "linux-x64" } else { "windows-x64" };
write_offline_native_driver_zip_for_platform(&zip_path, "kingbase", "0.1.34", other_platform);
let err = import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap_err();
assert!(err.contains("no drivers compatible"));
assert!(err.contains(AgentManager::current_platform()));
assert!(!manager.is_driver_installed("kingbase"));
}
#[test]
fn offline_zip_import_preserves_existing_native_driver_when_binary_is_invalid() {
let manager = test_manager("offline-invalid-native-preserves-existing");
let root = test_path("offline-invalid-native-preserves-existing-zip");
let valid_zip = root.join("valid.zip");
let invalid_zip = root.join("invalid.zip");
std::fs::create_dir_all(&root).unwrap();
write_offline_native_driver_zip(&valid_zip, "kingbase", "0.1.34");
import_agents_from_zip(&manager, &valid_zip, |_| {}).unwrap();
let original = std::fs::read(manager.driver_native_path("kingbase")).unwrap();
write_offline_native_driver_zip_with_bytes(
&invalid_zip,
"kingbase",
"0.1.35",
AgentManager::current_platform(),
b"not-a-native-agent".to_vec(),
);
let err = import_agents_from_zip(&manager, &invalid_zip, |_| {}).unwrap_err();
assert!(err.contains("not a"));
assert_eq!(std::fs::read(manager.driver_native_path("kingbase")).unwrap(), original);
assert_eq!(manager.load_state().installed_drivers["kingbase"].version, "0.1.34");
}
#[test]
fn offline_zip_inspection_reports_drivers_and_jre() {
let zip_path = test_path("offline-inspection").join("agents.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
write_offline_driver_zip_with_jre(&zip_path, "h2", "0.2.0", "21.0.12");
let plan = inspect_offline_zip(&zip_path).unwrap();
assert_eq!(plan.driver_keys, vec!["h2"]);
assert!(plan.includes_jre);
}
#[test]
fn offline_zip_import_rejects_unknown_driver_type() {
let zip_path = test_path("offline-unknown-driver").join("agents.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
write_offline_driver_zip_with_jar(&zip_path, "unknown-driver", "0.1.0", test_agent_jar_bytes());
let err = inspect_offline_zip(&zip_path).unwrap_err();
assert!(err.contains("unknown driver type"));
}
#[test]
fn offline_zip_import_rejects_unsafe_entry_path() {
let zip_path = test_path("offline-unsafe-entry").join("agents.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
let file = std::fs::File::create(&zip_path).unwrap();
let mut zip = zip::ZipWriter::new(file);
let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
zip.start_file("agent-registry.json", options).unwrap();
std::io::Write::write_all(&mut zip, br#"{"drivers":{}}"#).unwrap();
zip.start_file("../escape", options).unwrap();
std::io::Write::write_all(&mut zip, b"escape").unwrap();
zip.finish().unwrap();
let err = inspect_offline_zip(&zip_path).unwrap_err();
assert!(err.contains("unsafe path"));
}
#[test]
fn offline_zip_import_prefers_native_artifact_over_java_fallback() {
let manager = test_manager("offline-native-preferred");
let zip_path = test_path("offline-native-preferred-zip").join("kingbase.zip");
std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap();
write_offline_hybrid_driver_zip(&zip_path, "kingbase", "0.1.34");
let plan = inspect_offline_zip(&zip_path).unwrap();
let result = import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap();
assert_eq!(plan.driver_keys, vec!["kingbase"]);
assert_eq!(result.drivers_installed, vec!["kingbase"]);
assert!(manager.driver_native_path("kingbase").exists());
assert!(!manager.driver_jar_path("kingbase").exists());
}
fn write_offline_driver_zip(path: &std::path::Path, db_type: &str, version: &str) {
write_offline_driver_zip_with_jar(path, db_type, version, test_agent_jar_bytes());
}
fn write_offline_driver_zip_with_jre(path: &std::path::Path, db_type: &str, version: &str, jre_version: &str) {
fn write_offline_native_driver_zip(path: &std::path::Path, db_type: &str, version: &str) {
write_offline_native_driver_zip_for_platform(path, db_type, version, AgentManager::current_platform());
}
fn write_offline_native_driver_zip_for_platform(path: &std::path::Path, db_type: &str, version: &str, platform: &str) {
write_offline_native_driver_zip_with_bytes(path, db_type, version, platform, current_platform_native_binary());
}
fn write_offline_native_driver_zip_with_bytes(
path: &std::path::Path,
db_type: &str,
version: &str,
platform: &str,
native: Vec<u8>,
) {
let file = std::fs::File::create(path).unwrap();
let mut zip = zip::ZipWriter::new(file);
let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
let extension = if platform.starts_with("windows-") { ".exe" } else { "" };
let filename = format!("dbx-agent-{db_type}-{version}-{platform}{extension}");
let registry = serde_json::json!({
"jres": {},
"drivers": {
db_type: {
"version": version,
"label": db_type,
"min_app_version": "0.6.0",
"jre": DEFAULT_JRE_KEY,
"native": {
platform: {
"url": format!("https://example.com/{filename}"),
"size": native.len()
}
}
}
}
});
zip.start_file("agent-registry.json", options).unwrap();
std::io::Write::write_all(&mut zip, registry.to_string().as_bytes()).unwrap();
zip.start_file(format!("drivers/{filename}"), options).unwrap();
std::io::Write::write_all(&mut zip, &native).unwrap();
zip.finish().unwrap();
}
fn write_offline_hybrid_driver_zip(path: &std::path::Path, db_type: &str, version: &str) {
let file = std::fs::File::create(path).unwrap();
let mut zip = zip::ZipWriter::new(file);
let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
let jar = test_agent_jar_bytes();
let jre_archive = test_jre_archive_bytes();
let native = current_platform_native_binary();
let platform = AgentManager::current_platform();
let extension = if platform.starts_with("windows-") { ".exe" } else { "" };
let jar_filename = format!("dbx-agent-{db_type}-{version}.jar");
let native_filename = format!("dbx-agent-{db_type}-{version}-{platform}{extension}");
let registry = serde_json::json!({
"jres": {},
"drivers": {
db_type: {
"version": version,
"label": db_type,
"min_app_version": "0.1.0",
"jre": DEFAULT_JRE_KEY,
"jar": { "url": format!("https://example.com/{jar_filename}"), "size": jar.len() },
"native": {
platform: { "url": format!("https://example.com/{native_filename}"), "size": native.len() }
}
}
}
});
zip.start_file("agent-registry.json", options).unwrap();
std::io::Write::write_all(&mut zip, registry.to_string().as_bytes()).unwrap();
zip.start_file(format!("drivers/{jar_filename}"), options).unwrap();
std::io::Write::write_all(&mut zip, &jar).unwrap();
zip.start_file(format!("drivers/{native_filename}"), options).unwrap();
std::io::Write::write_all(&mut zip, &native).unwrap();
zip.finish().unwrap();
}
fn write_offline_driver_zip_with_jre(path: &std::path::Path, db_type: &str, version: &str, jre_version: &str) {
write_offline_driver_zip_with_jre_bytes(path, db_type, version, jre_version, test_jre_archive_bytes());
}
fn write_offline_driver_zip_with_jre_bytes(
path: &std::path::Path,
db_type: &str,
version: &str,
jre_version: &str,
jre_archive: Vec<u8>,
) {
write_offline_driver_zip_with_jre_and_jar_bytes(
path,
db_type,
version,
jre_version,
jre_archive,
test_agent_jar_bytes(),
);
}
fn write_offline_driver_zip_with_jre_and_jar_bytes(
path: &std::path::Path,
db_type: &str,
version: &str,
jre_version: &str,
jre_archive: Vec<u8>,
jar: Vec<u8>,
) {
let file = std::fs::File::create(path).unwrap();
let mut zip = zip::ZipWriter::new(file);
let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
let registry = serde_json::json!({
"jres": {
DEFAULT_JRE_KEY: {
@@ -548,7 +908,7 @@ fn write_offline_driver_zip_with_jre(path: &std::path::Path, db_type: &str, vers
"label": db_type,
"min_app_version": "0.1.0",
"jre": DEFAULT_JRE_KEY,
"jar": { "url": format!("https://example.com/dbx-agent-{db_type}.jar"), "size": jar.len() }
"jar": { "url": format!("https://example.com/dbx-agent-{db_type}-{version}.jar"), "size": jar.len() }
}
}
});
@@ -558,7 +918,7 @@ fn write_offline_driver_zip_with_jre(path: &std::path::Path, db_type: &str, vers
zip.start_file(format!("jre/dbx-jre-{DEFAULT_JRE_KEY}-{}.tar.gz", AgentManager::current_platform()), options)
.unwrap();
std::io::Write::write_all(&mut zip, &jre_archive).unwrap();
zip.start_file(format!("drivers/dbx-agent-{db_type}.jar"), options).unwrap();
zip.start_file(format!("drivers/dbx-agent-{db_type}-{version}.jar"), options).unwrap();
std::io::Write::write_all(&mut zip, &jar).unwrap();
zip.finish().unwrap();
}
@@ -598,14 +958,14 @@ fn write_offline_driver_zip_with_jar(path: &std::path::Path, db_type: &str, vers
"label": db_type,
"min_app_version": "0.1.0",
"jre": DEFAULT_JRE_KEY,
"jar": { "url": format!("https://example.com/dbx-agent-{db_type}.jar"), "size": jar.len() }
"jar": { "url": format!("https://example.com/dbx-agent-{db_type}-{version}.jar"), "size": jar.len() }
}
}
});
zip.start_file("agent-registry.json", options).unwrap();
std::io::Write::write_all(&mut zip, registry.to_string().as_bytes()).unwrap();
zip.start_file(format!("drivers/dbx-agent-{db_type}.jar"), options).unwrap();
zip.start_file(format!("drivers/dbx-agent-{db_type}-{version}.jar"), options).unwrap();
std::io::Write::write_all(&mut zip, &jar).unwrap();
zip.finish().unwrap();
}
@@ -623,3 +983,34 @@ fn test_agent_jar_bytes() -> Vec<u8> {
std::io::Write::write_all(&mut zip, b"Manifest-Version: 1.0\nMain-Class: com.dbx.agent.TestAgent\n\n").unwrap();
zip.finish().unwrap().into_inner()
}
fn current_platform_native_binary() -> Vec<u8> {
native_binary_for_arch(cfg!(target_arch = "aarch64"))
}
fn native_binary_for_arch(aarch64: bool) -> Vec<u8> {
if cfg!(windows) {
let mut bytes = vec![0_u8; 0x48];
bytes[..2].copy_from_slice(b"MZ");
bytes[0x3c..0x40].copy_from_slice(&(0x40_u32).to_le_bytes());
bytes[0x40..0x44].copy_from_slice(b"PE\0\0");
let machine = if aarch64 { 0xaa64_u16 } else { 0x8664_u16 };
bytes[0x44..0x46].copy_from_slice(&machine.to_le_bytes());
bytes
} else if cfg!(target_os = "linux") {
let mut bytes = vec![0_u8; 20];
bytes[..4].copy_from_slice(b"\x7fELF");
bytes[4] = 2;
bytes[5] = 1;
let machine = if aarch64 { 183_u16 } else { 62_u16 };
bytes[18..20].copy_from_slice(&machine.to_le_bytes());
bytes
} else if cfg!(target_os = "macos") {
let mut bytes = vec![0xcf, 0xfa, 0xed, 0xfe];
let cpu_type = if aarch64 { 0x0100_000c_u32 } else { 0x0100_0007_u32 };
bytes.extend_from_slice(&cpu_type.to_le_bytes());
bytes
} else {
Vec::new()
}
}
+2 -1
View File
@@ -245,7 +245,8 @@ async fn main() {
.route("/agents/upgrade-all", post(routes::agents::upgrade_all_agents))
.route("/agents/uninstall", post(routes::agents::uninstall_agent))
.route("/agents/import-offline", post(routes::agents::import_agents_from_zip))
.route("/agents/import-jar", post(routes::agents::import_agent_jar))
.route("/agents/import-driver", post(routes::agents::import_agent_driver_file))
.route("/agents/import-jar", post(routes::agents::import_agent_driver_file))
.route(
"/agents/java-runtime",
get(routes::agents::get_agent_java_runtime_config).post(routes::agents::set_agent_java_runtime_config),
+53 -25
View File
@@ -7,9 +7,10 @@ use dbx_core::agent_manager::{
AgentDriverInfo, AgentState, DriverStoreUsage, JavaRuntimeConfig, JavaRuntimeMode, DEFAULT_JRE_KEY,
};
use dbx_core::agent_service::{
build_agent_list, clear_agent_download_cache, fetch_registry,
import_agents_from_zip as import_agents_from_zip_core, install_agent_driver, invalidate_registry_cache,
reinstall_agent_jre, uninstall_agent_driver, uninstall_agent_jre, upgrade_all_agent_drivers, AgentProgressEvent,
build_agent_list, clear_agent_download_cache, fetch_registry, import_agent_driver,
import_agents_from_zip as import_agents_from_zip_core, inspect_offline_zip, install_agent_driver,
invalidate_registry_cache, reinstall_agent_jre, uninstall_agent_driver, uninstall_agent_jre,
upgrade_all_agent_drivers, AgentProgressEvent, OfflineImportPlan,
};
use dbx_core::driver_runtime::DriverRuntimeSummary;
use futures::Stream;
@@ -177,18 +178,22 @@ pub async fn import_agents_from_zip(
}
let zip_path = tmp_dir.join(format!("agent-offline-{}.zip", uuid::Uuid::new_v4()));
let mut upload = tokio::fs::File::create(&zip_path).await.map_err(|err| AppError(err.to_string()))?;
let mut field = field;
while let Some(chunk) = field.chunk().await.map_err(|err| AppError(err.to_string()))? {
upload.write_all(&chunk).await.map_err(|err| AppError(err.to_string()))?;
}
upload.flush().await.map_err(|err| AppError(err.to_string()))?;
drop(upload);
let tx = progress_sender(&state, "global").await;
let result =
let result = async {
let mut upload = tokio::fs::File::create(&zip_path).await.map_err(|err| AppError(err.to_string()))?;
let mut field = field;
while let Some(chunk) = field.chunk().await.map_err(|err| AppError(err.to_string()))? {
upload.write_all(&chunk).await.map_err(|err| AppError(err.to_string()))?;
}
upload.flush().await.map_err(|err| AppError(err.to_string()))?;
drop(upload);
let plan = inspect_offline_zip(&zip_path).map_err(AppError)?;
ensure_no_offline_import_blockers(&state.app, &plan).await.map_err(AppError)?;
import_agents_from_zip_core(&state.app.agent_manager, &zip_path, |event| send_progress_event(&tx, event))
.map_err(AppError);
.map_err(AppError)
}
.await;
let _ = std::fs::remove_file(&zip_path);
let result = result?;
@@ -199,37 +204,44 @@ pub async fn import_agents_from_zip(
Err(AppError("No file uploaded".to_string()))
}
pub async fn import_agent_jar(
pub async fn import_agent_driver_file(
State(state): State<Arc<WebState>>,
mut multipart: Multipart,
) -> Result<Json<serde_json::Value>, AppError> {
let mut db_type: Option<String> = None;
let mut jar_data: Option<Vec<u8>> = None;
let mut jar_name = String::new();
let mut driver_data: Option<Vec<u8>> = None;
let mut driver_name = String::new();
while let Ok(Some(field)) = multipart.next_field().await {
let name = field.name().unwrap_or("").to_string();
if name == "dbType" {
db_type = Some(field.text().await.map_err(|e| AppError(e.to_string()))?);
} else if name == "file" {
jar_name = field.file_name().unwrap_or("driver.jar").to_string();
if !jar_name.to_lowercase().ends_with(".jar") {
return Err(AppError("Only .jar files can be imported".to_string()));
}
jar_data = Some(field.bytes().await.map_err(|e| AppError(e.to_string()))?.to_vec());
driver_name = field.file_name().unwrap_or("agent").to_string();
driver_data = Some(field.bytes().await.map_err(|e| AppError(e.to_string()))?.to_vec());
}
}
let db_type = db_type.ok_or_else(|| AppError("Missing dbType field".to_string()))?;
let data = jar_data.ok_or_else(|| AppError("No file uploaded".to_string()))?;
let data = driver_data.ok_or_else(|| AppError("No file uploaded".to_string()))?;
let temp_dir = state.app.plugins.root_dir().join("jar_upload_tmp");
let temp_dir = state.app.plugins.root_dir().join("agent_upload_tmp");
std::fs::create_dir_all(&temp_dir).map_err(|e| AppError(e.to_string()))?;
let tmp_path = temp_dir.join(&jar_name);
let suffix = std::path::Path::new(&driver_name)
.extension()
.and_then(|extension| extension.to_str())
.map(|extension| format!(".{extension}"))
.unwrap_or_default();
let tmp_path = temp_dir.join(format!("agent-{}{}", uuid::Uuid::new_v4(), suffix));
std::fs::write(&tmp_path, &data).map_err(|e| AppError(e.to_string()))?;
dbx_core::agent_service::import_agent_jar(&state.app.agent_manager, &db_type, &tmp_path).map_err(AppError::from)?;
let result = async {
ensure_no_agent_update_blockers(&state.app, std::slice::from_ref(&db_type)).await.map_err(AppError)?;
import_agent_driver(&state.app.agent_manager, &db_type, &tmp_path).map_err(AppError::from)
}
.await;
let _ = std::fs::remove_file(&tmp_path);
result?;
Ok(Json(serde_json::json!({ "success": true })))
}
@@ -302,3 +314,19 @@ async fn ensure_no_agent_update_blockers(
Err(format!("请先关闭以下数据库连接后再更新驱动: {}", blockers.join(", ")))
}
}
async fn ensure_no_offline_import_blockers(
state: &dbx_core::connection::AppState,
plan: &OfflineImportPlan,
) -> Result<(), String> {
let mut driver_keys = plan.driver_keys.clone();
if plan.includes_jre {
// Replacing a managed JRE affects every running Java Agent, so include
// all active runtimes in the same connection-aware update preflight.
driver_keys.extend(state.agent_manager.active_daemon_keys().await);
driver_keys.extend(state.active_agent_connection_driver_keys().await);
driver_keys.sort();
driver_keys.dedup();
}
ensure_no_agent_update_blockers(state, &driver_keys).await
}