feat: add config to disable or restrict native UDF library loading, and load libraries without a global lock

docs/source/user-guide/latest/rust_udfs.md

@@ -178,9 +178,30 @@ rather than accept a function whose volatility it would go on to ignore.
 `libraryPath` is passed to the platform's dynamic loader. An absolute path is what you want in
 practice, and it is what the rest of this page assumes, but a bare library name resolves the same
 way it would for any other shared object, through `LD_LIBRARY_PATH` on Linux and
-`DYLD_LIBRARY_PATH` on macOS. That is a convenience, not a sandbox: Comet does not restrict which
-paths may be loaded, so it makes no difference to the trust decision described under
-[Limitations](#limitations).
+`DYLD_LIBRARY_PATH` on macOS. That is a convenience, not a sandbox: unless you set the configs
+below, Comet does not restrict which paths may be loaded, so it makes no difference to the trust
+decision described under [Limitations](#limitations).
+
+## Restricting which libraries can be loaded
+
+Loading a library runs its code with the privileges of the JVM, so an operator can limit it:
+
+| Config                               | Default | Effect                                                       |
+| ------------------------------------ | ------- | ------------------------------------------------------------ |
+| `spark.comet.nativeUdf.enabled`      | `true`  | When `false`, no native UDF library is loaded.               |
+| `spark.comet.nativeUdf.allowedPaths` | empty   | Comma-separated directories a library must be under, if set. |
+
+With `allowedPaths` set, a library is loaded only if its path is absolute and, once symlinks and
+`..` are resolved, lies under one of the listed directories (which are resolved the same way). A bare
+library name is refused, because the dynamic loader's search path is not something Comet can check.
+The check applies on the driver, where `register` fails with `CometNativeUdfNotAllowedException`,
+and again on each executor before it loads a library named in a plan.
+
+These are ordinary Spark configs, so they are set in the cluster's defaults and are not a defense
+against a user who can change session configuration: such a user can set them back. The allow-list
+limits which files Comet will load, not what the code in them does, and a file that is replaced
+between the check and the load is not detected. List only directories that are writable by users
+you trust.
 
 ### When Spark evaluates the call
 
@@ -310,7 +331,9 @@ This feature is at an early stage. The current limitations are:
   layout.
 - **Loading a library is loading native code.** It runs with the full privileges of the executor
   process and Comet cannot sandbox it: a bug in a UDF can corrupt memory or crash the executor.
-  Only register libraries you trust and control.
+  Only register libraries you trust and control. See
+  [Restricting which libraries can be loaded](#restricting-which-libraries-can-be-loaded) for the
+  configs that turn the feature off or limit the directories libraries may come from.
 - Once loaded, a library stays loaded for the life of the process, and `register` loads it on the
   driver as well as on the executors. To deploy a new build, write it to a new path and register
   that, or restart the executors and the driver. Never copy a new build over a library that is

native/core/src/comet_native_udf_bridge.rs

@@ -20,11 +20,14 @@
 
 use crate::errors::{try_unwrap_or_throw, CometError};
 use crate::execution::c_udf::cache::get_or_load;
+use crate::execution::c_udf::policy::NativeUdfPolicy;
 use jni::objects::{JClass, JString};
+use jni::sys::jboolean;
 use jni::EnvUnowned;
 
 /// Validate that `library_path` loads and exposes a UDF named
-/// `expected_name`.
+/// `expected_name`, once `enabled` and `allowed_paths` (the driver's
+/// `spark.comet.nativeUdf.*` configs) have allowed the path.
 ///
 /// Returns normally when it does and throws otherwise. The driver only
 /// needs that yes-or-no answer: everything else it knows about the UDF
@@ -37,6 +40,8 @@ pub extern "system" fn Java_org_apache_comet_udf_CometNativeUdfBridge_validateLi
     _class: JClass,
     library_path: JString,
     expected_name: JString,
+    enabled: jboolean,
+    allowed_paths: JString,
 ) {
     try_unwrap_or_throw(&e, |env| {
         let path: String = library_path
@@ -45,6 +50,12 @@ pub extern "system" fn Java_org_apache_comet_udf_CometNativeUdfBridge_validateLi
         let name: String = expected_name
             .try_to_string(env)
             .map_err(|e| CometError::Internal(e.to_string()))?;
+        let allowed_paths: String = allowed_paths
+            .try_to_string(env)
+            .map_err(|e| CometError::Internal(e.to_string()))?;
+        NativeUdfPolicy::new(enabled, &allowed_paths)
+            .check(&path)
+            .map_err(|e| CometError::Internal(e.to_string()))?;
         let lib = get_or_load(&path).map_err(|e| CometError::Internal(e.to_string()))?;
         if !lib.udfs.iter().any(|u| u.name == name) {
             // `CometNativeUDF.classifyNativeError` keys on "' not found in " to turn this into a

native/core/src/execution/c_udf/cache.rs

@@ -25,58 +25,112 @@
 
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
-use std::sync::{Arc, OnceLock, RwLock};
+use std::sync::{Arc, Mutex, MutexGuard, OnceLock, PoisonError, RwLock};
 
 use super::loader::{load, LoadedLibrary, LoaderError};
 
-static CACHE: OnceLock<RwLock<HashMap<PathBuf, Arc<LoadedLibrary>>>> = OnceLock::new();
+/// One library's slot in the cache. The mutex is held while the library loads, so concurrent
+/// requests for the same library wait for the one load, and a request for any other library does
+/// not wait at all.
+type Slot<T> = Arc<Mutex<Option<Arc<T>>>>;
 
-fn cache() -> &'static RwLock<HashMap<PathBuf, Arc<LoadedLibrary>>> {
-    CACHE.get_or_init(|| RwLock::new(HashMap::new()))
+/// Loaded libraries keyed by path. Both the path as the caller wrote it and its canonical form map
+/// to the same slot, so a symlink and its target share one library.
+///
+/// The map lock is only ever held for a lookup or an insert, never across a load. Neither it nor a
+/// slot is treated as lost when a thread panicked while holding it: a map operation cannot leave
+/// the map half-updated, and a slot is only filled once its load has succeeded, so a poisoned slot
+/// still reads as "not loaded" and the next request retries. A library whose load panics therefore
+/// fails its own queries and nothing else.
+struct Cache<T> {
+    slots: RwLock<HashMap<PathBuf, Slot<T>>>,
 }
 
-/// Get an already-loaded library, or load and cache it.
-pub fn get_or_load(path: impl AsRef<Path>) -> Result<Arc<LoadedLibrary>, LoaderError> {
-    let raw = path.as_ref().to_path_buf();
+impl<T> Cache<T> {
+    fn new() -> Self {
+        Self {
+            slots: RwLock::new(HashMap::new()),
+        }
+    }
 
-    if let Some(lib) = cache().read().unwrap().get(&raw).cloned() {
-        return Ok(lib);
+    fn existing(&self, path: &Path) -> Option<Slot<T>> {
+        self.slots
+            .read()
+            .unwrap_or_else(PoisonError::into_inner)
+            .get(path)
+            .cloned()
     }
 
-    let canonical = raw.canonicalize().unwrap_or_else(|_| raw.clone());
-    if canonical != raw {
-        // A statement of its own, so the read guard is dropped before the write lock is taken. In
-        // an `if let` scrutinee it would live until the end of the block, and this thread would
-        // wait on its own read lock forever.
-        let hit = cache().read().unwrap().get(&canonical).cloned();
-        if let Some(lib) = hit {
-            cache().write().unwrap().insert(raw, Arc::clone(&lib));
-            return Ok(lib);
+    fn slot(&self, path: &Path) -> Slot<T> {
+        if let Some(slot) = self.existing(path) {
+            return slot;
         }
+        let mut slots = self.slots.write().unwrap_or_else(PoisonError::into_inner);
+        Arc::clone(slots.entry(path.to_path_buf()).or_default())
     }
 
-    // The write lock is held across `load`, which runs the cdylib's static initializers, so a slow
-    // load of one library blocks lookups of every other path. See
-    // <https://github.com/apache/datafusion-comet/issues/5297>.
-    let mut w = cache().write().unwrap();
-    if let Some(lib) = w.get(&canonical).cloned() {
-        if canonical != raw {
-            w.insert(raw, Arc::clone(&lib));
-        }
-        return Ok(lib);
+    fn alias(&self, path: PathBuf, slot: &Slot<T>) {
+        self.slots
+            .write()
+            .unwrap_or_else(PoisonError::into_inner)
+            .entry(path)
+            .or_insert_with(|| Arc::clone(slot));
     }
-    let loaded = Arc::new(load(&canonical)?);
-    w.insert(canonical.clone(), Arc::clone(&loaded));
-    if canonical != raw {
-        w.insert(raw, Arc::clone(&loaded));
+
+    fn get_or_load_with<E>(
+        &self,
+        path: &Path,
+        load: impl FnOnce(&Path) -> Result<T, E>,
+    ) -> Result<Arc<T>, E> {
+        if let Some(slot) = self.existing(path) {
+            if let Some(lib) = lock(&slot).clone() {
+                return Ok(lib);
+            }
+        }
+
+        let canonical = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
+        let slot = self.slot(&canonical);
+        let lib = {
+            let mut guard = lock(&slot);
+            match guard.as_ref() {
+                Some(lib) => Arc::clone(lib),
+                None => {
+                    // An error leaves the slot empty, so a later request tries again.
+                    let lib = Arc::new(load(&canonical)?);
+                    *guard = Some(Arc::clone(&lib));
+                    lib
+                }
+            }
+        };
+        if canonical != path {
+            self.alias(path.to_path_buf(), &slot);
+        }
+        Ok(lib)
     }
-    Ok(loaded)
+}
+
+fn lock<T>(slot: &Slot<T>) -> MutexGuard<'_, Option<Arc<T>>> {
+    slot.lock().unwrap_or_else(PoisonError::into_inner)
+}
+
+static CACHE: OnceLock<Cache<LoadedLibrary>> = OnceLock::new();
+
+/// Get an already-loaded library, or load and cache it.
+///
+/// Loading one library does not block lookups of any other; see
+/// <https://github.com/apache/datafusion-comet/issues/5297>.
+pub fn get_or_load(path: impl AsRef<Path>) -> Result<Arc<LoadedLibrary>, LoaderError> {
+    CACHE
+        .get_or_init(Cache::new)
+        .get_or_load_with(path.as_ref(), |p| load(p))
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::execution::c_udf::test_support::{test_udfs_path, BUILD_HINT};
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::time::Duration;
 
     #[test]
     fn same_path_returns_same_arc() {
@@ -92,6 +146,108 @@ mod tests {
         assert!(matches!(err, LoaderError::Open { .. }));
     }
 
+    /// While one library is loading, a request for another library is answered. With a single
+    /// lock held across the load, the second request would wait for the first to finish, and this
+    /// test would time out.
+    #[test]
+    fn a_slow_load_does_not_block_another_library() {
+        let cache = Arc::new(Cache::<&'static str>::new());
+        let (started_tx, started_rx) = std::sync::mpsc::channel();
+        let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
+
+        let slow = {
+            let cache = Arc::clone(&cache);
+            std::thread::spawn(move || {
+                cache.get_or_load_with::<()>(Path::new("/no/such/slow.so"), |_| {
+                    started_tx.send(()).unwrap();
+                    release_rx.recv().unwrap();
+                    Ok("slow")
+                })
+            })
+        };
+        started_rx.recv_timeout(Duration::from_secs(30)).unwrap();
+
+        let (tx, rx) = std::sync::mpsc::channel();
+        {
+            let cache = Arc::clone(&cache);
+            std::thread::spawn(move || {
+                let _ = tx.send(
+                    cache.get_or_load_with::<()>(Path::new("/no/such/fast.so"), |_| Ok("fast")),
+                );
+            });
+        }
+        let fast = rx
+            .recv_timeout(Duration::from_secs(30))
+            .expect("loading a second library waited for the first")
+            .unwrap();
+        assert_eq!(*fast, "fast");
+
+        release_tx.send(()).unwrap();
+        assert_eq!(*slow.join().unwrap().unwrap(), "slow");
+    }
+
+    /// Requests for one library that arrive while it loads wait for that load and share its result.
+    #[test]
+    fn concurrent_requests_for_one_library_load_it_once() {
+        let cache = Arc::new(Cache::<usize>::new());
+        let loads = Arc::new(AtomicUsize::new(0));
+        let threads: Vec<_> = (0..8)
+            .map(|_| {
+                let cache = Arc::clone(&cache);
+                let loads = Arc::clone(&loads);
+                std::thread::spawn(move || {
+                    cache
+                        .get_or_load_with::<()>(Path::new("/no/such/one.so"), |_| {
+                            std::thread::sleep(Duration::from_millis(50));
+                            Ok(loads.fetch_add(1, Ordering::SeqCst))
+                        })
+                        .unwrap()
+                })
+            })
+            .collect();
+        let results: Vec<_> = threads.into_iter().map(|t| t.join().unwrap()).collect();
+        assert_eq!(loads.load(Ordering::SeqCst), 1);
+        assert!(results.iter().all(|r| Arc::ptr_eq(r, &results[0])));
+    }
+
+    /// A load that panics fails only that library: the cache keeps serving every other path, and a
+    /// later request for the same path retries instead of panicking on a poisoned lock.
+    #[test]
+    fn a_load_that_panics_does_not_disable_the_cache() {
+        let cache = Arc::new(Cache::<&'static str>::new());
+        let bad = Path::new("/no/such/bad.so");
+
+        let panicked = {
+            let cache = Arc::clone(&cache);
+            std::thread::spawn(move || {
+                let _ = cache.get_or_load_with::<()>(bad, |_| panic!("static initializer failed"));
+            })
+            .join()
+        };
+        assert!(panicked.is_err());
+
+        let other = cache
+            .get_or_load_with::<()>(Path::new("/no/such/other.so"), |_| Ok("other"))
+            .unwrap();
+        assert_eq!(*other, "other");
+        let retried = cache
+            .get_or_load_with::<()>(bad, |_| Ok("recovered"))
+            .unwrap();
+        assert_eq!(*retried, "recovered");
+    }
+
+    /// A failed load is not remembered.
+    #[test]
+    fn a_failed_load_is_retried() {
+        let cache = Cache::<&'static str>::new();
+        let path = Path::new("/no/such/flaky.so");
+        assert!(cache.get_or_load_with(path, |_| Err("boom")).is_err());
+        assert_eq!(
+            *cache.get_or_load_with::<()>(path, |_| Ok("ok")).unwrap(),
+            "ok"
+        );
+    }
+
     /// A path seen for the first time that resolves to an already-loaded library, here a symlink,
     /// takes the branch that records the new path under the write lock. That branch used to take
     /// the lock while still holding its own read guard, and hung. The lookup runs on a thread of

native/core/src/execution/c_udf/mod.rs

@@ -32,6 +32,7 @@
 pub mod cache;
 pub mod imported_c;
 pub mod loader;
+pub mod policy;
 
 #[cfg(test)]
 pub(crate) mod test_support;