From 8aae08ec052cb5107403c3dc69dd88581fd350a3 Mon Sep 17 00:00:00 2001 From: David Li Date: Fri, 9 Oct 2026 09:42:27 +0900 Subject: [PATCH] fix(rust/ffi): avoid potential overlapping &mut in exporter Assisted-by: GPT-6.1 Sol --- rust/driver_manager/tests/common/dummy.rs | 62 +++++-- .../tests/exporter_cancellation.rs | 156 +++++++++++++++++- rust/ffi/src/driver_exporter.rs | 123 ++++++++++---- 3 files changed, 293 insertions(+), 48 deletions(-) diff --git a/rust/driver_manager/tests/common/dummy.rs b/rust/driver_manager/tests/common/dummy.rs index acac42995c..5d9b1e64df 100644 --- a/rust/driver_manager/tests/common/dummy.rs +++ b/rust/driver_manager/tests/common/dummy.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; + use adbc_core::error::{Error, Result, Status}; use adbc_core::options::{ InfoCode, ObjectDepth, OptionConnection, OptionDatabase, OptionStatement, OptionValue, @@ -51,29 +53,69 @@ impl Driver for DummyDriver { } #[derive(Default)] -pub struct DummyDatabase {} +pub struct DummyDatabase { + options: HashMap, +} + +impl DummyDatabase { + fn option(&self, key: OptionDatabase) -> Result<&OptionValue> { + // Let concurrent exported getters overlap while their driver borrows are live. + std::thread::yield_now(); + self.options.get(&key).ok_or_else(|| { + Error::with_message_and_status( + format!("Option key not found: {key:?}"), + Status::NotFound, + ) + }) + } +} impl Optionable for DummyDatabase { type Option = OptionDatabase; - fn set_option(&mut self, _key: Self::Option, _value: OptionValue) -> Result<()> { + fn set_option(&mut self, key: Self::Option, value: OptionValue) -> Result<()> { + self.options.insert(key, value); Ok(()) } - fn get_option_bytes(&self, _key: Self::Option) -> Result> { - todo!() + fn get_option_bytes(&self, key: Self::Option) -> Result> { + match self.option(key)? { + OptionValue::Bytes(value) => Ok(value.clone()), + _ => Err(Error::with_message_and_status( + "Expected bytes option", + Status::InvalidState, + )), + } } - fn get_option_double(&self, _key: Self::Option) -> Result { - todo!() + fn get_option_double(&self, key: Self::Option) -> Result { + match self.option(key)? { + OptionValue::Double(value) => Ok(*value), + _ => Err(Error::with_message_and_status( + "Expected double option", + Status::InvalidState, + )), + } } - fn get_option_int(&self, _key: Self::Option) -> Result { - todo!() + fn get_option_int(&self, key: Self::Option) -> Result { + match self.option(key)? { + OptionValue::Int(value) => Ok(*value), + _ => Err(Error::with_message_and_status( + "Expected integer option", + Status::InvalidState, + )), + } } - fn get_option_string(&self, _key: Self::Option) -> Result { - todo!() + fn get_option_string(&self, key: Self::Option) -> Result { + match self.option(key)? { + OptionValue::String(value) => Ok(value.clone()), + _ => Err(Error::with_message_and_status( + "Expected string option", + Status::InvalidState, + )), + } } } diff --git a/rust/driver_manager/tests/exporter_cancellation.rs b/rust/driver_manager/tests/exporter_cancellation.rs index 186dd47e3a..7acd3b7d53 100644 --- a/rust/driver_manager/tests/exporter_cancellation.rs +++ b/rust/driver_manager/tests/exporter_cancellation.rs @@ -18,13 +18,15 @@ // Test that the driver exporter and cancellation do not have undefined // behavior by introducing overlapping mutable borrows. Run under Miri. -use std::sync::{Arc, Mutex}; +use std::sync::atomic::{AtomicPtr, Ordering}; +use std::sync::{Arc, Barrier, Mutex}; -use adbc_core::error::{Result, Status}; +use adbc_core::constants::ADBC_STATUS_OK; +use adbc_core::error::{Error, Result, Status}; use adbc_core::options::AdbcVersion; use adbc_core::{CancelHandle, Connection, Database, Driver, Statement}; use adbc_driver_manager::ManagedDriver; -use adbc_ffi::FFI_AdbcDriverInitFunc; +use adbc_ffi::{FFI_AdbcDatabase, FFI_AdbcDriverInitFunc, FFI_AdbcError, FFIDriver}; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::{ArrowError, Schema, SchemaRef}; @@ -105,3 +107,151 @@ fn connection_cancellation() -> Result<()> { guard.join().unwrap(); Ok(()) } + +// Exercise database getters directly (avoiding driver manager's lock) to test +// for potential UB in the exporter under concurrent access. +fn concurrent_database_getters(initialized: bool) -> Result<()> { + let driver = common::dummy::DummyDriver::ffi_driver(); + let new = driver.DatabaseNew.unwrap(); + let init = driver.DatabaseInit.unwrap(); + let release = driver.DatabaseRelease.unwrap(); + let set_string = driver.DatabaseSetOption.unwrap(); + let set_bytes = driver.DatabaseSetOptionBytes.unwrap(); + let set_int = driver.DatabaseSetOptionInt.unwrap(); + let set_double = driver.DatabaseSetOptionDouble.unwrap(); + let get_string = driver.DatabaseGetOption.unwrap(); + let get_bytes = driver.DatabaseGetOptionBytes.unwrap(); + let get_int = driver.DatabaseGetOptionInt.unwrap(); + let get_double = driver.DatabaseGetOptionDouble.unwrap(); + + let mut database = FFI_AdbcDatabase::default(); + let mut error = FFI_AdbcError::default(); + // Use the same raw pointer throughout: workers must not create &mut borrows + // of the shared handle, and passing through AtomicPtr preserves provenance. + let database_ptr = &raw mut database; + unsafe { + assert_eq!(new(database_ptr, &mut error), ADBC_STATUS_OK); + assert_eq!( + set_string( + database_ptr, + c"string".as_ptr(), + c"hello".as_ptr(), + &mut error + ), + ADBC_STATUS_OK + ); + assert_eq!( + set_bytes( + database_ptr, + c"bytes".as_ptr(), + b"hello".as_ptr(), + 5, + &mut error + ), + ADBC_STATUS_OK + ); + assert_eq!( + set_int(database_ptr, c"int".as_ptr(), 42, &mut error), + ADBC_STATUS_OK + ); + assert_eq!( + set_double(database_ptr, c"double".as_ptr(), 1.25, &mut error), + ADBC_STATUS_OK + ); + if initialized { + assert_eq!(init(database_ptr, &mut error), ADBC_STATUS_OK); + // Also exercise the setter against the initialized database mutex. + assert_eq!( + set_int(database_ptr, c"int".as_ptr(), 42, &mut error), + ADBC_STATUS_OK + ); + } + } + + let shared = AtomicPtr::new(database_ptr); + let barrier = Barrier::new(2); + let result = std::thread::scope(|scope| -> Result<()> { + let mut workers = Vec::new(); + for _ in 0..2 { + workers.push(scope.spawn(|| { + let database = shared.load(Ordering::Relaxed); + let mut error = FFI_AdbcError::default(); + barrier.wait(); + for _ in 0..4 { + let mut string = [0; 6]; + let mut length = string.len(); + unsafe { + assert_eq!( + get_string( + database, + c"string".as_ptr(), + string.as_mut_ptr().cast(), + &mut length, + &mut error + ), + ADBC_STATUS_OK + ); + } + assert_eq!(length, 6); + assert_eq!(&string, b"hello\0"); + + let mut bytes = [0; 5]; + let mut length = bytes.len(); + unsafe { + assert_eq!( + get_bytes( + database, + c"bytes".as_ptr(), + bytes.as_mut_ptr(), + &mut length, + &mut error + ), + ADBC_STATUS_OK + ); + } + assert_eq!(length, 5); + assert_eq!(&bytes, b"hello"); + + let mut int = 0; + unsafe { + assert_eq!( + get_int(database, c"int".as_ptr(), &mut int, &mut error), + ADBC_STATUS_OK + ); + } + assert_eq!(int, 42); + + let mut double = 0.0; + unsafe { + assert_eq!( + get_double(database, c"double".as_ptr(), &mut double, &mut error), + ADBC_STATUS_OK + ); + } + assert_eq!(double, 1.25); + } + })); + } + for worker in workers { + worker.join().map_err(|_| { + Error::with_message_and_status("Database getter worker panicked", Status::Internal) + })?; + } + Ok(()) + }); + + unsafe { + assert_eq!(release(database_ptr, &mut error), ADBC_STATUS_OK); + } + result +} + +#[test] +fn database_getters_before_init() -> Result<()> { + concurrent_database_getters(false) +} + +#[test] +fn database_getters_after_init() -> Result<()> { + concurrent_database_getters(true) +} diff --git a/rust/ffi/src/driver_exporter.rs b/rust/ffi/src/driver_exporter.rs index cd356768f5..72aa43add2 100644 --- a/rust/ffi/src/driver_exporter.rs +++ b/rust/ffi/src/driver_exporter.rs @@ -43,21 +43,31 @@ type StatementType = enum ExportedDatabase { /// Pre-init options - Options(HashMap), + Options(Mutex>), /// Initialized database - Database(DatabaseType), + Database(Mutex>), } +type DatabaseTuple<'a, DriverType> = ( + Option>>, + Option>>, +); + impl ExportedDatabase { - fn tuple( - &mut self, - ) -> ( - Option<&mut HashMap>, - Option<&mut DatabaseType>, - ) { + fn tuple(&self) -> Result> { match self { - Self::Options(options) => (Some(options), None), - Self::Database(database) => (None, Some(database)), + Self::Options(options) => options + .lock() + .map_err(|_| { + Error::with_message_and_status("Database mutex poisoned", Status::Internal) + }) + .map(|guard| (Some(guard), None)), + Self::Database(database) => database + .lock() + .map_err(|_| { + Error::with_message_and_status("Database mutex poisoned", Status::Internal) + }) + .map(|guard| (None, Some(guard))), } } } @@ -575,10 +585,10 @@ fn check_poison() -> Result<()> { // Database unsafe fn database_private_data<'a, DriverType: Driver>( - database: &mut FFI_AdbcDatabase, -) -> Result<&'a mut ExportedDatabase> { - let exported = database.private_data as *mut ExportedDatabase; - exported.as_mut().ok_or(Error::with_message_and_status( + database: *mut FFI_AdbcDatabase, +) -> Result<&'a ExportedDatabase> { + let exported = (*database).private_data as *mut ExportedDatabase; + exported.as_ref().ok_or(Error::with_message_and_status( "Uninitialized database", Status::InvalidState, )) @@ -586,7 +596,7 @@ unsafe fn database_private_data<'a, DriverType: Driver>( // SAFETY: Will panic if `key` is null. unsafe fn database_set_option_impl>( - database: &mut FFI_AdbcDatabase, + database: *mut FFI_AdbcDatabase, key: *const c_char, value: Value, error: *mut FFI_AdbcError, @@ -598,9 +608,21 @@ unsafe fn database_set_option_impl> match exported { ExportedDatabase::Options(options) => { + let mut options = check_err!( + options.lock().map_err(|_| { + Error::with_message_and_status("Database mutex poisoned", Status::Internal) + }), + error + ); options.insert(key.into(), value.into()); } ExportedDatabase::Database(database) => { + let mut database = check_err!( + database.lock().map_err(|_| { + Error::with_message_and_status("Database mutex poisoned", Status::Internal) + }), + error + ); check_err!(database.set_option(key.into(), value.into()), error); } } @@ -614,7 +636,9 @@ extern "C" fn database_new( ) -> AdbcStatusCode { catch_panic(error, || { let database = pointer_as_mut!(database, error); - let exported = Box::new(ExportedDatabase::::Options(HashMap::new())); + let exported = Box::new(ExportedDatabase::::Options(Mutex::new( + HashMap::new(), + ))); database.private_data = Box::into_raw(exported) as *mut c_void; ADBC_STATUS_OK @@ -628,15 +652,28 @@ extern "C" fn database_init( catch_panic(error, || { let database = pointer_as_mut!(database, error); + // Initialization requires exclusive access to the database handle. let exported = check_err!( - unsafe { database_private_data::(database) }, + unsafe { + let exported = database.private_data as *mut ExportedDatabase; + exported.as_mut().ok_or(Error::with_message_and_status( + "Uninitialized database", + Status::InvalidState, + )) + }, error ); if let ExportedDatabase::Options(options) = exported { + let options = check_err!( + options.get_mut().map_err(|_| { + Error::with_message_and_status("Database mutex poisoned", Status::Internal) + }), + error + ); let mut driver = DriverType::default(); let database = check_err!(driver.new_database_with_opts(options.clone()), error); - *exported = ExportedDatabase::Database(database); + *exported = ExportedDatabase::Database(Mutex::new(database)); } else { check_err!( Err(Error::with_message_and_status( @@ -682,7 +719,7 @@ extern "C" fn database_set_option( error: *mut FFI_AdbcError, ) -> AdbcStatusCode { catch_panic(error, || { - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); check_not_null!(key, error); check_not_null!(value, error); @@ -698,7 +735,7 @@ extern "C" fn database_set_option_int( error: *mut FFI_AdbcError, ) -> AdbcStatusCode { catch_panic(error, || { - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); check_not_null!(key, error); unsafe { database_set_option_impl::(database, key, value, error) } @@ -712,7 +749,7 @@ extern "C" fn database_set_option_double( error: *mut FFI_AdbcError, ) -> AdbcStatusCode { catch_panic(error, || { - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); check_not_null!(key, error); unsafe { database_set_option_impl::(database, key, value, error) } @@ -727,7 +764,7 @@ extern "C" fn database_set_option_bytes( error: *mut FFI_AdbcError, ) -> AdbcStatusCode { catch_panic(error, || { - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); check_not_null!(key, error); check_not_null!(value, error); @@ -744,7 +781,7 @@ extern "C" fn database_get_option( error: *mut FFI_AdbcError, ) -> AdbcStatusCode { catch_panic(error, || { - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); check_not_null!(key, error); check_not_null!(value, error); check_not_null!(length, error); @@ -753,9 +790,9 @@ extern "C" fn database_get_option( unsafe { database_private_data::(database) }, error ); - let (options, database) = exported.tuple(); + let (mut options, mut database) = check_err!(exported.tuple(), error); - let optvalue = unsafe { get_option(database, options, key) }; + let optvalue = unsafe { get_option(database.as_deref_mut(), options.as_deref_mut(), key) }; let optvalue = check_err!(optvalue, error); check_err!(unsafe { copy_string(&optvalue, value, length) }, error); @@ -770,7 +807,7 @@ extern "C" fn database_get_option_int( error: *mut FFI_AdbcError, ) -> AdbcStatusCode { catch_panic(error, || { - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); check_not_null!(key, error); check_not_null!(value, error); @@ -778,9 +815,12 @@ extern "C" fn database_get_option_int( unsafe { database_private_data::(database) }, error ); - let (options, database) = exported.tuple(); + let (mut options, mut database) = check_err!(exported.tuple(), error); - let optvalue = check_err!(unsafe { get_option_int(database, options, key) }, error); + let optvalue = check_err!( + unsafe { get_option_int(database.as_deref_mut(), options.as_deref_mut(), key) }, + error + ); unsafe { std::ptr::write_unaligned(value, optvalue) }; ADBC_STATUS_OK @@ -794,7 +834,7 @@ extern "C" fn database_get_option_double( error: *mut FFI_AdbcError, ) -> AdbcStatusCode { catch_panic(error, || { - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); check_not_null!(key, error); check_not_null!(value, error); @@ -802,9 +842,12 @@ extern "C" fn database_get_option_double( unsafe { database_private_data::(database) }, error ); - let (options, database) = exported.tuple(); + let (mut options, mut database) = check_err!(exported.tuple(), error); - let optvalue = check_err!(unsafe { get_option_double(database, options, key) }, error); + let optvalue = check_err!( + unsafe { get_option_double(database.as_deref_mut(), options.as_deref_mut(), key) }, + error + ); unsafe { std::ptr::write_unaligned(value, optvalue) }; ADBC_STATUS_OK @@ -819,7 +862,7 @@ extern "C" fn database_get_option_bytes( error: *mut FFI_AdbcError, ) -> AdbcStatusCode { catch_panic(error, || { - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); check_not_null!(key, error); check_not_null!(value, error); check_not_null!(length, error); @@ -828,9 +871,10 @@ extern "C" fn database_get_option_bytes( unsafe { database_private_data::(database) }, error ); - let (options, database) = exported.tuple(); + let (mut options, mut database) = check_err!(exported.tuple(), error); - let optvalue = unsafe { get_option_bytes(database, options, key) }; + let optvalue = + unsafe { get_option_bytes(database.as_deref_mut(), options.as_deref_mut(), key) }; let optvalue = check_err!(optvalue, error); unsafe { copy_bytes(&optvalue, value, length) }; @@ -917,7 +961,7 @@ extern "C" fn connection_init( ) -> AdbcStatusCode { catch_panic(error, || { let connection = pointer_as_mut!(connection, error); - let database = pointer_as_mut!(database, error); + check_not_null!(database, error); // It's OK to have a &mut here since the caller can't yet get a cancel handle let exported_connection = check_err!( @@ -948,6 +992,15 @@ extern "C" fn connection_init( error ); + let database = check_err!( + database.lock().map_err(|_| { + Error::with_message_and_status( + "Database mutex poisoned", + Status::Internal, + ) + }), + error + ); database.new_connection_with_opts(options.clone()) } _ => Err(Error::with_message_and_status(