Skip to main content

fui/
worker_host_services.rs

1use crate::worker_runtime::WorkerRuntime;
2use std::collections::HashMap;
3use std::panic::{catch_unwind, AssertUnwindSafe};
4use std::sync::{Arc, Mutex, OnceLock};
5
6#[derive(Clone, Debug, PartialEq)]
7#[doc(hidden)]
8pub enum NativeWorkerHostServiceValue {
9    String(String),
10    Bool(bool),
11    I32(i32),
12    U32(u32),
13    I64(i64),
14    U64(u64),
15    F64(f64),
16    Bytes(Vec<u8>),
17    I32Array(Vec<i32>),
18    U32Array(Vec<u32>),
19    I64Array(Vec<i64>),
20    U64Array(Vec<u64>),
21    F64Array(Vec<f64>),
22    Void,
23}
24
25#[doc(hidden)]
26pub trait NativeWorkerHostServiceType: Default + Sized {
27    const TYPE_NAME: &'static str;
28    fn into_native_worker_host_service_value(self) -> NativeWorkerHostServiceValue;
29    fn from_native_worker_host_service_value(
30        value: NativeWorkerHostServiceValue,
31    ) -> Result<Self, String>;
32}
33
34macro_rules! native_worker_host_service_type {
35    ($type:ty, $variant:ident, $name:literal) => {
36        impl NativeWorkerHostServiceType for $type {
37            const TYPE_NAME: &'static str = $name;
38
39            fn into_native_worker_host_service_value(self) -> NativeWorkerHostServiceValue {
40                NativeWorkerHostServiceValue::$variant(self)
41            }
42
43            fn from_native_worker_host_service_value(
44                value: NativeWorkerHostServiceValue,
45            ) -> Result<Self, String> {
46                match value {
47                    NativeWorkerHostServiceValue::$variant(value) => Ok(value),
48                    _ => Err(format!("expected {}", Self::TYPE_NAME)),
49                }
50            }
51        }
52    };
53}
54
55native_worker_host_service_type!(String, String, "string");
56native_worker_host_service_type!(bool, Bool, "bool");
57native_worker_host_service_type!(i32, I32, "i32");
58native_worker_host_service_type!(u32, U32, "u32");
59native_worker_host_service_type!(i64, I64, "i64");
60native_worker_host_service_type!(u64, U64, "u64");
61native_worker_host_service_type!(f64, F64, "f64");
62native_worker_host_service_type!(Vec<u8>, Bytes, "bytes");
63native_worker_host_service_type!(Vec<i32>, I32Array, "i32_array");
64native_worker_host_service_type!(Vec<u32>, U32Array, "u32_array");
65native_worker_host_service_type!(Vec<i64>, I64Array, "i64_array");
66native_worker_host_service_type!(Vec<u64>, U64Array, "u64_array");
67native_worker_host_service_type!(Vec<f64>, F64Array, "f64_array");
68
69impl NativeWorkerHostServiceType for () {
70    const TYPE_NAME: &'static str = "void";
71
72    fn into_native_worker_host_service_value(self) -> NativeWorkerHostServiceValue {
73        NativeWorkerHostServiceValue::Void
74    }
75
76    fn from_native_worker_host_service_value(
77        value: NativeWorkerHostServiceValue,
78    ) -> Result<Self, String> {
79        match value {
80            NativeWorkerHostServiceValue::Void => Ok(()),
81            _ => Err(format!("expected {}", Self::TYPE_NAME)),
82        }
83    }
84}
85
86type Handler = dyn Fn(Vec<NativeWorkerHostServiceValue>) -> Result<NativeWorkerHostServiceValue, String>
87    + Send
88    + Sync
89    + 'static;
90
91struct RegisteredHandler {
92    generation: u64,
93    handler: Arc<Handler>,
94}
95
96#[derive(Default)]
97struct Registry {
98    next_generation: u64,
99    handlers: HashMap<&'static str, RegisteredHandler>,
100}
101
102fn registry() -> &'static Mutex<Registry> {
103    static REGISTRY: OnceLock<Mutex<Registry>> = OnceLock::new();
104    REGISTRY.get_or_init(|| Mutex::new(Registry::default()))
105}
106
107pub struct NativeWorkerHostServiceRegistration {
108    import_name: &'static str,
109    generation: u64,
110}
111
112impl Drop for NativeWorkerHostServiceRegistration {
113    fn drop(&mut self) {
114        let mut registry = registry()
115            .lock()
116            .unwrap_or_else(|poisoned| poisoned.into_inner());
117        if registry
118            .handlers
119            .get(self.import_name)
120            .is_some_and(|registered| registered.generation == self.generation)
121        {
122            registry.handlers.remove(self.import_name);
123        }
124    }
125}
126
127#[doc(hidden)]
128pub fn register_native_worker_host_service(
129    import_name: &'static str,
130    handler: impl Fn(Vec<NativeWorkerHostServiceValue>) -> Result<NativeWorkerHostServiceValue, String>
131        + Send
132        + Sync
133        + 'static,
134) -> Result<NativeWorkerHostServiceRegistration, String> {
135    let mut registry = registry()
136        .lock()
137        .unwrap_or_else(|poisoned| poisoned.into_inner());
138    if registry.handlers.contains_key(import_name) {
139        return Err(format!(
140            "Native Worker host service {import_name} is already registered."
141        ));
142    }
143    registry.next_generation = registry.next_generation.wrapping_add(1).max(1);
144    let generation = registry.next_generation;
145    registry.handlers.insert(
146        import_name,
147        RegisteredHandler {
148            generation,
149            handler: Arc::new(handler),
150        },
151    );
152    Ok(NativeWorkerHostServiceRegistration {
153        import_name,
154        generation,
155    })
156}
157
158#[doc(hidden)]
159pub fn native_worker_host_service_arg<T: NativeWorkerHostServiceType>(
160    args: &mut std::vec::IntoIter<NativeWorkerHostServiceValue>,
161    import_name: &str,
162    index: usize,
163) -> Result<T, String> {
164    let value = args.next().ok_or_else(|| {
165        format!("Native Worker host service {import_name} is missing argument {index}.")
166    })?;
167    T::from_native_worker_host_service_value(value).map_err(|message| {
168        format!("Native Worker host service {import_name} argument {index} {message}.")
169    })
170}
171
172#[cfg(not(target_arch = "wasm32"))]
173unsafe extern "C" {
174    fn fui_native_worker_host_service_is_allowed(name: *const u8, length: u32) -> bool;
175}
176
177fn is_allowed(import_name: &str) -> bool {
178    #[cfg(target_arch = "wasm32")]
179    {
180        let _ = import_name;
181        false
182    }
183    #[cfg(not(target_arch = "wasm32"))]
184    unsafe {
185        fui_native_worker_host_service_is_allowed(
186            import_name.as_ptr(),
187            import_name.len().try_into().unwrap_or(u32::MAX),
188        )
189    }
190}
191
192fn invoke(
193    import_name: &'static str,
194    args: Vec<NativeWorkerHostServiceValue>,
195) -> Result<NativeWorkerHostServiceValue, String> {
196    if !is_allowed(import_name) {
197        return Err(format!(
198            "Native Worker host service {import_name} is not allowed by the active Worker declaration."
199        ));
200    }
201    let handler = {
202        let registry = registry()
203            .lock()
204            .unwrap_or_else(|poisoned| poisoned.into_inner());
205        registry
206            .handlers
207            .get(import_name)
208            .map(|registered| Arc::clone(&registered.handler))
209    }
210    .ok_or_else(|| format!("Native Worker host service {import_name} is not registered."))?;
211    match catch_unwind(AssertUnwindSafe(|| handler(args))) {
212        Ok(result) => result,
213        Err(_) => Err(format!(
214            "Native Worker host service {import_name} panicked."
215        )),
216    }
217}
218
219#[doc(hidden)]
220pub fn invoke_native_worker_host_service<T: NativeWorkerHostServiceType>(
221    import_name: &'static str,
222    args: Vec<NativeWorkerHostServiceValue>,
223) -> T {
224    match invoke(import_name, args).and_then(|value| {
225        T::from_native_worker_host_service_value(value).map_err(|message| {
226            format!("Native Worker host service {import_name} returned {message}.")
227        })
228    }) {
229        Ok(value) => value,
230        Err(message) => {
231            WorkerRuntime::fail(message);
232            T::default()
233        }
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
241    use std::thread;
242
243    static ALLOWED: AtomicBool = AtomicBool::new(true);
244
245    #[no_mangle]
246    extern "C" fn fui_native_worker_host_service_is_allowed(
247        _name: *const u8,
248        _length: u32,
249    ) -> bool {
250        ALLOWED.load(Ordering::SeqCst)
251    }
252
253    #[test]
254    fn native_registry_normalizes_all_values_and_releases_owned_buffers() {
255        ALLOWED.store(true, Ordering::SeqCst);
256        let registration = register_native_worker_host_service("allValues", |args| {
257            assert_eq!(args.len(), 13);
258            assert_eq!(
259                args[0],
260                NativeWorkerHostServiceValue::String("hello".into())
261            );
262            assert_eq!(args[1], NativeWorkerHostServiceValue::Bool(true));
263            assert_eq!(
264                args[12],
265                NativeWorkerHostServiceValue::F64Array(vec![1.5, 2.5])
266            );
267            Ok(NativeWorkerHostServiceValue::Bytes(vec![7, 8, 9]))
268        })
269        .expect("register service");
270        let result: Vec<u8> = invoke_native_worker_host_service(
271            "allValues",
272            vec![
273                NativeWorkerHostServiceValue::String("hello".into()),
274                NativeWorkerHostServiceValue::Bool(true),
275                NativeWorkerHostServiceValue::I32(-2),
276                NativeWorkerHostServiceValue::U32(2),
277                NativeWorkerHostServiceValue::I64(-3),
278                NativeWorkerHostServiceValue::U64(3),
279                NativeWorkerHostServiceValue::F64(4.5),
280                NativeWorkerHostServiceValue::Bytes(vec![1]),
281                NativeWorkerHostServiceValue::I32Array(vec![-1]),
282                NativeWorkerHostServiceValue::U32Array(vec![1]),
283                NativeWorkerHostServiceValue::I64Array(vec![-2]),
284                NativeWorkerHostServiceValue::U64Array(vec![2]),
285                NativeWorkerHostServiceValue::F64Array(vec![1.5, 2.5]),
286            ],
287        );
288        assert_eq!(result, [7, 8, 9]);
289        drop(registration);
290    }
291
292    #[test]
293    fn native_registry_enforces_allowlist_before_invocation() {
294        let invoked = Arc::new(AtomicUsize::new(0));
295        let invoked_by_handler = Arc::clone(&invoked);
296        let _registration = register_native_worker_host_service("denied", move |_| {
297            invoked_by_handler.fetch_add(1, Ordering::SeqCst);
298            Ok(NativeWorkerHostServiceValue::Void)
299        })
300        .expect("register service");
301        ALLOWED.store(false, Ordering::SeqCst);
302        let _: () = invoke_native_worker_host_service("denied", vec![]);
303        ALLOWED.store(true, Ordering::SeqCst);
304        assert_eq!(invoked.load(Ordering::SeqCst), 0);
305    }
306
307    #[test]
308    fn native_registry_normalizes_unknown_malformed_error_and_panic_results() {
309        ALLOWED.store(true, Ordering::SeqCst);
310        assert!(invoke("unknown", vec![])
311            .unwrap_err()
312            .contains("not registered"));
313
314        let mut missing = Vec::new().into_iter();
315        assert!(
316            native_worker_host_service_arg::<u32>(&mut missing, "malformed", 0)
317                .unwrap_err()
318                .contains("missing argument 0")
319        );
320        let mut wrong = vec![NativeWorkerHostServiceValue::String("wrong".into())].into_iter();
321        assert!(
322            native_worker_host_service_arg::<u32>(&mut wrong, "malformed", 0)
323                .unwrap_err()
324                .contains("expected u32")
325        );
326
327        let malformed = register_native_worker_host_service("malformed", |_| {
328            Ok(NativeWorkerHostServiceValue::String("wrong".into()))
329        })
330        .expect("register malformed service");
331        assert!(
332            <u32 as NativeWorkerHostServiceType>::from_native_worker_host_service_value(
333                invoke("malformed", vec![]).expect("invoke malformed service")
334            )
335            .unwrap_err()
336            .contains("expected u32")
337        );
338        drop(malformed);
339
340        let error = register_native_worker_host_service("error", |_| Err("unavailable".into()))
341            .expect("register error service");
342        assert_eq!(invoke("error", vec![]), Err("unavailable".into()));
343        drop(error);
344
345        let panicking = register_native_worker_host_service("panicking", |_| panic!("boom"))
346            .expect("register panicking service");
347        assert!(invoke("panicking", vec![])
348            .unwrap_err()
349            .contains("panicked"));
350        drop(panicking);
351    }
352
353    #[test]
354    fn native_registry_supports_concurrent_worker_calls() {
355        ALLOWED.store(true, Ordering::SeqCst);
356        let _registration = register_native_worker_host_service("concurrent", |args| {
357            let mut args = args.into_iter();
358            let value: u32 = native_worker_host_service_arg(&mut args, "concurrent", 0)?;
359            Ok(NativeWorkerHostServiceValue::U32(value + 1))
360        })
361        .expect("register concurrent service");
362        let threads: Vec<_> = (0..8)
363            .map(|value| {
364                thread::spawn(move || {
365                    invoke("concurrent", vec![NativeWorkerHostServiceValue::U32(value)])
366                })
367            })
368            .collect();
369        for (value, thread) in threads.into_iter().enumerate() {
370            assert_eq!(
371                thread.join().expect("join worker"),
372                Ok(NativeWorkerHostServiceValue::U32(value as u32 + 1))
373            );
374        }
375    }
376}