Repository: launchbadge/sqlx
File: sqlx-sqlite/src/connection/mod.rs (Rust)

Code before:

    /// Sets a progress handler that is invoked periodically during long running calls. If the progress callback
    /// returns `false`, then the operation is interrupted.
    ///
    /// `num_ops` is the approximate number of [virtual machine instructions](https://www.sqlite.org/opcode.html)
    /// that are evaluated between successive invocations of the callback. If `num_ops` is less than one then the
    /// progress handler is disabled.
    ///
    /// Only a single progress handler may be defined at one time per database connection; setting a new progress
    /// handler cancels the old one.
    ///
    /// The progress handler callback must not do anything that will modify the database connection that invoked
    /// the progress handler. Note that sqlite3_prepare_v2() and sqlite3_step() both modify their database connections
    /// in this context.
    pub async fn set_progress_handler<F>(&mut self, num_ops: i32, callback: F)
    where
        F: FnMut() -> bool + Send + 'static,
    {
        unsafe {
            let callback = Box::new(callback);
            if let Ok(mut lock_conn) = self.lock_handle().await {
                sqlite3_progress_handler(
                    lock_conn.as_raw_handle().as_mut(),
                    num_ops,
                    Some(progress_callback::<F>),
                    &*callback as *const F as *mut F as *mut _,
                );
            }
        }
    }

    /// Removes a previously set progress handler on a database connection.
    pub async fn remove_progress_handler(&mut self) {
        unsafe {
            if let Ok(mut lock_conn) = self.lock_handle().await {
                sqlite3_progress_handler(lock_conn.as_raw_handle().as_mut(), 0, None, 0 as *mut _);
            }
        }
    }
}

impl Debug for SqliteConnection {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("SqliteConnection")
            .field("row_channel_size", &self.row_channel_size)
            .field("cached_statements_size", &self.cached_statements_size())
            .finish()
    }
}

impl Connection for SqliteConnection {

Review comment (bug):
:stop_sign: This is undefined behavior: the `Box` and its contents will be dropped at the end of the function, leading to a use-after free when the handler is invoked. It doesn't segfault in your simple test because the callback is zero-sized, so `Box::new()` doesn't actually allocate.

You should use [`Box::leak()`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html#method.leak) to ensure the allocation is not freed prematurely.

At the same time, the function should be _very_ clear that **this means leaking memory** as there's no way to give SQLite a callback to free the pointer. Because of this, I would almost just take a `fn()` pointer directly instead.

Alternatively, you could store the `Box` in the `ConnectionState` struct so that it can be freed when appropriate, but be very careful with regards to drop order, as you _must_ ensure it's not dropped before the `ConnectionHandle`: https://doc.rust-lang.org/reference/destructors.html

Code after:

impl SqliteConnection {
    pub(crate) async fn establish(options: &SqliteConnectOptions) -> Result<Self, Error> {
        let params = EstablishParams::from_options(options)?;
        let worker = ConnectionWorker::establish(params).await?;
        Ok(Self {
            optimize_on_close: options.optimize_on_close.clone(),
            worker,
            row_channel_size: options.row_channel_size,
        })
    }

    /// Lock the SQLite database handle out from the worker thread so direct SQLite API calls can
    /// be made safely.
    ///
    /// Returns an error if the worker thread crashed.
    pub async fn lock_handle(&mut self) -> Result<LockedSqliteHandle<'_>, Error> {
        let guard = self.worker.unlock_db().await?;

        Ok(LockedSqliteHandle { guard })
    }
}

impl Debug for SqliteConnection {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("SqliteConnection")
            .field("row_channel_size", &self.row_channel_size)
            .field("cached_statements_size", &self.cached_statements_size())
            .finish()
    }
}

impl Connection for SqliteConnection {
    type Database = Sqlite;

    type Options = SqliteConnectOptions;

    fn close(mut self) -> BoxFuture<'static, Result<(), Error>> {
        Box::pin(async move {
            if let OptimizeOnClose::Enabled { analysis_limit } = self.optimize_on_close {
                let mut pragma_string = String::new();
                if let Some(limit) = analysis_limit {
                    write!(pragma_string, "PRAGMA analysis_limit = {}; ", limit).ok();
                }
                pragma_string.push_str("PRAGMA optimize;");
                self.execute(&*pragma_string).await?;
            }
            let shutdown = self.worker.shutdown();
            // Drop the statement worker, which should
            // cover all references to the connection handle outside of the worker thread
            drop(self);