diesel/sqlite/
auto_extension.rs

1#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
2extern crate libsqlite3_sys as ffi;
3#[cfg(all(target_family = "wasm", target_os = "unknown"))]
4use sqlite_wasm_rs as ffi;
5
6use crate::result::Error::DatabaseError;
7use crate::result::*;
8use crate::sqlite::SqliteConnection;
9use alloc::boxed::Box;
10use alloc::string::{String, ToString};
11use core::ffi::{c_char, c_int};
12#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
13type RawApiPointer = *const core::ffi::c_void;
14#[cfg(all(target_family = "wasm", target_os = "unknown"))]
15type RawApiPointer = *const ffi::sqlite3_api_routines;
16
17/// SQLite's auto-extension callback type.
18type RawAutoExtension = unsafe extern "C" fn(
19    db: *mut ffi::sqlite3,
20    pz_err_msg: *mut *mut c_char,
21    p_api: RawApiPointer,
22) -> c_int;
23
24// TODO: diesel 3.0 use `libsqlite3-sys` declarations after raising the minimum to 0.29.
25#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
26mod auto_extension_ffi {
27    use super::{RawAutoExtension, c_int};
28
29    // SAFETY: These declarations match SQLite's C ABI and bypass incompatible callback types in older bindings.
30    #[allow(unsafe_code)]
31    unsafe extern "C" {
32        pub(super) fn sqlite3_auto_extension(entry_point: Option<RawAutoExtension>) -> c_int;
33        pub(super) fn sqlite3_cancel_auto_extension(entry_point: Option<RawAutoExtension>)
34        -> c_int;
35    }
36}
37
38#[cfg(all(target_family = "wasm", target_os = "unknown"))]
39use ffi as auto_extension_ffi;
40
41use super::BorrowedSqliteConnection;
42
43/// Registers an auto-extension that runs for every SQLite connection opened in
44/// this process, including non-Diesel ones.
45///
46/// This is a safe wrapper around [`sqlite3_auto_extension`][docs]. The callback
47/// receives the [`SqliteConnection`] being opened and returns `Ok(())` to
48/// continue or an error to fail the open. Use it to register SQL functions,
49/// collations, or aggregates through the usual connection API, or to initialize
50/// a statically linked C extension such as Spatialite or sqlite-vec via
51/// [`SqliteConnection::with_raw_connection`].
52///
53/// Call this before opening any connection. Extensions run in registration
54/// order, and the first error aborts the open. The callback must be a `fn` item
55/// or a closure that captures only zero-sized values (enforced at compile
56/// time), and registering the same `fn` twice is a no-op. It may run on several
57/// threads at once and must not open another connection (which would re-enter
58/// the auto-extensions and recurse) or call [`register_auto_extension`],
59/// [`cancel_auto_extension`], or [`reset_auto_extension`]. Panics are caught and
60/// turned into a failed open.
61///
62/// [docs]: https://www.sqlite.org/c3ref/auto_extension.html
63///
64/// # Example
65///
66/// ```rust
67/// # // This test registers Rust callbacks that sqlite invokes through the
68/// # // FFI boundary, which is not supported by miri `-Zmiri-native-lib`.
69/// # #[cfg(not(miri))]
70/// # fn main() {
71/// use diesel::dsl::sql;
72/// use diesel::prelude::*;
73/// use diesel::sql_types::Integer;
74/// use diesel::sqlite::{register_auto_extension, reset_auto_extension, SqliteConnection, BorrowedSqliteConnection};
75/// use std::pin::Pin;
76///
77/// // Registers a case-insensitive collation on every new connection.
78/// fn my_ext(conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
79///     unsafe {
80///         // SAFETY: This is safe as long as we don't move or replace out the
81///         // inner SqliteConnection. None of the functions defined on SqliteConnection
82///         // do that
83///         conn.inner_mut().register_collation("RUSTNOCASE", |a, b| a.to_lowercase().cmp(&b.to_lowercase()))
84///     }
85/// }
86///
87/// register_auto_extension(my_ext).unwrap();
88///
89/// // Every future connection now has the collation.
90/// let mut conn = SqliteConnection::establish(":memory:").unwrap();
91/// let equal: i32 = sql::<Integer>("SELECT 'a' = 'A' COLLATE RUSTNOCASE")
92///     .get_result(&mut conn)
93///     .unwrap();
94/// assert_eq!(equal, 1);
95/// # reset_auto_extension();
96/// # }
97/// # #[cfg(miri)]
98/// # fn main() {}
99/// ```
100#[allow(unsafe_code)]
101pub fn register_auto_extension<F>(extension: F) -> QueryResult<()>
102where
103    F: Fn(&'_ mut BorrowedSqliteConnection) -> QueryResult<()> + Sync + 'static,
104{
105    // SAFETY: `entry_point` returns a stable function pointer with SQLite's callback ABI.
106    let result =
107        unsafe { auto_extension_ffi::sqlite3_auto_extension(Some(entry_point(extension))) };
108    if result == ffi::SQLITE_OK {
109        Ok(())
110    } else {
111        Err(DatabaseError(
112            DatabaseErrorKind::Unknown,
113            Box::new(ffi::code_to_str(result).to_string()),
114        ))
115    }
116}
117
118/// Removes a previously registered auto-extension, returning `true` if it was
119/// found ([docs][cancel_docs]).
120///
121/// Pass the same `fn` item given to [`register_auto_extension`]. A closure
122/// cannot be cancelled this way, because its type cannot be named again. Use
123/// [`reset_auto_extension`] to clear everything instead.
124///
125/// [cancel_docs]: https://www.sqlite.org/c3ref/cancel_auto_extension.html
126#[allow(unsafe_code)]
127pub fn cancel_auto_extension<F>(extension: F) -> bool
128where
129    F: Fn(&'_ mut BorrowedSqliteConnection) -> QueryResult<()> + Sync + 'static,
130{
131    // SAFETY: `entry_point` returns the same stable pointer used for registration.
132    unsafe { auto_extension_ffi::sqlite3_cancel_auto_extension(Some(entry_point(extension))) != 0 }
133}
134
135/// Clears **all** registered auto-extensions ([docs][reset_docs]).
136///
137/// After this call, no auto-extensions will run for newly opened connections.
138///
139/// [reset_docs]: https://www.sqlite.org/c3ref/reset_auto_extension.html
140#[allow(unsafe_code)]
141pub fn reset_auto_extension() {
142    unsafe { ffi::sqlite3_reset_auto_extension() }
143}
144
145/// Returns the trampoline for `F`. `extension` is taken by value to infer `F`,
146/// then `forget`-ten so the zero-sized callback stays conceptually alive for the
147/// process and its destructor never runs.
148fn entry_point<F>(extension: F) -> RawAutoExtension
149where
150    F: Fn(&'_ mut BorrowedSqliteConnection) -> QueryResult<()> + Sync + 'static,
151{
152    core::mem::forget(extension);
153    trampoline::<F>
154}
155
156/// The C entry point handed to SQLite, monomorphized per callback type `F` so
157/// each distinct callback maps to a distinct, stable address. SQLite's
158/// pointer-based deduplication and [`cancel_auto_extension`] rely on that.
159#[allow(unsafe_code)]
160unsafe extern "C" fn trampoline<F>(
161    db: *mut ffi::sqlite3,
162    pz_err_msg: *mut *mut c_char,
163    _p_api: RawApiPointer,
164) -> c_int
165where
166    F: Fn(&'_ mut BorrowedSqliteConnection) -> QueryResult<()> + Sync + 'static,
167{
168    const {
169        if !(core::mem::size_of::<F>() == 0) {
    {
        ::core::panicking::panic_fmt(format_args!("an auto-extension callback must not capture non-zero-sized state. Use a `fn` item or a closure that captures only zero-sized values"));
    }
};assert!(
170            core::mem::size_of::<F>() == 0,
171            "an auto-extension callback must not capture non-zero-sized state. \
172             Use a `fn` item or a closure that captures only zero-sized values"
173        );
174    }
175
176    // `_p_api` matters only for runtime-loaded shared libraries. Statically
177    // linked extensions link the SQLite symbols directly, so we ignore it.
178    let result: Result<(), String> =
179        crate::util::std_compat::catch_unwind(core::panic::AssertUnwindSafe(|| {
180            let Some(db) = core::ptr::NonNull::new(db) else {
181                return Err(String::from(
182                    "auto-extension received a null database handle",
183                ));
184            };
185            // Reconstruct a *reference* to the zero-sized callback, never an
186            // owned value, so its destructor never runs. `&F: Fn` because `F: Fn`.
187            // SAFETY: `F` is zero-sized (asserted above), so a dangling, aligned,
188            // non-null pointer is a valid `&F`, which `NonNull::dangling` provides.
189            let extension: &F = unsafe { core::ptr::NonNull::<F>::dangling().as_ref() };
190            // SAFETY: `db` is a valid handle for the duration of this call, and
191            // the borrowed connection does not take ownership of it.
192            unsafe { SqliteConnection::with_borrowed_connection(db, extension) }
193                .map_err(|e| e.to_string())
194        }))
195        .unwrap_or_else(|panic| {
196            Err(match panic_detail(panic) {
197                Some(message) => ::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("auto-extension panicked: {0}",
                message))
    })alloc::format!("auto-extension panicked: {message}"),
198                None => String::from("auto-extension panicked"),
199            })
200        });
201
202    match result {
203        Ok(()) => ffi::SQLITE_OK,
204        Err(message) => {
205            set_error_message(pz_err_msg, &message);
206            ffi::SQLITE_ERROR
207        }
208    }
209}
210
211/// Best-effort message from a caught panic payload. The no_std `catch_unwind`
212/// carries no payload, so only the `std` variant can recover the text.
213#[cfg(feature = "std")]
214fn panic_detail(panic: alloc::boxed::Box<dyn core::any::Any + Send>) -> Option<String> {
215    panic
216        .downcast_ref::<&str>()
217        .map(|s| (*s).to_owned())
218        .or_else(|| panic.downcast_ref::<String>().cloned())
219}
220
221#[cfg(not(feature = "std"))]
222fn panic_detail(_panic: ()) -> Option<String> {
223    None
224}
225
226/// Writes `message` into `*pz_err_msg` with `sqlite3_malloc`, which is the
227/// allocator SQLite later frees it with. The message is truncated at the first
228/// NUL byte to form a valid C string, and allocation failure is ignored.
229#[allow(unsafe_code)]
230fn set_error_message(pz_err_msg: *mut *mut c_char, message: &str) {
231    if pz_err_msg.is_null() {
232        return;
233    }
234
235    let bytes = message.as_bytes();
236    let len = bytes.iter().position(|&b| b == 0).unwrap_or(bytes.len());
237
238    // SQLite sizes allocations with a C `int`. A message that does not fit is
239    // dropped rather than truncated to a bogus length.
240    let Ok(size) = c_int::try_from(len + 1) else {
241        return;
242    };
243    let buffer = unsafe { ffi::sqlite3_malloc(size) } as *mut u8;
244    if buffer.is_null() {
245        return;
246    }
247
248    unsafe {
249        core::ptr::copy_nonoverlapping(bytes.as_ptr(), buffer, len);
250        *buffer.add(len) = 0;
251        *pz_err_msg = buffer as *mut c_char;
252    }
253}
254
255// These tests either rely on sqlite calling a registered auto extension, or
256// read an error message allocated by the native library, neither of which is
257// supported when running under miri with a native libsqlite3
258// (`-Zmiri-native-lib`), so the whole module is compiled out in that case.
259#[cfg(all(test, not(miri)))]
260mod tests {
261    use super::*;
262    use crate::dsl::sql;
263    use crate::prelude::*;
264    use crate::sql_types::Integer;
265    use crate::sqlite::BorrowedSqliteConnection;
266    use std::sync::Mutex;
267
268    // `sqlite3_auto_extension` is process-global, so these tests serialize on
269    // this lock and register only benign (never-failing) extensions, leaving
270    // connections opened by other tests unaffected. The failing path is covered
271    // by `trampoline_maps_result_to_return_code`, which calls the trampoline
272    // directly without touching the global registry.
273    static AUTO_EXT_TEST_LOCK: Mutex<()> = Mutex::new(());
274
275    // A benign auto-extension: registers a `TESTCOLL` collation through the
276    // normal connection API.
277    #[allow(unsafe_code, reason = "Test call to inner unsafe function")]
278    fn test_ext_init(conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
279        unsafe {
280            conn.inner_mut()
281                .register_collation("TESTCOLL", |a, b| a.cmp(b))
282        }
283    }
284
285    fn open_memory_connection() -> SqliteConnection {
286        SqliteConnection::establish(":memory:").expect("Failed to open :memory: connection")
287    }
288
289    // Errors out if `TESTCOLL` is not registered on a freshly opened connection.
290    fn probe_collation() -> QueryResult<i32> {
291        let mut conn = open_memory_connection();
292        sql::<Integer>("SELECT 'a' = 'a' COLLATE TESTCOLL").get_result(&mut conn)
293    }
294
295    /// RAII guard that calls `reset_auto_extension()` on drop, ensuring global
296    /// state is cleaned up even if a test panics.
297    struct TestResetGuard;
298
299    impl Drop for TestResetGuard {
300        fn drop(&mut self) {
301            reset_auto_extension();
302        }
303    }
304
305    #[test]
306    fn auto_extension_lifecycle() {
307        let _lock = AUTO_EXT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
308        let _guard = TestResetGuard;
309        reset_auto_extension();
310
311        // -- 1. register + new connection has the collation --
312        register_auto_extension(test_ext_init).unwrap();
313        assert_eq!(probe_collation().unwrap(), 1);
314
315        // -- 2. cancel + new connection does NOT have the collation --
316        let removed = cancel_auto_extension(test_ext_init);
317        assert!(
318            removed,
319            "cancel should return true for registered extension"
320        );
321        assert!(
322            probe_collation().is_err(),
323            "collation should not be available after cancel"
324        );
325
326        // -- 3. cancel returns false for unregistered --
327        let removed = cancel_auto_extension(test_ext_init);
328        assert!(
329            !removed,
330            "cancel should return false for unregistered extension"
331        );
332
333        // -- 4. reset clears all --
334        register_auto_extension(test_ext_init).unwrap();
335        reset_auto_extension();
336        assert!(
337            probe_collation().is_err(),
338            "collation should not be available after reset"
339        );
340
341        // -- 5. duplicate registration is idempotent --
342        register_auto_extension(test_ext_init).unwrap();
343        register_auto_extension(test_ext_init).unwrap();
344        assert_eq!(probe_collation().unwrap(), 1);
345        // _guard drops here, ensuring reset even on panic.
346    }
347
348    // Drives the trampoline directly (via `entry_point`), without registering
349    // it in SQLite's global list, so the Ok/Err/null paths can be checked
350    // deterministically without affecting connections opened by other tests.
351    #[test]
352    #[allow(unsafe_code)]
353    fn trampoline_maps_result_to_return_code() {
354        fn ok_ext(_conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
355            Ok(())
356        }
357        fn err_ext(_conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
358            Err(Error::QueryBuilderError("boom".into()))
359        }
360
361        let ok_tramp = entry_point(ok_ext);
362        let err_tramp = entry_point(err_ext);
363
364        let mut conn = open_memory_connection();
365        // SAFETY: the pointer is only used for the duration of the closure,
366        // while `conn` is alive.
367        unsafe {
368            conn.with_raw_connection(|db| {
369                let mut err: *mut c_char = core::ptr::null_mut();
370
371                // Ok -> SQLITE_OK, no error message allocated.
372                let rc = ok_tramp(db, &mut err, core::ptr::null());
373                assert_eq!(rc, ffi::SQLITE_OK);
374                assert!(err.is_null());
375
376                // Err -> SQLITE_ERROR, message written via sqlite3_malloc.
377                let rc = err_tramp(db, &mut err, core::ptr::null());
378                assert_eq!(rc, ffi::SQLITE_ERROR);
379                assert!(!err.is_null());
380                let message = core::ffi::CStr::from_ptr(err)
381                    .to_string_lossy()
382                    .into_owned();
383                assert_eq!(message, "boom");
384                ffi::sqlite3_free(err as *mut core::ffi::c_void);
385
386                // Null db handle -> SQLITE_ERROR, never dereferenced.
387                let mut err: *mut c_char = core::ptr::null_mut();
388                let rc = ok_tramp(core::ptr::null_mut(), &mut err, core::ptr::null());
389                assert_eq!(rc, ffi::SQLITE_ERROR);
390                if !err.is_null() {
391                    ffi::sqlite3_free(err as *mut core::ffi::c_void);
392                }
393            })
394        }
395    }
396
397    // Regression test for the closure-reconstruction soundness fix. The callback
398    // captures a zero-sized guard with a load-bearing `Drop` that must never run,
399    // because the trampoline only reproduces the callback behind a reference and
400    // `forget`s the registered value (the old `mem::zeroed()` ran it repeatedly).
401    #[test]
402    fn callback_zero_sized_capture_is_never_dropped() {
403        use std::sync::atomic::{AtomicUsize, Ordering};
404
405        static DROPS: AtomicUsize = AtomicUsize::new(0);
406
407        // Zero-sized, `Sync`, with a load-bearing `Drop`.
408        struct Guard;
409        impl Drop for Guard {
410            fn drop(&mut self) {
411                DROPS.fetch_add(1, Ordering::SeqCst);
412            }
413        }
414
415        let _lock = AUTO_EXT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
416        let _reset = TestResetGuard;
417        reset_auto_extension();
418
419        let guard = Guard;
420        // `move` captures the zero-sized `guard` by value, so the closure is a
421        // zero-sized type with a non-trivial destructor.
422        #[allow(unsafe_code, reason = "Test call to inner unsafe function")]
423        register_auto_extension(move |conn: &mut BorrowedSqliteConnection| {
424            let _ = &guard;
425            unsafe {
426                conn.inner_mut()
427                    .register_collation("TESTCOLL", |a, b| a.cmp(b))
428            }
429        })
430        .unwrap();
431
432        for _ in 0..5 {
433            assert_eq!(probe_collation().unwrap(), 1);
434        }
435        reset_auto_extension();
436
437        assert_eq!(
438            DROPS.load(Ordering::SeqCst),
439            0,
440            "the captured guard's destructor must never run"
441        );
442    }
443
444    #[test]
445    fn leaked_failed_cursor_in_auto_extension_fails_open_cleanly() {
446        use std::cell::Cell;
447
448        // Connections other tests open concurrently skip the extension.
449        thread_local!(static ARMED: Cell<bool> = const { Cell::new(false) });
450
451        #[allow(unsafe_code, reason = "Test code")]
452        fn extension(conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
453            if ARMED.get() {
454                unsafe {
455                    crate::sqlite::connection::leak_failed_cursor(conn.inner_mut());
456                }
457            }
458            Ok(())
459        }
460
461        let _lock = AUTO_EXT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
462        let _reset = TestResetGuard;
463        register_auto_extension(extension).unwrap();
464        ARMED.set(true);
465
466        // `sqlite3_open_v2` reports the error the finalized statement left on the handle.
467        let conn = SqliteConnection::establish(":memory:");
468        assert!(matches!(conn, Err(ConnectionError::BadConnection(_))));
469    }
470
471    #[test]
472    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))] // there is no catch unwind for wasm
473    fn panic_in_borrowed_connection_teardown_keeps_handle_open() {
474        use crate::connection::{Instrumentation, InstrumentationEvent};
475        use std::cell::Cell;
476
477        // Connections other tests open concurrently skip the extension.
478        thread_local!(static ARMED: Cell<bool> = const { Cell::new(false) });
479
480        struct PanicOnDrop;
481
482        impl Instrumentation for PanicOnDrop {
483            fn on_connection_event(&mut self, _event: InstrumentationEvent<'_>) {}
484        }
485
486        impl Drop for PanicOnDrop {
487            fn drop(&mut self) {
488                if ARMED.get() {
489                    panic!("instrumentation dropped");
490                }
491            }
492        }
493
494        fn extension(conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
495            if ARMED.get() {
496                conn.set_instrumentation(PanicOnDrop);
497            }
498            Ok(())
499        }
500
501        let _lock = AUTO_EXT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
502        let _reset = TestResetGuard;
503        register_auto_extension(extension).unwrap();
504        ARMED.set(true);
505        let conn = SqliteConnection::establish(":memory:");
506        ARMED.set(false);
507
508        // The caught panic fails the open, and the handle is closed exactly once.
509        assert!(matches!(conn, Err(ConnectionError::BadConnection(_))));
510    }
511
512    // Counts calls and destructor runs of the hook closures registered by
513    // `register_all_hooks`. Tests using them hold `AUTO_EXT_TEST_LOCK`.
514    mod hook_tokens {
515        use std::sync::atomic::{AtomicUsize, Ordering};
516        pub(super) static FIRED: AtomicUsize = AtomicUsize::new(0);
517        pub(super) static DROPPED: AtomicUsize = AtomicUsize::new(0);
518        /// Dropped together with its hook closure; panics on drop if the flag is set.
519        pub(super) struct Token(pub(super) bool);
520        impl Token {
521            pub(super) fn fire(&self) {
522                FIRED.fetch_add(1, Ordering::SeqCst);
523            }
524        }
525        impl Drop for Token {
526            fn drop(&mut self) {
527                DROPPED.fetch_add(1, Ordering::SeqCst);
528                if self.0 {
529                    panic!("hook closure dropped");
530                }
531            }
532        }
533    }
534
535    // Registers all nine hooks through the borrowed connection. The closure of hook
536    // number `panicking` (in `RawConnection` field order) panics when dropped.
537    #[allow(unsafe_code)]
538    fn register_all_hooks(conn: &mut BorrowedSqliteConnection, panicking: usize) {
539        use crate::sqlite::{
540            AuthorizerDecision, BusyDecision, CommitDecision, ProgressDecision, SqliteChangeOps,
541            SqliteTraceFlags, SqliteUpdateRouter,
542        };
543        use core::num::NonZeroU32;
544        use hook_tokens::Token;
545        let conn = unsafe { conn.inner_mut() };
546        let t = Token(panicking == 0);
547        conn.on_update(SqliteUpdateRouter::new().on_any(SqliteChangeOps::ALL, move |_| t.fire()));
548        let t = Token(panicking == 1);
549        conn.on_commit(move || {
550            t.fire();
551            CommitDecision::Proceed
552        });
553        let t = Token(panicking == 2);
554        conn.on_rollback(move || t.fire());
555        let t = Token(panicking == 3);
556        conn.on_progress(NonZeroU32::new(1).unwrap(), move || {
557            t.fire();
558            ProgressDecision::Continue
559        });
560        let t = Token(panicking == 4);
561        conn.on_wal(move |_, _, _| t.fire());
562        let t = Token(panicking == 5);
563        conn.on_busy(move |_| {
564            t.fire();
565            BusyDecision::GiveUp
566        });
567        let t = Token(panicking == 6);
568        conn.on_authorize(move |_| {
569            t.fire();
570            AuthorizerDecision::Allow
571        });
572        let t = Token(panicking == 7);
573        conn.on_trace(SqliteTraceFlags::all(), move |_| t.fire());
574        let t = Token(panicking == 8);
575        conn.on_collation_needed(move |_| {
576            t.fire();
577            None
578        });
579    }
580
581    // Runs statements that would trigger the hooks `register_all_hooks` registers,
582    // except the busy handler and the WAL hook (those need lock contention or WAL mode).
583    fn trigger_hooks(conn: &mut SqliteConnection) {
584        use crate::connection::SimpleConnection;
585        conn.batch_execute(
586            "CREATE TABLE IF NOT EXISTS t (x); INSERT INTO t VALUES (1);
587             BEGIN; INSERT INTO t VALUES (2); ROLLBACK;",
588        )
589        .unwrap();
590        let _ = conn.batch_execute("SELECT 'a' = 'b' COLLATE nosuch");
591    }
592
593    // All nine hooks an auto-extension registers are unregistered and freed.
594    #[test]
595    fn hooks_registered_in_auto_extension_are_unregistered_and_freed() {
596        use std::cell::Cell;
597        use std::sync::atomic::Ordering;
598        thread_local!(static ARMED: Cell<bool> = const { Cell::new(false) });
599        fn extension(conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
600            if ARMED.get() {
601                register_all_hooks(conn, usize::MAX);
602            }
603            Ok(())
604        }
605        let _lock = AUTO_EXT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
606        let _reset = TestResetGuard;
607        hook_tokens::FIRED.store(0, Ordering::SeqCst);
608        hook_tokens::DROPPED.store(0, Ordering::SeqCst);
609        register_auto_extension(extension).unwrap();
610        ARMED.set(true);
611        let conn = SqliteConnection::establish(":memory:");
612        ARMED.set(false);
613        let mut conn = conn.unwrap();
614        assert_eq!(hook_tokens::DROPPED.load(Ordering::SeqCst), 9);
615        trigger_hooks(&mut conn);
616        assert_eq!(hook_tokens::FIRED.load(Ordering::SeqCst), 0);
617    }
618
619    // A panic while freeing hook number `n` must not close the handle, and must
620    // not leak the closures of the other hooks. Gated off WASM, where
621    // `catch_unwind` aborts because `panic = "abort"`.
622    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
623    fn assert_panicking_hook_drop_frees_all_hooks(n: usize) {
624        use std::cell::Cell;
625        use std::sync::atomic::Ordering;
626        thread_local!(static PANIC_AT: Cell<usize> = const { Cell::new(usize::MAX) });
627        thread_local!(static ARMED: Cell<bool> = const { Cell::new(false) });
628        fn extension(conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
629            if ARMED.get() {
630                register_all_hooks(conn, PANIC_AT.get());
631            }
632            Ok(())
633        }
634        let _lock = AUTO_EXT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
635        let _reset = TestResetGuard;
636        hook_tokens::DROPPED.store(0, Ordering::SeqCst);
637        register_auto_extension(extension).unwrap();
638        PANIC_AT.set(n);
639        ARMED.set(true);
640        let conn = SqliteConnection::establish(":memory:");
641        ARMED.set(false);
642        assert!(matches!(conn, Err(ConnectionError::BadConnection(_))));
643        assert_eq!(hook_tokens::DROPPED.load(Ordering::SeqCst), 9);
644    }
645
646    #[test]
647    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
648    fn panicking_first_hook_drop_frees_all_hooks() {
649        assert_panicking_hook_drop_frees_all_hooks(0);
650    }
651
652    #[test]
653    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
654    fn panicking_last_hook_drop_frees_all_hooks() {
655        assert_panicking_hook_drop_frees_all_hooks(8);
656    }
657
658    // A panicking callback becomes `SQLITE_ERROR` with the payload recovered
659    // into the message (`&str` and `String` payloads, generic fallback
660    // otherwise), and drives the panic-unwind path through the drop guard. Gated
661    // off WASM, where `catch_unwind` aborts because `panic = "abort"`.
662    #[test]
663    #[allow(unsafe_code)]
664    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
665    fn trampoline_reports_panic_message() {
666        fn panic_str(_conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
667            panic!("boom-str");
668        }
669        fn panic_string(_conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
670            panic!("boom-{}", 7);
671        }
672        fn panic_other(_conn: &mut BorrowedSqliteConnection) -> QueryResult<()> {
673            std::panic::panic_any(7_u8);
674        }
675
676        let cases: [(RawAutoExtension, &str); 3] = [
677            (entry_point(panic_str), "auto-extension panicked: boom-str"),
678            (entry_point(panic_string), "auto-extension panicked: boom-7"),
679            (entry_point(panic_other), "auto-extension panicked"),
680        ];
681
682        let mut conn = open_memory_connection();
683        // SAFETY: the pointer is only used while `conn` is alive.
684        unsafe {
685            conn.with_raw_connection(|db| {
686                for (tramp, expected) in cases {
687                    let mut err: *mut c_char = core::ptr::null_mut();
688                    let rc = tramp(db, &mut err, core::ptr::null());
689                    assert_eq!(rc, ffi::SQLITE_ERROR);
690                    assert!(!err.is_null());
691                    let message = core::ffi::CStr::from_ptr(err)
692                        .to_string_lossy()
693                        .into_owned();
694                    assert_eq!(message, expected);
695                    ffi::sqlite3_free(err as *mut core::ffi::c_void);
696                }
697            })
698        }
699    }
700
701    #[test]
702    #[allow(unsafe_code)]
703    fn error_message_truncates_at_interior_nul() {
704        let mut err: *mut c_char = core::ptr::null_mut();
705        set_error_message(&mut err, "before\0after");
706        assert!(!err.is_null());
707        // SAFETY: `err` is a sqlite-allocated C string we own until we free it.
708        unsafe {
709            let truncated = core::ffi::CStr::from_ptr(err).to_str().unwrap();
710            assert_eq!(truncated, "before");
711            ffi::sqlite3_free(err as *mut core::ffi::c_void);
712        }
713
714        // A null out-pointer is a no-op (must not write through it or crash).
715        set_error_message(core::ptr::null_mut(), "ignored");
716    }
717}
diesel/sqlite/auto_extension.rs

diesel/sqlite/
auto_extension.rs
