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
41/// Registers an auto-extension that runs for every SQLite connection opened in
42/// this process, including non-Diesel ones.
43///
44/// This is a safe wrapper around [`sqlite3_auto_extension`][docs]. The callback
45/// receives the [`SqliteConnection`] being opened and returns `Ok(())` to
46/// continue or an error to fail the open. Use it to register SQL functions,
47/// collations, or aggregates through the usual connection API, or to initialize
48/// a statically linked C extension such as Spatialite or sqlite-vec via
49/// [`SqliteConnection::with_raw_connection`].
50///
51/// Call this before opening any connection. Extensions run in registration
52/// order, and the first error aborts the open. The callback must be a `fn` item
53/// or a closure that captures only zero-sized values (enforced at compile
54/// time), and registering the same `fn` twice is a no-op. It may run on several
55/// threads at once and must not open another connection (which would re-enter
56/// the auto-extensions and recurse) or call [`register_auto_extension`],
57/// [`cancel_auto_extension`], or [`reset_auto_extension`]. Panics are caught and
58/// turned into a failed open.
59///
60/// [docs]: https://www.sqlite.org/c3ref/auto_extension.html
61///
62/// # Example
63///
64/// ```rust
65/// use diesel::dsl::sql;
66/// use diesel::prelude::*;
67/// use diesel::sql_types::Integer;
68/// use diesel::sqlite::{register_auto_extension, reset_auto_extension, SqliteConnection};
69///
70/// // Registers a case-insensitive collation on every new connection.
71/// fn my_ext(conn: &mut SqliteConnection) -> QueryResult<()> {
72///     conn.register_collation("RUSTNOCASE", |a, b| a.to_lowercase().cmp(&b.to_lowercase()))
73/// }
74///
75/// register_auto_extension(my_ext).unwrap();
76///
77/// // Every future connection now has the collation.
78/// let mut conn = SqliteConnection::establish(":memory:").unwrap();
79/// let equal: i32 = sql::<Integer>("SELECT 'a' = 'A' COLLATE RUSTNOCASE")
80///     .get_result(&mut conn)
81///     .unwrap();
82/// assert_eq!(equal, 1);
83/// # reset_auto_extension();
84/// ```
85#[allow(unsafe_code)]
86pub fn register_auto_extension<F>(extension: F) -> QueryResult<()>
87where
88    F: Fn(&mut SqliteConnection) -> QueryResult<()> + Sync + 'static,
89{
90    // SAFETY: `entry_point` returns a stable function pointer with SQLite's callback ABI.
91    let result =
92        unsafe { auto_extension_ffi::sqlite3_auto_extension(Some(entry_point(extension))) };
93    if result == ffi::SQLITE_OK {
94        Ok(())
95    } else {
96        Err(DatabaseError(
97            DatabaseErrorKind::Unknown,
98            Box::new(ffi::code_to_str(result).to_string()),
99        ))
100    }
101}
102
103/// Removes a previously registered auto-extension, returning `true` if it was
104/// found ([docs][cancel_docs]).
105///
106/// Pass the same `fn` item given to [`register_auto_extension`]. A closure
107/// cannot be cancelled this way, because its type cannot be named again. Use
108/// [`reset_auto_extension`] to clear everything instead.
109///
110/// [cancel_docs]: https://www.sqlite.org/c3ref/cancel_auto_extension.html
111#[allow(unsafe_code)]
112pub fn cancel_auto_extension<F>(extension: F) -> bool
113where
114    F: Fn(&mut SqliteConnection) -> QueryResult<()> + Sync + 'static,
115{
116    // SAFETY: `entry_point` returns the same stable pointer used for registration.
117    unsafe { auto_extension_ffi::sqlite3_cancel_auto_extension(Some(entry_point(extension))) != 0 }
118}
119
120/// Clears **all** registered auto-extensions ([docs][reset_docs]).
121///
122/// After this call, no auto-extensions will run for newly opened connections.
123///
124/// [reset_docs]: https://www.sqlite.org/c3ref/reset_auto_extension.html
125#[allow(unsafe_code)]
126pub fn reset_auto_extension() {
127    unsafe { ffi::sqlite3_reset_auto_extension() }
128}
129
130/// Returns the trampoline for `F`. `extension` is taken by value to infer `F`,
131/// then `forget`-ten so the zero-sized callback stays conceptually alive for the
132/// process and its destructor never runs.
133fn entry_point<F>(extension: F) -> RawAutoExtension
134where
135    F: Fn(&mut SqliteConnection) -> QueryResult<()> + Sync + 'static,
136{
137    core::mem::forget(extension);
138    trampoline::<F>
139}
140
141/// The C entry point handed to SQLite, monomorphized per callback type `F` so
142/// each distinct callback maps to a distinct, stable address. SQLite's
143/// pointer-based deduplication and [`cancel_auto_extension`] rely on that.
144#[allow(unsafe_code)]
145unsafe extern "C" fn trampoline<F>(
146    db: *mut ffi::sqlite3,
147    pz_err_msg: *mut *mut c_char,
148    _p_api: RawApiPointer,
149) -> c_int
150where
151    F: Fn(&mut SqliteConnection) -> QueryResult<()> + Sync + 'static,
152{
153    const {
154        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!(
155            core::mem::size_of::<F>() == 0,
156            "an auto-extension callback must not capture non-zero-sized state. \
157             Use a `fn` item or a closure that captures only zero-sized values"
158        );
159    }
160
161    // `_p_api` matters only for runtime-loaded shared libraries. Statically
162    // linked extensions link the SQLite symbols directly, so we ignore it.
163    let result: Result<(), String> =
164        crate::util::std_compat::catch_unwind(core::panic::AssertUnwindSafe(|| {
165            let Some(db) = core::ptr::NonNull::new(db) else {
166                return Err(String::from(
167                    "auto-extension received a null database handle",
168                ));
169            };
170            // Reconstruct a *reference* to the zero-sized callback, never an
171            // owned value, so its destructor never runs. `&F: Fn` because `F: Fn`.
172            // SAFETY: `F` is zero-sized (asserted above), so a dangling, aligned,
173            // non-null pointer is a valid `&F`, which `NonNull::dangling` provides.
174            let extension: &F = unsafe { core::ptr::NonNull::<F>::dangling().as_ref() };
175            // SAFETY: `db` is a valid handle for the duration of this call, and
176            // the borrowed connection does not take ownership of it.
177            unsafe { SqliteConnection::with_borrowed_connection(db, extension) }
178                .map_err(|e| e.to_string())
179        }))
180        .unwrap_or_else(|panic| {
181            Err(match panic_detail(panic) {
182                Some(message) => ::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("auto-extension panicked: {0}",
                message))
    })alloc::format!("auto-extension panicked: {message}"),
183                None => String::from("auto-extension panicked"),
184            })
185        });
186
187    match result {
188        Ok(()) => ffi::SQLITE_OK,
189        Err(message) => {
190            set_error_message(pz_err_msg, &message);
191            ffi::SQLITE_ERROR
192        }
193    }
194}
195
196/// Best-effort message from a caught panic payload. The no_std `catch_unwind`
197/// carries no payload, so only the `std` variant can recover the text.
198#[cfg(feature = "std")]
199fn panic_detail(panic: alloc::boxed::Box<dyn core::any::Any + Send>) -> Option<String> {
200    panic
201        .downcast_ref::<&str>()
202        .map(|s| (*s).to_owned())
203        .or_else(|| panic.downcast_ref::<String>().cloned())
204}
205
206#[cfg(not(feature = "std"))]
207fn panic_detail(_panic: ()) -> Option<String> {
208    None
209}
210
211/// Writes `message` into `*pz_err_msg` with `sqlite3_malloc`, which is the
212/// allocator SQLite later frees it with. The message is truncated at the first
213/// NUL byte to form a valid C string, and allocation failure is ignored.
214#[allow(unsafe_code)]
215fn set_error_message(pz_err_msg: *mut *mut c_char, message: &str) {
216    if pz_err_msg.is_null() {
217        return;
218    }
219
220    let bytes = message.as_bytes();
221    let len = bytes.iter().position(|&b| b == 0).unwrap_or(bytes.len());
222
223    // SQLite sizes allocations with a C `int`. A message that does not fit is
224    // dropped rather than truncated to a bogus length.
225    let Ok(size) = c_int::try_from(len + 1) else {
226        return;
227    };
228    let buffer = unsafe { ffi::sqlite3_malloc(size) } as *mut u8;
229    if buffer.is_null() {
230        return;
231    }
232
233    unsafe {
234        core::ptr::copy_nonoverlapping(bytes.as_ptr(), buffer, len);
235        *buffer.add(len) = 0;
236        *pz_err_msg = buffer as *mut c_char;
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243    use crate::dsl::sql;
244    use crate::prelude::*;
245    use crate::sql_types::Integer;
246    use std::sync::Mutex;
247
248    // `sqlite3_auto_extension` is process-global, so these tests serialize on
249    // this lock and register only benign (never-failing) extensions, leaving
250    // connections opened by other tests unaffected. The failing path is covered
251    // by `trampoline_maps_result_to_return_code`, which calls the trampoline
252    // directly without touching the global registry.
253    static AUTO_EXT_TEST_LOCK: Mutex<()> = Mutex::new(());
254
255    // A benign auto-extension: registers a `TESTCOLL` collation through the
256    // normal connection API.
257    fn test_ext_init(conn: &mut SqliteConnection) -> QueryResult<()> {
258        conn.register_collation("TESTCOLL", |a, b| a.cmp(b))
259    }
260
261    fn open_memory_connection() -> SqliteConnection {
262        SqliteConnection::establish(":memory:").expect("Failed to open :memory: connection")
263    }
264
265    // Errors out if `TESTCOLL` is not registered on a freshly opened connection.
266    fn probe_collation() -> QueryResult<i32> {
267        let mut conn = open_memory_connection();
268        sql::<Integer>("SELECT 'a' = 'a' COLLATE TESTCOLL").get_result(&mut conn)
269    }
270
271    /// RAII guard that calls `reset_auto_extension()` on drop, ensuring global
272    /// state is cleaned up even if a test panics.
273    struct TestResetGuard;
274
275    impl Drop for TestResetGuard {
276        fn drop(&mut self) {
277            reset_auto_extension();
278        }
279    }
280
281    #[test]
282    fn auto_extension_lifecycle() {
283        let _lock = AUTO_EXT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
284        let _guard = TestResetGuard;
285        reset_auto_extension();
286
287        // -- 1. register + new connection has the collation --
288        register_auto_extension(test_ext_init).unwrap();
289        assert_eq!(probe_collation().unwrap(), 1);
290
291        // -- 2. cancel + new connection does NOT have the collation --
292        let removed = cancel_auto_extension(test_ext_init);
293        assert!(
294            removed,
295            "cancel should return true for registered extension"
296        );
297        assert!(
298            probe_collation().is_err(),
299            "collation should not be available after cancel"
300        );
301
302        // -- 3. cancel returns false for unregistered --
303        let removed = cancel_auto_extension(test_ext_init);
304        assert!(
305            !removed,
306            "cancel should return false for unregistered extension"
307        );
308
309        // -- 4. reset clears all --
310        register_auto_extension(test_ext_init).unwrap();
311        reset_auto_extension();
312        assert!(
313            probe_collation().is_err(),
314            "collation should not be available after reset"
315        );
316
317        // -- 5. duplicate registration is idempotent --
318        register_auto_extension(test_ext_init).unwrap();
319        register_auto_extension(test_ext_init).unwrap();
320        assert_eq!(probe_collation().unwrap(), 1);
321        // _guard drops here, ensuring reset even on panic.
322    }
323
324    // Drives the trampoline directly (via `entry_point`), without registering
325    // it in SQLite's global list, so the Ok/Err/null paths can be checked
326    // deterministically without affecting connections opened by other tests.
327    #[test]
328    #[allow(unsafe_code)]
329    fn trampoline_maps_result_to_return_code() {
330        fn ok_ext(_conn: &mut SqliteConnection) -> QueryResult<()> {
331            Ok(())
332        }
333        fn err_ext(_conn: &mut SqliteConnection) -> QueryResult<()> {
334            Err(Error::QueryBuilderError("boom".into()))
335        }
336
337        let ok_tramp = entry_point(ok_ext);
338        let err_tramp = entry_point(err_ext);
339
340        let mut conn = open_memory_connection();
341        // SAFETY: the pointer is only used for the duration of the closure,
342        // while `conn` is alive.
343        unsafe {
344            conn.with_raw_connection(|db| {
345                let mut err: *mut c_char = core::ptr::null_mut();
346
347                // Ok -> SQLITE_OK, no error message allocated.
348                let rc = ok_tramp(db, &mut err, core::ptr::null());
349                assert_eq!(rc, ffi::SQLITE_OK);
350                assert!(err.is_null());
351
352                // Err -> SQLITE_ERROR, message written via sqlite3_malloc.
353                let rc = err_tramp(db, &mut err, core::ptr::null());
354                assert_eq!(rc, ffi::SQLITE_ERROR);
355                assert!(!err.is_null());
356                let message = core::ffi::CStr::from_ptr(err)
357                    .to_string_lossy()
358                    .into_owned();
359                assert_eq!(message, "boom");
360                ffi::sqlite3_free(err as *mut core::ffi::c_void);
361
362                // Null db handle -> SQLITE_ERROR, never dereferenced.
363                let mut err: *mut c_char = core::ptr::null_mut();
364                let rc = ok_tramp(core::ptr::null_mut(), &mut err, core::ptr::null());
365                assert_eq!(rc, ffi::SQLITE_ERROR);
366                if !err.is_null() {
367                    ffi::sqlite3_free(err as *mut core::ffi::c_void);
368                }
369            })
370        }
371    }
372
373    // Regression test for the closure-reconstruction soundness fix. The callback
374    // captures a zero-sized guard with a load-bearing `Drop` that must never run,
375    // because the trampoline only reproduces the callback behind a reference and
376    // `forget`s the registered value (the old `mem::zeroed()` ran it repeatedly).
377    #[test]
378    fn callback_zero_sized_capture_is_never_dropped() {
379        use std::sync::atomic::{AtomicUsize, Ordering};
380
381        static DROPS: AtomicUsize = AtomicUsize::new(0);
382
383        // Zero-sized, `Sync`, with a load-bearing `Drop`.
384        struct Guard;
385        impl Drop for Guard {
386            fn drop(&mut self) {
387                DROPS.fetch_add(1, Ordering::SeqCst);
388            }
389        }
390
391        let _lock = AUTO_EXT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
392        let _reset = TestResetGuard;
393        reset_auto_extension();
394
395        let guard = Guard;
396        // `move` captures the zero-sized `guard` by value, so the closure is a
397        // zero-sized type with a non-trivial destructor.
398        register_auto_extension(move |conn: &mut SqliteConnection| {
399            let _ = &guard;
400            conn.register_collation("TESTCOLL", |a, b| a.cmp(b))
401        })
402        .unwrap();
403
404        for _ in 0..5 {
405            assert_eq!(probe_collation().unwrap(), 1);
406        }
407        reset_auto_extension();
408
409        assert_eq!(
410            DROPS.load(Ordering::SeqCst),
411            0,
412            "the captured guard's destructor must never run"
413        );
414    }
415
416    // A panicking callback becomes `SQLITE_ERROR` with the payload recovered
417    // into the message (`&str` and `String` payloads, generic fallback
418    // otherwise), and drives the panic-unwind path through the drop guard. Gated
419    // off WASM, where `catch_unwind` aborts because `panic = "abort"`.
420    #[test]
421    #[allow(unsafe_code)]
422    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
423    fn trampoline_reports_panic_message() {
424        fn panic_str(_conn: &mut SqliteConnection) -> QueryResult<()> {
425            panic!("boom-str");
426        }
427        fn panic_string(_conn: &mut SqliteConnection) -> QueryResult<()> {
428            panic!("boom-{}", 7);
429        }
430        fn panic_other(_conn: &mut SqliteConnection) -> QueryResult<()> {
431            std::panic::panic_any(7_u8);
432        }
433
434        let cases: [(RawAutoExtension, &str); 3] = [
435            (entry_point(panic_str), "auto-extension panicked: boom-str"),
436            (entry_point(panic_string), "auto-extension panicked: boom-7"),
437            (entry_point(panic_other), "auto-extension panicked"),
438        ];
439
440        let mut conn = open_memory_connection();
441        // SAFETY: the pointer is only used while `conn` is alive.
442        unsafe {
443            conn.with_raw_connection(|db| {
444                for (tramp, expected) in cases {
445                    let mut err: *mut c_char = core::ptr::null_mut();
446                    let rc = tramp(db, &mut err, core::ptr::null());
447                    assert_eq!(rc, ffi::SQLITE_ERROR);
448                    assert!(!err.is_null());
449                    let message = core::ffi::CStr::from_ptr(err)
450                        .to_string_lossy()
451                        .into_owned();
452                    assert_eq!(message, expected);
453                    ffi::sqlite3_free(err as *mut core::ffi::c_void);
454                }
455            })
456        }
457    }
458
459    #[test]
460    #[allow(unsafe_code)]
461    fn error_message_truncates_at_interior_nul() {
462        let mut err: *mut c_char = core::ptr::null_mut();
463        set_error_message(&mut err, "before\0after");
464        assert!(!err.is_null());
465        // SAFETY: `err` is a sqlite-allocated C string we own until we free it.
466        unsafe {
467            let truncated = core::ffi::CStr::from_ptr(err).to_str().unwrap();
468            assert_eq!(truncated, "before");
469            ffi::sqlite3_free(err as *mut core::ffi::c_void);
470        }
471
472        // A null out-pointer is a no-op (must not write through it or crash).
473        set_error_message(core::ptr::null_mut(), "ignored");
474    }
475}
diesel/sqlite/auto_extension.rs

diesel/sqlite/
auto_extension.rs
