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(®istered.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}