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

Code before:
use crate::connection::LogSettings;
use crate::error::Error;
use crate::sqlite::connection::handle::ConnectionHandle;
use crate::sqlite::connection::{ConnectionState, Statements};
use crate::sqlite::{SqliteConnectOptions, SqliteError};
use indexmap::IndexMap;
use libc::c_void;
use libsqlite3_sys::{
    sqlite3_busy_timeout, sqlite3_db_config, sqlite3_extended_result_codes, sqlite3_free,
    sqlite3_load_extension, sqlite3_open_v2, SQLITE_DBCONFIG_ENABLE_LOAD_EXTENSION, SQLITE_OK,
    SQLITE_OPEN_CREATE, SQLITE_OPEN_FULLMUTEX, SQLITE_OPEN_MEMORY, SQLITE_OPEN_NOMUTEX,
    SQLITE_OPEN_PRIVATECACHE, SQLITE_OPEN_READONLY, SQLITE_OPEN_READWRITE, SQLITE_OPEN_SHAREDCACHE,
};
use std::ffi::{CStr, CString};
use std::io;
use std::ptr::{addr_of_mut, null, null_mut};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

static THREAD_ID: AtomicU64 = AtomicU64::new(0);

pub struct EstablishParams {
    filename: CString,
    open_flags: i32,
    busy_timeout: Duration,
    statement_cache_capacity: usize,
    log_settings: LogSettings,
    extensions: IndexMap<CString, Option<CString>>,
    pub(crate) thread_name: String,
    pub(crate) command_channel_size: usize,
}

impl EstablishParams {
    pub fn from_options(options: &SqliteConnectOptions) -> Result<Self, Error> {
        let mut filename = options
            .filename
            .to_str()
            .ok_or_else(|| {
                io::Error::new(
                    io::ErrorKind::InvalidData,
                    "filename passed to SQLite must be valid UTF-8",
                )
            })?
            .to_owned();

        // By default, we connect to an in-memory database.
        // [SQLITE_OPEN_NOMUTEX] will instruct [sqlite3_open_v2] to return an error if it
        // cannot satisfy our wish for a thread-safe, lock-free connection object

        let mut flags = if options.serialized {
            SQLITE_OPEN_FULLMUTEX
        } else {
            SQLITE_OPEN_NOMUTEX
        };

        flags |= if options.read_only {
            SQLITE_OPEN_READONLY
        } else if options.create_if_missing {
            SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE
        } else {
            SQLITE_OPEN_READWRITE
        };

        if options.in_memory {
            flags |= SQLITE_OPEN_MEMORY;
        }

        flags |= if options.shared_cache {
            SQLITE_OPEN_SHAREDCACHE
        } else {
            SQLITE_OPEN_PRIVATECACHE
        };

        let mut query_params: Vec<String> = vec![];

        if options.immutable {
            query_params.push("immutable=true".into())
        }

        if let Some(vfs) = &options.vfs {
            query_params.push(format!("vfs={}", vfs))
        }

        if !query_params.is_empty() {
            filename = format!("file:{}?{}", filename, query_params.join("&"));
            flags |= libsqlite3_sys::SQLITE_OPEN_URI;
        }

        let filename = CString::new(filename).map_err(|_| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                "filename passed to SQLite must not contain nul bytes",
            )
        })?;

        let extensions = options
            .extensions
            .iter()
            .map(|(name, entry)| {
                let entry = entry
                    .as_ref()
                    .map(|e| {
                        CString::new(e.as_bytes()).map_err(|_| {
                            io::Error::new(
                                io::ErrorKind::InvalidData,
                                "extension entrypoint names passed to SQLite must not contain nul bytes"
                            )
                        })
                    })
                    .transpose()?;
                Ok((
                    CString::new(name.as_bytes()).map_err(|_| {
                        io::Error::new(
                            io::ErrorKind::InvalidData,
                            "extension names passed to SQLite must not contain nul bytes",
                        )
                    })?,
                    entry,
                ))
            })
            .collect::<Result<IndexMap<CString, Option<CString>>, io::Error>>()?;

        Ok(Self {
            filename,
            open_flags: flags,
            busy_timeout: options.busy_timeout,
            statement_cache_capacity: options.statement_cache_capacity,
            log_settings: options.log_settings.clone(),
            extensions,
            thread_name: (options.thread_name)(THREAD_ID.fetch_add(1, Ordering::AcqRel)),
            command_channel_size: options.command_channel_size,
        })
    }

    pub(crate) fn establish(&self) -> Result<ConnectionState, Error> {
        let mut handle = null_mut();

        // <https://www.sqlite.org/c3ref/open.html>
        let mut status = unsafe {
            sqlite3_open_v2(self.filename.as_ptr(), &mut handle, self.open_flags, null())
        };

        if handle.is_null() {
            // Failed to allocate memory
            return Err(Error::Io(io::Error::new(
                io::ErrorKind::OutOfMemory,
                "SQLite is unable to allocate memory to hold the sqlite3 object",
            )));
        }

        // SAFE: tested for NULL just above
        // This allows any returns below to close this handle with RAII
        let handle = unsafe { ConnectionHandle::new(handle) };

        if status != SQLITE_OK {
            return Err(Error::Database(Box::new(SqliteError::new(handle.as_ptr()))));
        }

        // Enable extended result codes
        // https://www.sqlite.org/c3ref/extended_result_codes.html
        unsafe {
            // NOTE: ignore the failure here
            sqlite3_extended_result_codes(handle.as_ptr(), 1);
        }

        if !self.extensions.is_empty() {
            // Enable extension loading via the db_config function, as recommended by the docs rather
            // than the more obvious `sqlite3_enable_load_extension`
            // https://www.sqlite.org/c3ref/db_config.html
            // https://www.sqlite.org/c3ref/c_dbconfig_defensive.html#sqlitedbconfigenableloadextension
            status = unsafe {
                sqlite3_db_config(
                    handle.as_ptr(),
                    SQLITE_DBCONFIG_ENABLE_LOAD_EXTENSION,
                    1, // '1' here means enable the C API function but not the SQL function
                    null::<i32>(),
                )
            };
            if status != SQLITE_OK {
                return Err(Error::Database(Box::new(SqliteError::new(handle.as_ptr()))));
            }

            for ext in self.extensions.iter() {
                // `sqlite3_load_extension` is unusual as it returns its errors via an out-pointer
                // rather than by calling `sqlite3_errmsg`
                let mut error = null_mut();
                status = unsafe {
                    sqlite3_load_extension(
                        handle.as_ptr(),
                        ext.0.as_ptr(),
                        ext.1.as_ref().map_or(null(), |e| e.as_ptr()),
                        addr_of_mut!(error),
                    )
                };

                if status != SQLITE_OK {
                    // SAFETY: We become responsible for any memory allocation at `&error`, so test
                    // for null and take an RAII version for returns
                    let err_msg = if !error.is_null() {
                        unsafe {
                            let e = CStr::from_ptr(error).into();
                            sqlite3_free(error as *mut c_void);
                            e
                        }
                    } else {
                        CString::new("Unknown error when loading extension")
                            .expect("text should be representable as a CString")
                    };
                    return Err(Error::Database(Box::new(SqliteError::extension(
                        handle.as_ptr(),
                        &err_msg,
                    ))));
                }
            }
        }

        // Configure a busy timeout
        // This causes SQLite to automatically sleep in increasing intervals until the time
        // when there is something locked during [sqlite3_step].
        //
        // We also need to convert the u128 value to i32, checking we're not overflowing.
        let ms = i32::try_from(self.busy_timeout.as_millis())
            .expect("Given busy timeout value is too big.");

        status = unsafe { sqlite3_busy_timeout(handle.as_ptr(), ms) };

        if status != SQLITE_OK {
            return Err(Error::Database(Box::new(SqliteError::new(handle.as_ptr()))));
        }

        Ok(ConnectionState {
            handle,
            statements: Statements::new(self.statement_cache_capacity),
            transaction_depth: 0,
            log_settings: self.log_settings.clone(),
        })
    }
}


Review comment (security):
For sanity reasons, this should probably set `SQLITE_DBCONFIG_ENABLE_LOAD_EXTENSION` back to `0` after the extensions have been loaded.

Without that, I could see someone coming along, seeing this and saying "SQLx leaves extension loading enabled after startup, that's a security vulnerability," even if it's technically not since using the C API would be required.

Code after:
                            let e = CStr::from_ptr(error).into();
                            sqlite3_free(error as *mut c_void);
                            e
                        }
                    } else {
                        CString::new("Unknown error when loading extension")
                            .expect("text should be representable as a CString")
                    };
                    return Err(Error::Database(Box::new(SqliteError::extension(
                        handle.as_ptr(),
                        &err_msg,
                    ))));
                }
            }

            // Preempt any hypothetical security issues arising from leaving ENABLE_LOAD_EXTENSION
            // on by disabling the flag again once we've loaded all the requested modules.
            // Fail-fast (via `?`) if disabling the extension loader didn't work for some reason,
            // avoids an unexpected state going undetected.
            unsafe {
                Self::sqlite3_set_load_extension(
                    handle.as_ptr(),
                    SqliteLoadExtensionMode::DisableAll,
                )?;
            }
        }

        // Configure a busy timeout
        // This causes SQLite to automatically sleep in increasing intervals until the time
        // when there is something locked during [sqlite3_step].
        //
        // We also need to convert the u128 value to i32, checking we're not overflowing.
        let ms = i32::try_from(self.busy_timeout.as_millis())
            .expect("Given busy timeout value is too big.");

        status = unsafe { sqlite3_busy_timeout(handle.as_ptr(), ms) };

        if status != SQLITE_OK {
            return Err(Error::Database(Box::new(SqliteError::new(handle.as_ptr()))));
        }

        Ok(ConnectionState {
            handle,
            statements: Statements::new(self.statement_cache_capacity),
            transaction_depth: 0,
            log_settings: self.log_settings.clone(),
        })
    }
}