Fix the custom sql function implementation for Sqlite

weiznich · weiznich · commit da4447b10e83 · 2021-09-13T09:11:07.000+02:00
diff --git a/diesel/src/sqlite/connection/functions.rs b/diesel/src/sqlite/connection/functions.rs
@@ -1,14 +1,21 @@
 extern crate libsqlite3_sys as ffi;
 
 use super::raw::RawConnection;
+use super::row::PrivateSqliteRow;
 use super::serialized_value::SerializedValue;
 use super::{Sqlite, SqliteAggregateFunction};
 use crate::deserialize::{FromSqlRow, StaticallySizedRow};
 use crate::result::{DatabaseErrorKind, Error, QueryResult};
 use crate::row::{Field, PartialRow, Row, RowIndex};
 use crate::serialize::{IsNull, Output, ToSql};
 use crate::sql_types::HasSqlType;
+use crate::sqlite::connection::sqlite_value::OwnedSqliteValue;
+use crate::sqlite::SqliteValue;
+use std::cell::{Ref, RefCell};
 use std::marker::PhantomData;
+use std::mem::ManuallyDrop;
+use std::ops::DerefMut;
+use std::rc::Rc;
 
 pub fn register<ArgsSqlType, RetSqlType, Args, Ret, F>(
     conn: &RawConnection,
@@ -85,7 +92,7 @@ where
 }
 
 pub(crate) fn build_sql_function_args<ArgsSqlType, Args>(
-    args: &[*mut ffi::sqlite3_value],
+    args: &mut [*mut ffi::sqlite3_value],
 ) -> Result<Args, Error>
 where
     Args: FromSqlRow<ArgsSqlType, Sqlite>,
@@ -117,14 +124,67 @@ where
     })
 }
 
-#[derive(Clone)]
 struct FunctionRow<'a> {
-    args: &'a [*mut ffi::sqlite3_value],
+    // we use `ManuallyDrop` to prevent dropping the content of the internal vector
+    // as this buffer is owned by sqlite not by diesel
+    args: Rc<RefCell<ManuallyDrop<PrivateSqliteRow<'a>>>>,
+    field_count: usize,
+    marker: PhantomData<&'a ffi::sqlite3_value>,
+}
+
+impl<'a> Drop for FunctionRow<'a> {
+    fn drop(&mut self) {
+        if let Some(args) = Rc::get_mut(&mut self.args) {
+            if let PrivateSqliteRow::Duplicated { column_names, .. } =
+                DerefMut::deref_mut(RefCell::get_mut(args))
+            {
+                if let Some(inner) = Rc::get_mut(column_names) {
+                    // an empty Vector does not allocate according to the documentation
+                    // so this prevents leaking memory
+                    std::mem::drop(std::mem::replace(inner, Vec::new()));
+                }
+            }
+        }
+    }
 }
 
 impl<'a> FunctionRow<'a> {
-    fn new(args: &'a [*mut ffi::sqlite3_value]) -> Self {
-        Self { args }
+    fn new(args: &mut [*mut ffi::sqlite3_value]) -> Self {
+        let lenghts = args.len();
+        let args = unsafe {
+            Vec::from_raw_parts(
+                // This cast is safe because:
+                // * Casting from a pointer to an arry to a pointer to the first array
+                // element is safe
+                // * Casting from a raw pointer to `NonNull<T>` is safe,
+                // because `NonNull` is #[repr(transparent)]
+                // * Casting from `NonNull<T>` to `OwnedSqliteValue` is safe,
+                // as the struct is `#[repr(transparent)]
+                // * Casting from `NonNull<T>` to `Option<NonNull<T>>` as the documentation
+                // states: "This is so that enums may use this forbidden value as a discriminant –
+                // Option<NonNull<T>> has the same size as *mut T"
+                // * The last point remains true for `OwnedSqliteValue` as `#[repr(transparent)]
+                // guarantees the same layout as the inner type
+                // * It's unsafe to drop the vector (and the vector elements)
+                // because of this we wrap the vector (or better the Row)
+                // Into `ManualDrop` to prevent the dropping
+                args as *mut [*mut ffi::sqlite3_value] as *mut ffi::sqlite3_value
+                    as *mut Option<OwnedSqliteValue>,
+                lenghts,
+                lenghts,
+            )
+        };
+
+        Self {
+            field_count: lenghts,
+            args: Rc::new(RefCell::new(ManuallyDrop::new(
+                PrivateSqliteRow::Duplicated {
+                    values: args,
+                    column_names: Rc::new(vec![None; lenghts]),
+                },
+            ))),
+            marker: PhantomData,
+        }
     }
 }
 
@@ -133,18 +193,17 @@ impl<'a> Row<'a, Sqlite> for FunctionRow<'a> {
     type InnerPartialRow = Self;
 
     fn field_count(&self) -> usize {
-        self.args.len()
+        self.field_count
     }
 
     fn get<I>(&self, idx: I) -> Option<Self::Field>
     where
         Self: crate::row::RowIndex<I>,
     {
         let idx = self.idx(idx)?;
-
-        self.args.get(idx).map(|arg| FunctionArgument {
-            arg: *arg,
-            p: PhantomData,
+        Some(FunctionArgument {
+            args: self.args.clone(),
+            col_idx: idx as i32,
         })
     }
 
@@ -155,7 +214,7 @@ impl<'a> Row<'a, Sqlite> for FunctionRow<'a> {
 
 impl<'a> RowIndex<usize> for FunctionRow<'a> {
     fn idx(&self, idx: usize) -> Option<usize> {
-        if idx < self.args.len() {
+        if idx < self.field_count() {
             Some(idx)
         } else {
             None
@@ -170,8 +229,8 @@ impl<'a, 'b> RowIndex<&'a str> for FunctionRow<'b> {
 }
 
 struct FunctionArgument<'a> {
-    arg: *mut ffi::sqlite3_value,
-    p: PhantomData<&'a ()>,
+    args: Rc<RefCell<ManuallyDrop<PrivateSqliteRow<'a>>>>,
+    col_idx: i32,
 }
 
 impl<'a> Field<'a, Sqlite> for FunctionArgument<'a> {
@@ -187,7 +246,9 @@ impl<'a> Field<'a, Sqlite> for FunctionArgument<'a> {
     where
         'a: 'b,
     {
-        todo!()
-        //        unsafe { SqliteValue::new(self.arg) }
+        SqliteValue::new(
+            Ref::map(self.args.borrow(), |drop| std::ops::Deref::deref(drop)),
+            self.col_idx,
+        )
     }
 }
diff --git a/diesel/src/sqlite/connection/raw.rs b/diesel/src/sqlite/connection/raw.rs
@@ -91,7 +91,7 @@ impl RawConnection {
         f: F,
     ) -> QueryResult<()>
     where
-        F: FnMut(&Self, &[*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue>
+        F: FnMut(&Self, &mut [*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue>
             + std::panic::UnwindSafe
             + Send
             + 'static,
@@ -269,7 +269,7 @@ extern "C" fn run_custom_function<F>(
     num_args: libc::c_int,
     value_ptr: *mut *mut ffi::sqlite3_value,
 ) where
-    F: FnMut(&RawConnection, &[*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue>
+    F: FnMut(&RawConnection, &mut [*mut ffi::sqlite3_value]) -> QueryResult<SerializedValue>
         + std::panic::UnwindSafe
         + Send
         + 'static,
@@ -278,7 +278,6 @@ extern "C" fn run_custom_function<F>(
     static NULL_DATA_ERR: &str = "An unknown error occurred. sqlite3_user_data returned a null pointer. This should never happen.";
     static NULL_CONN_ERR: &str = "An unknown error occurred. sqlite3_context_db_handle returned a null pointer. This should never happen.";
 
-    let args = unsafe { slice::from_raw_parts(value_ptr, num_args as _) };
     let conn = match unsafe { NonNull::new(ffi::sqlite3_context_db_handle(ctx)) } {
         // We use `ManuallyDrop` here because we do not want to run the
         // Drop impl of `RawConnection` as this would close the connection
@@ -306,13 +305,16 @@ extern "C" fn run_custom_function<F>(
     // this is sound as `F` itself and the stored string is `UnwindSafe`
     let callback = std::panic::AssertUnwindSafe(&mut data_ptr.callback);
 
-    let result =
-        std::panic::catch_unwind(move || Ok((callback.0)(&*conn, args)?)).unwrap_or_else(|p| {
-            Err(SqliteCallbackError::Panic(
-                p,
-                data_ptr.function_name.clone(),
-            ))
-        });
+    let result = std::panic::catch_unwind(move || {
+        let args = unsafe { slice::from_raw_parts_mut(value_ptr, num_args as _) };
+        Ok((callback.0)(&*conn, args)?)
+    })
+    .unwrap_or_else(|p| {
+        Err(SqliteCallbackError::Panic(
+            p,
+            data_ptr.function_name.clone(),
+        ))
+    });
     match result {
         Ok(value) => value.result_of(ctx),
         Err(e) => {
@@ -342,15 +344,16 @@ extern "C" fn run_aggregator_step_function<ArgsSqlType, RetSqlType, Args, Ret, A
     Ret: ToSql<RetSqlType, Sqlite>,
     Sqlite: HasSqlType<RetSqlType>,
 {
-    let args = unsafe { slice::from_raw_parts(value_ptr, num_args as _) };
-    let result =
-        std::panic::catch_unwind(move || run_aggregator_step::<A, Args, ArgsSqlType>(ctx, args))
-            .unwrap_or_else(|e| {
-                Err(SqliteCallbackError::Panic(
-                    e,
-                    format!("{}::step() paniced", std::any::type_name::<A>()),
-                ))
-            });
+    let result = std::panic::catch_unwind(move || {
+        let args = unsafe { slice::from_raw_parts_mut(value_ptr, num_args as _) };
+        run_aggregator_step::<A, Args, ArgsSqlType>(ctx, args)
+    })
+    .unwrap_or_else(|e| {
+        Err(SqliteCallbackError::Panic(
+            e,
+            format!("{}::step() paniced", std::any::type_name::<A>()),
+        ))
+    });
 
     match result {
         Ok(()) => {}
@@ -360,7 +363,7 @@ extern "C" fn run_aggregator_step_function<ArgsSqlType, RetSqlType, Args, Ret, A
 
 fn run_aggregator_step<A, Args, ArgsSqlType>(
     ctx: *mut ffi::sqlite3_context,
-    args: &[*mut ffi::sqlite3_value],
+    args: &mut [*mut ffi::sqlite3_value],
 ) -> Result<(), SqliteCallbackError>
 where
     A: SqliteAggregateFunction<Args>,
diff --git a/diesel/src/sqlite/connection/row.rs b/diesel/src/sqlite/connection/row.rs
@@ -164,3 +164,119 @@ impl<'a> Field<'a, Sqlite> for SqliteField<'a> {
         SqliteValue::new(self.row.inner.borrow(), self.col_idx)
     }
 }
+
+#[test]
+fn fun_with_row_iters() {
+    crate::table! {
+        #[allow(unused_parens)]
+        users(id) {
+            id -> Integer,
+            name -> Text,
+        }
+    }
+
+    use crate::deserialize::{FromSql, FromSqlRow};
+    use crate::prelude::*;
+    use crate::row::{Field, Row};
+    use crate::sql_types;
+
+    let conn = &mut crate::test_helpers::connection();
+
+    crate::sql_query("CREATE TABLE users(id INTEGER PRIMARY KEY, name TEXT NOT NULL);")
+        .execute(conn)
+        .unwrap();
+
+    crate::insert_into(users::table)
+        .values(vec![
+            (users::id.eq(1), users::name.eq("Sean")),
+            (users::id.eq(2), users::name.eq("Tess")),
+        ])
+        .execute(conn)
+        .unwrap();
+
+    let query = users::table.select((users::id, users::name));
+
+    let expected = vec![(1, String::from("Sean")), (2, String::from("Tess"))];
+
+    let row_iter = conn.load(&query).unwrap();
+    for (row, expected) in row_iter.zip(&expected) {
+        let row = row.unwrap();
+
+        let deserialized = <(i32, String) as FromSqlRow<
+            (sql_types::Integer, sql_types::Text),
+            _,
+        >>::build_from_row(&row)
+        .unwrap();
+
+        assert_eq!(&deserialized, expected);
+    }
+
+    {
+        let collected_rows = conn.load(&query).unwrap().collect::<Vec<_>>();
+
+        for (row, expected) in collected_rows.iter().zip(&expected) {
+            let deserialized = row
+                .as_ref()
+                .map(|row| {
+                    <(i32, String) as FromSqlRow<
+                            (sql_types::Integer, sql_types::Text),
+                        _,
+                        >>::build_from_row(row).unwrap()
+                })
+                .unwrap();
+
+            assert_eq!(&deserialized, expected);
+        }
+    }
+
+    let mut row_iter = conn.load(&query).unwrap();
+
+    let first_row = row_iter.next().unwrap().unwrap();
+    let first_fields = (first_row.get(0).unwrap(), first_row.get(1).unwrap());
+    let first_values = (first_fields.0.value(), first_fields.1.value());
+
+    assert!(row_iter.next().unwrap().is_err());
+    std::mem::drop(first_values);
+
+    let second_row = row_iter.next().unwrap().unwrap();
+    let second_fields = (second_row.get(0).unwrap(), second_row.get(1).unwrap());
+    let second_values = (second_fields.0.value(), second_fields.1.value());
+
+    assert!(row_iter.next().unwrap().is_err());
+    std::mem::drop(second_values);
+
+    assert!(row_iter.next().is_none());
+
+    let first_values = (first_fields.0.value(), first_fields.1.value());
+    let second_values = (second_fields.0.value(), second_fields.1.value());
+
+    assert_eq!(
+        <i32 as FromSql<sql_types::Integer, Sqlite>>::from_nullable_sql(first_values.0).unwrap(),
+        expected[0].0
+    );
+    assert_eq!(
+        <String as FromSql<sql_types::Text, Sqlite>>::from_nullable_sql(first_values.1).unwrap(),
+        expected[0].1
+    );
+
+    assert_eq!(
+        <i32 as FromSql<sql_types::Integer, Sqlite>>::from_nullable_sql(second_values.0).unwrap(),
+        expected[1].0
+    );
+    assert_eq!(
+        <String as FromSql<sql_types::Text, Sqlite>>::from_nullable_sql(second_values.1).unwrap(),
+        expected[1].1
+    );
+
+    let first_fields = (first_row.get(0).unwrap(), first_row.get(1).unwrap());
+    let first_values = (first_fields.0.value(), first_fields.1.value());
+
+    assert_eq!(
+        <i32 as FromSql<sql_types::Integer, Sqlite>>::from_nullable_sql(first_values.0).unwrap(),
+        expected[0].0
+    );
+    assert_eq!(
+        <String as FromSql<sql_types::Text, Sqlite>>::from_nullable_sql(first_values.1).unwrap(),
+        expected[0].1
+    );
+}
diff --git a/diesel/src/sqlite/connection/sqlite_value.rs b/diesel/src/sqlite/connection/sqlite_value.rs
@@ -23,6 +23,7 @@ pub struct SqliteValue<'a, 'b> {
     col_idx: i32,
 }
 
+#[repr(transparent)]
 pub struct OwnedSqliteValue {
     pub(super) value: NonNull<ffi::sqlite3_value>,
 }
diff --git a/diesel/src/sqlite/connection/statement_iterator.rs b/diesel/src/sqlite/connection/statement_iterator.rs
@@ -88,9 +88,23 @@ impl<'a> Iterator for StatementIterator<'a> {
                     // a user stored the row in some long time container before calling next another time
                     // In this case we copy out the current values into a temporary store and advance
                     // the statement iterator internally afterwards
-                    if let PrivateSqliteRow::Direct(stmt) =
-                        last_row.replace_with(|inner| inner.duplicate(&mut self.column_names))
-                    {
+                    let last_row = {
+                        let mut last_row = match last_row.try_borrow_mut() {
+                            Ok(o) => o,
+                            Err(_e) => {
+                                self.inner = Started(last_row.clone());
+                                return Some(Err(crate::result::Error::DeserializationError(
+                                    "Failed to reborrow row. Try to release any `SqliteValue` \
+                                     that exists at this point"
+                                        .into(),
+                                )));
+                            }
+                        };
+                        let last_row = &mut *last_row;
+                        let duplicated = last_row.duplicate(&mut self.column_names);
+                        std::mem::replace(last_row, duplicated)
+                    };
+                    if let PrivateSqliteRow::Direct(stmt) = last_row {
                         match stmt.step() {
                             Err(e) => Some(Err(e)),
                             Ok(None) => None,
diff --git a/diesel/src/util.rs b/diesel/src/util.rs
diff --git a/diesel_tests/tests/deserialization.rs b/diesel_tests/tests/deserialization.rs

Original file line number	Diff line number	Diff line change
`@@ -23,6 +23,7 @@ pub struct SqliteValue<'a, 'b> {`
`23`	`23`	`col_idx: i32,`
`24`	`24`	`}`
`25`	`25`
	`26`	`+#[repr(transparent)]`
`26`	`27`	`pub struct OwnedSqliteValue {`
`27`	`28`	`pub(super) value: NonNull<ffi::sqlite3_value>,`
`28`	`29`	`}`