diff --git a/crates/dbx-drivers/src/db/sqlite_worker.rs b/crates/dbx-drivers/src/db/sqlite_worker.rs index 2a097815eb..37b339179f 100644 --- a/crates/dbx-drivers/src/db/sqlite_worker.rs +++ b/crates/dbx-drivers/src/db/sqlite_worker.rs @@ -1,4 +1,4 @@ -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -24,6 +24,11 @@ const DEFAULT_PERSIST_DIR: &str = "~/.cache/dbx/sqlite-worker"; const CONSENT_FILE_NAME: &str = "sqlite-worker-consent.json"; const SQLITE_WORKER_MAX_RESPONSE_BYTES: usize = 16 * 1024 * 1024; static SQLITE_SSH_RUNTIME_ENABLED: AtomicBool = AtomicBool::new(false); +// ponytail: one mutex per host+path for the process lifetime. Opening a table +// prewarms and queries at once; both used to `cat` the same `.part` and publish +// a corrupt worker. Upgrade path: evict idle keys if the map ever matters. +static SQLITE_WORKER_INSTALLS: std::sync::OnceLock>>>> = + std::sync::OnceLock::new(); pub fn sqlite_worker_chain_id(connection_id: &str) -> String { format!("{connection_id}:sqlite-worker") @@ -401,6 +406,11 @@ pub async fn connect_sqlite_worker( ensure_remote_sqlite_file_exists(session.as_ref(), &expanded_db).await?; if placement != SqliteWorkerPlacement::Preplaced { ensure_worker_consent(data_dir, &identity, &remote_path, &digest).await?; + } + let _install = remote_worker_install_lock(&identity, &remote_path).await; + if placement != SqliteWorkerPlacement::Preplaced + && verify_remote_digest(session.as_ref(), &remote_path, &digest).await.is_err() + { upload_worker(session.as_ref(), &remote_path, &local_worker.bytes).await?; } verify_remote_digest(session.as_ref(), &remote_path, &digest).await?; @@ -690,10 +700,27 @@ async fn remove_uploaded_session_worker(session: &Handle, path: &str) } } -async fn upload_worker(session: &Handle, dest: &str, bytes: &[u8]) -> Result<(), String> { +fn worker_upload_command(dest: &str, byte_len: usize) -> String { let quoted = shell_quote(dest); - let command = - format!("mkdir -p \"$(dirname {quoted})\" && cat > {quoted}.part && chmod 700 {quoted}.part && mv {quoted}.part {quoted}"); + // Unique part file: two connects must not `cat` into the same path. + let part = shell_quote(&format!("{dest}.part-{}", uuid::Uuid::new_v4().simple())); + format!( + "mkdir -p \"$(dirname {quoted})\" && cat > {part} && test \"$(wc -c < {part})\" -eq {byte_len} && chmod 700 {part} && mv {part} {quoted}" + ) +} + +async fn remote_worker_install_lock(identity: &str, remote_path: &str) -> tokio::sync::OwnedMutexGuard<()> { + let key = format!("{identity}\n{remote_path}"); + let installs = SQLITE_WORKER_INSTALLS.get_or_init(|| AsyncMutex::new(HashMap::new())); + let slot = { + let mut installs = installs.lock().await; + installs.entry(key).or_insert_with(|| Arc::new(AsyncMutex::new(()))).clone() + }; + slot.lock_owned().await +} + +async fn upload_worker(session: &Handle, dest: &str, bytes: &[u8]) -> Result<(), String> { + let command = worker_upload_command(dest, bytes.len()); ssh_exec_with_stdin(session, command, bytes).await } @@ -707,29 +734,41 @@ async fn ssh_exec_with_stdin( channel.data(stdin).await.map_err(|e| e.to_string())?; channel.eof().await.map_err(|e| e.to_string())?; let mut exit_status = None; + let mut stderr = Vec::new(); loop { match channel.wait().await { Some(ChannelMsg::ExitStatus { exit_status: status }) => exit_status = Some(status), Some(ChannelMsg::Close) | None => break, - Some(ChannelMsg::Eof | ChannelMsg::Data { .. } | ChannelMsg::ExtendedData { .. }) => {} + Some(ChannelMsg::ExtendedData { data, .. }) => stderr.extend_from_slice(data.as_ref()), + Some(ChannelMsg::Eof | ChannelMsg::Data { .. }) => {} Some(_) => {} } } - remote_exec_status(exit_status) + remote_exec_status(exit_status, &String::from_utf8_lossy(&stderr)) } -fn remote_exec_status(exit_status: Option) -> Result<(), String> { +fn remote_exec_status(exit_status: Option, stderr: &str) -> Result<(), String> { match exit_status { Some(0) => Ok(()), - Some(code) => Err(format!("remote command exited with status {code}")), - None => Err("remote command closed without an exit status".to_string()), + Some(code) => Err(remote_command_failure(format!("remote command exited with status {code}"), stderr)), + None => Err(remote_command_failure("remote command closed without an exit status".to_string(), stderr)), + } +} + +fn remote_command_failure(message: String, stderr: &str) -> String { + let detail = stderr.split_whitespace().collect::>().join(" "); + let detail = detail.chars().take(400).collect::(); + if detail.is_empty() { + message + } else { + format!("{message}: {detail}") } } fn validate_remote_download(copied: u64, exit_status: Option) -> Result<(), String> { // Check exit status first: a failed `cat` with empty output is better // explained by the nonzero status than by the emptiness it caused. - remote_exec_status(exit_status)?; + remote_exec_status(exit_status, "")?; if copied == 0 { return Err("Downloaded SQLite backup was empty".to_string()); } @@ -840,6 +879,28 @@ mod tests { assert_eq!(sqlite_worker_chain_id("conn-1"), "conn-1:sqlite-worker"); } + #[test] + fn worker_upload_uses_a_private_part_file_and_checks_size() { + let command = worker_upload_command("/home/u/.cache/dbx/sqlite-worker/abc", 12); + assert!(command.contains("mkdir -p")); + assert!(command.contains(".part-")); + assert!(!command.contains(".part &&")); + assert!(command.contains("test \"$(wc -c < ")); + assert!(command.contains("-eq 12")); + assert!(command.contains("&& mv ")); + assert_ne!(worker_upload_command("/tmp/worker", 1), worker_upload_command("/tmp/worker", 1)); + } + + #[test] + fn remote_command_failure_keeps_stderr() { + assert_eq!(remote_exec_status(Some(0), "noise").unwrap(), ()); + assert_eq!( + remote_exec_status(Some(1), "cannot execute binary file").unwrap_err(), + "remote command exited with status 1: cannot execute binary file" + ); + assert!(remote_exec_status(None, "").unwrap_err().contains("without an exit status")); + } + #[test] fn remote_download_requires_a_successful_non_empty_transfer() { assert!(validate_remote_download(1, Some(0)).is_ok());