From 512d8faf62782ffd1a11e8d97e7f6ea10283c8ec Mon Sep 17 00:00:00 2001 From: atontini Date: Mon, 19 Oct 2026 07:59:12 +0200 Subject: [PATCH] fix(rust): release update-hook closures on final disconnect --- libsql/src/local/connection.rs | 127 ++++++++++++++++++++++++++++++--- 1 file changed, 117 insertions(+), 10 deletions(-) diff --git a/libsql/src/local/connection.rs b/libsql/src/local/connection.rs index 4c54f5acd0..b9dc195f62 100644 --- a/libsql/src/local/connection.rs +++ b/libsql/src/local/connection.rs @@ -33,6 +33,8 @@ pub struct Connection { pub(crate) writer: Option, authorizer: Arc>>, + // SQLite has no destructor for update-hook contexts. Share ownership with clones. + update_hook: Arc>>>, } impl Drop for Connection { @@ -76,6 +78,7 @@ impl Connection { #[cfg(feature = "replication")] writer: db.writer()?, authorizer: Arc::new(RwLock::new(None)), + update_hook: Arc::new(RwLock::new(None)), }; #[cfg(feature = "sync")] if let Some(_) = db.sync_ctx { @@ -101,12 +104,19 @@ impl Connection { #[cfg(feature = "replication")] writer: None, authorizer: Arc::new(RwLock::new(None)), + update_hook: Arc::new(RwLock::new(None)), } } /// Disconnect from the database. pub fn disconnect(&mut self) { - if Arc::get_mut(&mut self.drop_ref).is_some() { + if Arc::get_mut(&mut self.drop_ref).is_some() && !self.raw.is_null() { + let owned_hook = self.update_hook.write().take(); + if owned_hook.is_some() { + // Unregister before releasing the callback and its captures. + unsafe { ffi::sqlite3_update_hook(self.raw, None, std::ptr::null_mut()) }; + } + drop(owned_hook); unsafe { libsql_sys::ffi::sqlite3_close_v2(self.raw); self.raw = std::ptr::null_mut(); @@ -412,20 +422,21 @@ impl Connection { /// Installs update hook pub fn add_update_hook(&self, cb: Box) { - let c = Box::new(Container { cb }); - let ptr: *mut Container = std::ptr::from_mut(Box::leak(c)); - - let old_data = unsafe { + let mut c = Box::new(Container { cb }); + let ptr: *mut Container = &mut *c; + let mut owner = self.update_hook.write(); + unsafe { ffi::sqlite3_update_hook( self.raw, Some(update_hook_cb), ptr as *mut ::std::os::raw::c_void, - ) - }; - - if !old_data.is_null() { - let _ = unsafe { Box::from_raw(old_data as *mut Container) }; + ); } + // Only free contexts owned by this wrapper, including when a raw-handle + // caller has installed a different hook. Release captures outside the lock. + let previous = owner.replace(c); + drop(owner); + drop(previous); } pub fn enable_load_extension(&self, onoff: bool) -> Result<()> { @@ -821,6 +832,102 @@ mod tests { OpenFlags, }; + #[test] + fn update_hook_replacement_and_final_disconnect_drop_captures() { + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + struct Capture(Arc); + impl Drop for Capture { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + let drops = Arc::new(AtomicUsize::new(0)); + let db = Database::new(":memory:".to_string(), OpenFlags::default()); + let conn = Connection::connect(&db).unwrap(); + let first = Capture(drops.clone()); + conn.add_update_hook(Box::new(move |_, _, _, _| { + let _ = &first; + })); + let second = Capture(drops.clone()); + let calls = Arc::new(AtomicUsize::new(0)); + let callback_calls = calls.clone(); + conn.add_update_hook(Box::new(move |_, _, _, _| { + let _ = &second; + callback_calls.fetch_add(1, Ordering::SeqCst); + })); + conn.execute("CREATE TABLE t(x)", Params::None).unwrap(); + conn.execute("INSERT INTO t VALUES (1)", Params::None) + .unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!( + drops.load(Ordering::SeqCst), + 1, + "replaced hook must be released" + ); + let mut last = conn.clone(); + drop(conn); + assert_eq!( + drops.load(Ordering::SeqCst), + 1, + "clones retain the active hook" + ); + last.execute("INSERT INTO t VALUES (2)", Params::None) + .unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 2); + last.disconnect(); + assert_eq!( + drops.load(Ordering::SeqCst), + 2, + "disconnect releases the final hook" + ); + last.disconnect(); + drop(last); + assert_eq!(drops.load(Ordering::SeqCst), 2, "cleanup is idempotent"); + } + + #[test] + fn update_hook_does_not_free_foreign_context() { + let db = Database::new(":memory:".to_string(), OpenFlags::default()); + let conn = Connection::connect(&db).unwrap(); + // Raw-handle callers may install a context that is not a Box. + let mut foreign = 42_u64; + unsafe { + libsql_sys::ffi::sqlite3_update_hook( + conn.handle(), + None, + (&mut foreign as *mut u64).cast(), + ); + } + conn.add_update_hook(Box::new(|_, _, _, _| {})); + conn.execute("CREATE TABLE t(x)", Params::None).unwrap(); + conn.execute("INSERT INTO t VALUES (1)", Params::None) + .unwrap(); + drop(conn); + assert_eq!(foreign, 42); + } + + #[test] + fn prepared_statement_retains_update_hook_until_final_drop() { + use std::sync::Arc; + let db = Database::new(":memory:".to_string(), OpenFlags::default()); + let conn = Connection::connect(&db).unwrap(); + conn.execute("CREATE TABLE t(x)", Params::None).unwrap(); + let held = Arc::new(()); + let weak = Arc::downgrade(&held); + conn.add_update_hook(Box::new(move |_, _, _, _| { + let _ = &held; + })); + let statement = conn.prepare("INSERT INTO t VALUES (1)").unwrap(); + drop(conn); + assert!(weak.upgrade().is_some()); + statement.execute(&Params::None).unwrap(); + drop(statement); + assert!(weak.upgrade().is_none()); + } + #[tokio::test] pub async fn test_kek() { let temp_dir = tempfile::tempdir().unwrap();