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

diesel/sqlite/
auto_extension.rs
