Add a wrapper for strings allocated on sqlite heap

src/raw_statement.rs

@@ -1,6 +1,8 @@
 use super::ffi;
 use super::unlock_notify;
 use super::StatementStatus;
+#[cfg(feature = "modern_sqlite")]
+use crate::util::SqliteMallocString;
 use std::ffi::CStr;
 use std::os::raw::c_int;
 use std::ptr;
@@ -124,15 +126,9 @@ impl RawStatement {
         unsafe { ffi::sqlite3_stmt_readonly(self.0) != 0 }
     }
 
-    /// `CStr` must be freed
     #[cfg(feature = "modern_sqlite")] // 3.14.0
-    pub unsafe fn expanded_sql(&self) -> Option<&CStr> {
-        let ptr = ffi::sqlite3_expanded_sql(self.0);
-        if ptr.is_null() {
-            None
-        } else {
-            Some(CStr::from_ptr(ptr))
-        }
+    pub(crate) fn expanded_sql(&self) -> Option<SqliteMallocString> {
+        unsafe { SqliteMallocString::from_raw(ffi::sqlite3_expanded_sql(self.0)) }
     }
 
     pub fn get_status(&self, status: StatementStatus, reset: bool) -> i32 {

src/statement.rs

@@ -668,16 +668,9 @@ impl Statement<'_> {
     /// bound parameters expanded.
     #[cfg(feature = "modern_sqlite")]
     pub fn expanded_sql(&self) -> Option<String> {
-        unsafe {
-            match self.stmt.expanded_sql() {
-                Some(s) => {
-                    let sql = String::from_utf8_lossy(s.to_bytes()).to_string();
-                    ffi::sqlite3_free(s.as_ptr() as *mut _);
-                    Some(sql)
-                }
-                _ => None,
-            }
-        }
+        self.stmt
+            .expanded_sql()
+            .map(|s| s.to_string_lossy().to_string())
     }
 
     /// Get the value for one of the status counters for this statement.

src/util/mod.rs

@@ -3,3 +3,9 @@ pub(crate) mod param_cache;
 mod small_cstr;
 pub(crate) use param_cache::ParamIndexCache;
 pub(crate) use small_cstr::SmallCString;
+
+// Doesn't use any modern features or vtab stuff, but is only used by them.
+#[cfg(any(feature = "modern_sqlite", feature = "vtab"))]
+mod sqlite_string;
+#[cfg(any(feature = "modern_sqlite", feature = "vtab"))]
+pub(crate) use sqlite_string::SqliteMallocString;

src/util/sqlite_string.rs

@@ -0,0 +1,227 @@
+use crate::ffi;
+use std::marker::PhantomData;
+use std::os::raw::{c_char, c_int};
+use std::ptr::NonNull;
+
+/// A string we own that's allocated on the SQLite heap. Automatically calls
+/// `sqlite3_free` when dropped, unless `into_raw` (or `into_inner`) is called
+/// on it. If constructed from a rust string, `sqlite3_malloc` is used.
+///
+/// It has identical representation to a nonnull `*mut c_char`, so you can use
+/// it transparently as one. It's nonnull, so Option<SqliteMallocString> is a
+/// nullable one.
+///
+/// Note that this is for Strings we manage that are owned by SQLite, like
+/// `sqlite3_extended_sql` results, some error message pointers. Most strings
+/// are not this! Misuse is extremely dangerous!
+///
+/// Note that this is not a lossless interface. Incoming strings with internal
+/// NULs are modified, and outgoing strings which are non-UTF8 are modified.
+/// This seems unavoidable.
+#[repr(transparent)]
+pub(crate) struct SqliteMallocString {
+    ptr: NonNull<c_char>,
+    _boo: PhantomData<Box<[c_char]>>,
+}
+// This is owned data for a primitive type, and thus it's safe to implement
+// these. That said, nothing needs them, and they make things easier to misuse.
+
+// unsafe impl Send for SqliteMallocString {}
+// unsafe impl Sync for SqliteMallocString {}
+
+impl SqliteMallocString {
+    /// SAFETY: Caller must be certain that `m` a nul-terminated c string
+    /// allocated by sqlite3_malloc, and that SQLite expects us to free it!
+    #[inline]
+    pub(crate) unsafe fn from_raw_nonnull(ptr: NonNull<c_char>) -> Self {
+        Self {
+            ptr,
+            _boo: PhantomData,
+        }
+    }
+
+    /// SAFETY: Caller must be certain that `m` a nul-terminated c string
+    /// allocated by sqlite3_malloc, and that SQLite expects us to free it!
+    #[inline]
+    pub(crate) unsafe fn from_raw(ptr: *mut c_char) -> Option<Self> {
+        NonNull::new(ptr).map(|p| Self::from_raw_nonnull(p))
+    }
+
+    /// Get the pointer behind `self`. After this is called, we no longer manage
+    /// it.
+    #[inline]
+    pub(crate) fn into_inner(self) -> NonNull<c_char> {
+        let p = self.ptr;
+        std::mem::forget(self);
+        p
+    }
+
+    /// Get the pointer behind `self`. After this is called, we no longer manage
+    /// it.
+    #[inline]
+    pub(crate) fn into_raw(self) -> *mut c_char {
+        self.into_inner().as_ptr()
+    }
+
+    /// Borrow the pointer behind `self`. We still manage it when this function
+    /// returns. If you want to relinquish ownership, use `into_raw`.
+    #[inline]
+    pub(crate) fn as_ptr(&self) -> *const c_char {
+        self.ptr.as_ptr()
+    }
+
+    #[inline]
+    pub(crate) fn as_cstr(&self) -> &std::ffi::CStr {
+        unsafe { std::ffi::CStr::from_ptr(self.as_ptr()) }
+    }
+
+    #[inline]
+    pub(crate) fn to_string_lossy(&self) -> std::borrow::Cow<'_, str> {
+        self.as_cstr().to_string_lossy()
+    }
+
+    /// Convert `s` into a SQLite string.
+    ///
+    /// This should almost never be done except for cases like error messages or
+    /// other strings that SQLite frees.
+    ///
+    /// If `s` contains internal NULs, we'll replace them with
+    /// `NUL_REPLACE_CHAR`.
+    ///
+    /// Except for debug_asserts which may trigger during testing, this function
+    /// never panics. If we hit integer overflow or the allocation fails, we
+    /// call `handle_alloc_error` which aborts the program after calling a
+    /// global hook.
+    ///
+    /// This means it's safe to use in extern "C" functions outside of
+    /// catch_unwind.
+    pub(crate) fn from_str(s: &str) -> Self {
+        use std::convert::TryFrom;
+        let s = if s.as_bytes().contains(&0) {
+            std::borrow::Cow::Owned(make_nonnull(s))
+        } else {
+            std::borrow::Cow::Borrowed(s)
+        };
+        debug_assert!(!s.as_bytes().contains(&0));
+        let bytes: &[u8] = s.as_ref().as_bytes();
+        let src_ptr: *const c_char = bytes.as_ptr().cast();
+        let src_len = bytes.len();
+        let maybe_len_plus_1 = s.len().checked_add(1).and_then(|v| c_int::try_from(v).ok());
+        unsafe {
+            let res_ptr = maybe_len_plus_1
+                .and_then(|len_to_alloc| {
+                    // `>` because we added 1.
+                    debug_assert!(len_to_alloc > 0);
+                    debug_assert_eq!((len_to_alloc - 1) as usize, src_len);
+                    NonNull::new(ffi::sqlite3_malloc(len_to_alloc) as *mut c_char)
+                })
+                .unwrap_or_else(|| {
+                    use std::alloc::{handle_alloc_error, Layout};
+                    // Report via handle_alloc_error so that it can be handled with any
+                    // other allocation errors and properly diagnosed.
+                    //
+                    // This is safe:
+                    // - `align` is never 0
+                    // - `align` is always a power of 2.
+                    // - `size` needs no realignment because it's guaranteed to be
+                    //   aligned (everything is aligned to 1)
+                    // - `size` is also never zero, although this function doesn't actually require it now.
+                    let layout = Layout::from_size_align_unchecked(s.len().saturating_add(1), 1);
+                    // Note: This call does not return.
+                    handle_alloc_error(layout);
+                });
+            let buf: *mut i8 = res_ptr.as_ptr();
+            src_ptr.copy_to_nonoverlapping(buf, src_len);
+            buf.add(src_len).write(0);
+            debug_assert_eq!(std::ffi::CStr::from_ptr(res_ptr.as_ptr()).to_bytes(), bytes);
+            Self::from_raw_nonnull(res_ptr)
+        }
+    }
+}
+
+const NUL_REPLACE: &str = "␀";
+
+#[cold]
+fn make_nonnull(v: &str) -> String {
+    v.replace('\0', NUL_REPLACE)
+}
+
+impl Drop for SqliteMallocString {
+    fn drop(&mut self) {
+        unsafe { ffi::sqlite3_free(self.ptr.as_ptr().cast()) };
+    }
+}
+
+impl std::fmt::Debug for SqliteMallocString {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.to_string_lossy().fmt(f)
+    }
+}
+
+impl std::fmt::Display for SqliteMallocString {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.to_string_lossy().fmt(f)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    #[test]
+    fn test_from_str() {
+        let to_check = [
+            ("", ""),
+            ("\0", "␀"),
+            ("␀", "␀"),
+            ("\0bar", "␀bar"),
+            ("foo\0bar", "foo␀bar"),
+            ("foo\0", "foo␀"),
+            ("a\0b\0c\0\0d", "a␀b␀c␀␀d"),
+            ("foobar0123", "foobar0123"),
+        ];
+
+        for &(input, output) in &to_check {
+            let s = SqliteMallocString::from_str(input);
+            assert_eq!(s.to_string_lossy(), output);
+            assert_eq!(s.as_cstr().to_str().unwrap(), output);
+        }
+    }
+
+    // This will trigger an asan error if into_raw still freed the ptr.
+    #[test]
+    fn test_lossy() {
+        let p = SqliteMallocString::from_str("abcd").into_raw();
+        // Make invalid
+        let s = unsafe {
+            p.cast::<u8>().write(b'\xff');
+            SqliteMallocString::from_raw(p).unwrap()
+        };
+        assert_eq!(s.to_string_lossy().as_ref(), "\u{FFFD}bcd");
+    }
+
+    // This will trigger an asan error if into_raw still freed the ptr.
+    #[test]
+    fn test_into_raw() {
+        let mut v = vec![];
+        for i in 0..1000 {
+            v.push(SqliteMallocString::from_str(&i.to_string()).into_raw());
+            v.push(SqliteMallocString::from_str(&format!("abc {} 😀", i)).into_raw());
+        }
+        unsafe {
+            for (i, s) in v.chunks_mut(2).enumerate() {
+                let s0 = std::mem::replace(&mut s[0], std::ptr::null_mut());
+                let s1 = std::mem::replace(&mut s[1], std::ptr::null_mut());
+                assert_eq!(
+                    std::ffi::CStr::from_ptr(s0).to_str().unwrap(),
+                    &i.to_string()
+                );
+                assert_eq!(
+                    std::ffi::CStr::from_ptr(s1).to_str().unwrap(),
+                    &format!("abc {} 😀", i)
+                );
+                let _ = SqliteMallocString::from_raw(s0).unwrap();
+                let _ = SqliteMallocString::from_raw(s1).unwrap();
+            }
+        }
+    }
+}

src/vtab/mod.rs

@@ -1026,23 +1026,8 @@ unsafe fn result_error<T>(ctx: *mut ffi::sqlite3_context, result: Result<T>) ->
 
 // Space to hold this string must be obtained
 // from an SQLite memory allocation function
-unsafe fn alloc<S: AsRef<[u8]>>(s: S) -> *mut c_char {
-    use std::convert::TryInto;
-    let s = s.as_ref();
-    if memchr::memchr(0, s).is_some() {
-        panic!("Null character found")
-    }
-    let len = s.len();
-    let total_len = len.checked_add(1).unwrap();
-
-    let dst = ffi::sqlite3_malloc(total_len.try_into().unwrap()) as *mut c_char;
-    if dst.is_null() {
-        panic!("Out of memory")
-    }
-    ptr::copy_nonoverlapping(s.as_ptr() as *const c_char, dst, len);
-    // null terminator
-    *dst.offset(len.try_into().unwrap()) = 0;
-    dst
+fn alloc(s: &str) -> *mut c_char {
+    crate::util::SqliteMallocString::from_str(s).into_raw()
 }
 
 #[cfg(feature = "array")]